diff --git a/giant/cli.py b/giant/cli.py index 2a9f7f7..f5bc144 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -67,6 +67,7 @@ def _write_prediction_ref( pred_uuid: str, out: Path, dataset_path: Path, + comment: str | None = None, ) -> Path: """Write a YAML sidecar in the checkpoint directory and return its path.""" ref = { @@ -76,6 +77,8 @@ def _write_prediction_ref( "checkpoint": str(checkpoint.resolve()), "timestamp": datetime.now(timezone.utc).isoformat(), } + if comment is not None: + ref["comment"] = comment ref_path = checkpoint.parent / f"{pred_uuid}.yaml" ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False)) return ref_path @@ -315,6 +318,14 @@ def predict( Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (default: auto)"), ] = None, + comment: Annotated[ + Optional[str], + typer.Option( + "--comment", + "-m", + help="Free-text note recorded in the prediction's YAML sidecar", + ), + ] = None, ) -> None: """Run trained model on a parquet file and save predictions.""" batch_size_auto = False @@ -532,7 +543,9 @@ def predict( if writer is not None: writer.close() - ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path) + ref_path = _write_prediction_ref( + checkpoint, pred_uuid, out, dataset_path, comment + ) typer.echo(f"reference: {ref_path}") if skipped: diff --git a/tests/test_cli_predict.py b/tests/test_cli_predict.py index d74ee93..a6f74fb 100644 --- a/tests/test_cli_predict.py +++ b/tests/test_cli_predict.py @@ -103,6 +103,25 @@ def test_ref_yaml_contains_expected_fields(tmp_path): assert data["dataset"] == str(dataset) assert data["checkpoint"] == str(checkpoint.resolve()) assert "timestamp" in data + assert "comment" not in data + + +def test_ref_yaml_includes_comment_when_provided(tmp_path): + ckpt_dir = tmp_path / "checkpoints" + ckpt_dir.mkdir() + checkpoint = ckpt_dir / "best.pt" + checkpoint.touch() + + out = tmp_path / "pred.parquet" + dataset = tmp_path / "full.manifest" + pred_uuid = str(uuid.uuid4()) + + ref_path = _write_prediction_ref( + checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3" + ) + data = yaml.safe_load(ref_path.read_text()) + + assert data["comment"] == "baseline sweep run 3" def test_ref_timestamp_is_iso_format(tmp_path):