v7: add --held-out-classes and --keep-only-classes for generalization eval

train_v7_finetune.py

@@ -112,7 +112,10 @@ class BlindspotDataset(Dataset):
     def __init__(self, airport, mode, data_dir,
                  past_max=256, past_min=60,
                  delta_min=30, delta_max=120,
-                 seed=0, epoch_multiplier=4):
+                 seed=0, epoch_multiplier=4,
+                 held_out_classes=None,  # None = no filter; list = exclude these classes
+                 keep_only_classes=None,  # None = no filter; list = keep ONLY these classes (overrides held_out)
+                 ):
         ensure_data(airport, data_dir)
         airport_dir = os.path.join(data_dir, airport)
         raw, labels = load_atfm(airport, mode, airport_dir)
@@ -134,6 +137,17 @@ class BlindspotDataset(Dataset):
             raise RuntimeError(
                 f"No trajectories of length >= {min_required} in {airport}/{mode}"
             )
+
+        # Class-based filtering (held-out generalization eval)
+        if keep_only_classes is not None:
+            keep_set = set(int(c) for c in keep_only_classes)
+            class_keep = np.array([int(c) in keep_set for c in labels])
+            keep = keep & class_keep
+        elif held_out_classes is not None:
+            held = set(int(c) for c in held_out_classes)
+            class_keep = np.array([int(c) not in held for c in labels])
+            keep = keep & class_keep
+
         raw = raw[keep]
         lengths = lengths[keep]
         self.labels = labels[keep].astype(np.int64)
@@ -838,6 +852,10 @@ def main():
                    help="If --pretrained-encoder is a HF repo, name of the file in it.")
     p.add_argument("--freeze-encoder", action="store_true",
                    help="Freeze tokenizer + encoder weights after loading pretrained.")
+    p.add_argument("--held-out-classes", default=None,
+                   help="Comma-separated class IDs to EXCLUDE from training (e.g., '6,18,28').")
+    p.add_argument("--keep-only-classes", default=None,
+                   help="Comma-separated class IDs to KEEP for evaluation (eval on these only).")
     p.add_argument("--trackio-name", default=None)
     args = p.parse_args()
 
@@ -862,23 +880,34 @@ def main():
         trackio.init(project="flight-jepa-v2", name=args.trackio_name,
                       config=vars(args))
 
+    held_out = (
+        [int(c) for c in args.held_out_classes.split(",")]
+        if args.held_out_classes else None
+    )
+    keep_only = (
+        [int(c) for c in args.keep_only_classes.split(",")]
+        if args.keep_only_classes else None
+    )
     train_ds = BlindspotDataset(
         airport=args.airport, mode="TRAIN", data_dir=args.data_dir,
         past_max=args.past_max, past_min=args.past_min,
         delta_min=args.delta_min, delta_max=args.delta_max,
         seed=args.seed, epoch_multiplier=args.epoch_multiplier,
+        held_out_classes=held_out, keep_only_classes=keep_only,
     )
     test_ds = BlindspotDataset(
         airport=args.airport, mode="TEST", data_dir=args.data_dir,
         past_max=args.past_max, past_min=args.past_min,
         delta_min=args.delta_min, delta_max=args.delta_max,
         seed=args.seed + 1, epoch_multiplier=1,
+        held_out_classes=held_out, keep_only_classes=keep_only,
     )
     extrap_ds = BlindspotDataset(
         airport=args.airport, mode="TEST", data_dir=args.data_dir,
         past_max=args.past_max, past_min=args.past_min,
         delta_min=args.delta_min, delta_max=args.extrap_delta_max,
         seed=args.seed + 99, epoch_multiplier=1,
+        held_out_classes=held_out, keep_only_classes=keep_only,
     )
 
     train_dl = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True,