auto-sync 2026-07-03T10:26:01Z workspace (part 2)

workspace/scripts/train_ctt.py

@@ -19,6 +19,7 @@ if str(PROJECT_ROOT) not in sys.path:
 import numpy as np  # noqa: E402
 import torch  # noqa: E402
 
+from cil.chart_features import CHART_FEATURE_MODES, build_chart_feature  # noqa: E402
 from cil.models import CTTConfig, CausalTangentTransport, ChartEncoder, TangentNormalizer, UtilityEnergy  # noqa: E402
 from cil.models.ctt import chamfer_to_target_set, diversity_loss, negative_boundary_loss  # noqa: E402
 from cil.models.utility_energy import listwise_ranking_loss, pairwise_ranking_loss  # noqa: E402
@@ -57,6 +58,12 @@ def main(argv: list[str] | None = None) -> int:
     parser.add_argument("--negative-margin", type=float, default=0.2)
     parser.add_argument("--transport-samples-per-pair", type=int, default=4)
     parser.add_argument("--diversity-temperature", type=float, default=1.0)
+    parser.add_argument(
+        "--chart-feature-mode",
+        choices=CHART_FEATURE_MODES,
+        default="base",
+        help="Deployment-visible chart feature family used by ChartEncoder.",
+    )
     args = parser.parse_args(argv)
     if args.config:
         args = _merge_config(args, _load_simple_yaml(args.config))
@@ -68,7 +75,11 @@ def main(argv: list[str] | None = None) -> int:
     out_dir = args.out_dir
     out_dir.mkdir(parents=True, exist_ok=True)
     _write_run_provenance(out_dir, args)
-    charts, index = load_charts(args.dataset, max_charts=args.max_charts)
+    charts, index = load_charts(
+        args.dataset,
+        max_charts=args.max_charts,
+        chart_feature_mode=args.chart_feature_mode,
+    )
     if len(charts) < 2:
         raise SystemExit("CTT training requires at least two charts with positive tangents")
 
@@ -164,6 +175,7 @@ def main(argv: list[str] | None = None) -> int:
         "source_index": str(args.dataset),
         "data_hash": index.get("content_hash"),
         "split_hash": index.get("split_hash"),
+        "chart_feature_mode": args.chart_feature_mode,
     }
     torch.save(checkpoint, out_dir / "model.pt")
     metrics = {
@@ -175,6 +187,7 @@ def main(argv: list[str] | None = None) -> int:
         "final_mean_loss": mean_loss,
         "data_hash": index.get("content_hash"),
         "split_hash": index.get("split_hash"),
+        "chart_feature_mode": args.chart_feature_mode,
         "loss_weights": {
             "pos_alignment": args.pos_alignment,
             "negative_boundary": args.negative_boundary,
@@ -198,7 +211,12 @@ def main(argv: list[str] | None = None) -> int:
     return 0
 
 
-def load_charts(index_path: Path, *, max_charts: int | None = None) -> tuple[list[Chart], dict[str, Any]]:
+def load_charts(
+    index_path: Path,
+    *,
+    max_charts: int | None = None,
+    chart_feature_mode: str = "base",
+) -> tuple[list[Chart], dict[str, Any]]:
     index = json.loads(index_path.read_text())
     if not index.get("include_outcomes", False):
         raise SystemExit(f"{index_path} does not include outcomes; CTT training requires train split")
@@ -215,14 +233,17 @@ def load_charts(index_path: Path, *, max_charts: int | None = None) -> tuple[lis
             utilities = data["utility"]
             outcomes = data["outcome_vector"]
             seeds = data["seed"]
+            metadata_values = data["metadata_json"]
             for row in range(chart_ids.shape[0]):
                 chart_id = str(chart_ids[row])
+                metadata = _json_loads(str(metadata_values[row]))
                 item = grouped.setdefault(
                     chart_id,
                     {
                         "task_id": str(task_ids[row]),
                         "seed": str(seeds[row]),
                         "base": None,
+                        "metadata": metadata,
                         "positives": [],
                         "negatives": [],
                         "tangents": [],
@@ -232,6 +253,7 @@ def load_charts(index_path: Path, *, max_charts: int | None = None) -> tuple[lis
                 )
                 if bool(is_base[row]):
                     item["base"] = base_actions[row].astype("float32")
+                    item["metadata"] = metadata
                 label = str(labels[row])
                 tangent = spline_codes[row].astype("float32")
                 if label == "positive":
@@ -256,7 +278,14 @@ def load_charts(index_path: Path, *, max_charts: int | None = None) -> tuple[lis
                 chart_id=chart_id,
                 task_id=str(item["task_id"]),
                 seed=str(item["seed"]),
-                feature=torch.as_tensor(item["base"], dtype=torch.float32),
+                feature=torch.as_tensor(
+                    build_chart_feature(
+                        item["base"],
+                        item.get("metadata", {}),
+                        mode=chart_feature_mode,
+                    ),
+                    dtype=torch.float32,
+                ),
                 positives=positives,
                 negatives=negatives,
                 tangents=tangents,
@@ -269,6 +298,14 @@ def load_charts(index_path: Path, *, max_charts: int | None = None) -> tuple[lis
     return charts, index
 
 
+def _json_loads(value: str) -> dict[str, Any]:
+    try:
+        payload = json.loads(value)
+    except json.JSONDecodeError:
+        return {}
+    return payload if isinstance(payload, dict) else {}
+
+
 def _neighbor_pairs(charts: list[Chart], *, neighbors: int) -> list[tuple[int, int]]:
     pairs: list[tuple[int, int]] = []
     for target_idx, target in enumerate(charts):