Add --comment option to predict, recorded in YAML sidecar

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

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:

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):