"""Run raw encoder-only heads: four epochs, no dev, fixed benchmark splits. This is not a capacity/context-matched Thinker ablation: native encoder vectors and validity flags are concatenated, so the existing MLP adapts to input width. This entry defaults to Ego4D; run_easycom_encoder_only.py shares the workflow. """ from __future__ import annotations import argparse import os from pathlib import Path import shlex import subprocess import sys from tqdm.auto import tqdm from easycom_head import DATASET_NAMES, load_session_data from local_turn_features import LOCAL_FORMAT, MODALITIES from prepare_encoder_only_features import ( DEFAULT_DATA_ROOTS, PACKING, encoder_only_mode, encoder_sample_split, sample_unit, ) REPO = Path(__file__).resolve().parent def parse_args(argv=None, *, dataset="ego4d"): if dataset not in DATASET_NAMES: raise ValueError(f"Unknown dataset: {dataset}") parser = argparse.ArgumentParser( description=f"{dataset}: raw encoder-only A/V/AV heads; four epochs, no dev, fixed train/test splits." ) parser.add_argument("--data-root", type=Path, default=DEFAULT_DATA_ROOTS[dataset]) parser.add_argument("--manifest", type=Path) parser.add_argument("--target-root", type=Path) parser.add_argument("--feature-root", type=Path, required=True, help="Prepared encoder-only inputs, not Thinker/local sidecars") parser.add_argument("--modalities", choices=tuple(MODALITIES), default="av") parser.add_argument("--checkpoint-root", type=Path, default=Path(f"checkpoints/{dataset}_encoder_only_1s")) parser.add_argument("--result-root", type=Path, default=Path(f"results/{dataset}_encoder_only_1s")) parser.add_argument("--seeds", type=int, nargs="+", default=[0, 1, 2]) parser.add_argument("--device", default="cuda:0") parser.add_argument("--dry-run", action="store_true", help="Validate all train/test inputs and print commands without training") parser.add_argument("--no-progress", action="store_true") args = parser.parse_args(argv) args.dataset = dataset if dataset == "ego4d": args.manifest = args.manifest or args.data_root / "ego4d_egospeak.tsv" elif args.manifest is not None: parser.error("EasyCom uses train sessions 4..12 and test sessions 3/2/1, not --manifest") args.target_root = args.target_root or args.data_root / f"{dataset}_target_frame" for key in ("data_root", "manifest", "target_root", "feature_root", "checkpoint_root", "result_root"): value = getattr(args, key) if value is not None: setattr(args, key, value.expanduser().resolve()) if len(set(args.seeds)) != len(args.seeds) or any(seed < 0 for seed in args.seeds): parser.error("Seeds must be unique and nonnegative") return args def run_paths(args, seed): name = f"encoder_{args.modalities}_seed{seed}" return args.checkpoint_root / f"{name}.pth", args.result_root / name def preflight(args): names, train, test = encoder_sample_split(args.dataset, args.manifest) unit = sample_unit(args.dataset) marker = args.feature_root / "complete.txt" if not marker.is_file() or marker.read_text().strip() != f"{len(names)} {unit}s": raise ValueError(f"Missing/incomplete encoder-only export: {marker}") for seed in args.seeds: checkpoint, result = run_paths(args, seed) if checkpoint.exists() or checkpoint.with_suffix(".pth.json").exists() or result.exists(): raise FileExistsError(f"Refusing to overwrite {checkpoint} or {result}; choose unfinished seeds/new output roots") expected, total, width = None, 0, None for name in tqdm(names, desc="encoder-only inputs", unit=unit, disable=args.no_progress): data = load_session_data(args.feature_root, args.target_root, name, dataset=args.dataset, expected_mode=encoder_only_mode(args.modalities)) config = data.feature_metadata.get("config") or {} contract = config.get("encoder_contract") or {} if (config.get("input_source") != "qwen-encoders-only" or config.get("uses_thinker_states") is not False or config.get("packing") != PACKING or contract.get("format") != LOCAL_FORMAT or contract.get("modalities") != args.modalities): raise ValueError(f"{name}: not a matching encoder-only export") encoder_config = contract.get("config") or {} if (encoder_config.get("causal_policy") != "half-open-trailing-window" or encoder_config.get("uses_speech_labels") is not False): raise ValueError(f"{name}: missing causal/no-speech-label encoder contract") dimensions = [encoder_config.get(f"{part}_dim") for part in MODALITIES[args.modalities]] if any(not isinstance(d, int) or d <= 0 for d in dimensions): raise ValueError(f"{name}: invalid encoder dimensions") width = sum(d + 1 for d in dimensions) if data.states.ndim != 2 or data.hidden_size != width or len(data.states) <= 10: raise ValueError(f"{name}: invalid native encoder-vector/validity-bit shape") if expected is not None and config != expected: raise ValueError(f"{name}: encoder export contract differs across {unit}s") expected = config total += len(data.states) del data print(f"{args.dataset} encoder-only {args.modalities}: {len(train)} train / {len(test)} test; " f"{total} states, input width={width}, MLP hidden width={width // 2}; epochs=4, dev=none", flush=True) print("No Thinker features or local-fusion branch. Native concatenation changes head capacity; " "context remains the saved encoder windows.", flush=True) def experiment_commands(args): common = ["--dataset", args.dataset, "--feature-root", str(args.feature_root), "--target-root", str(args.target_root), "--expected-mode", encoder_only_mode(args.modalities), "--device", args.device] if args.manifest is not None: common += ["--manifest", str(args.manifest)] if args.no_progress: common.append("--no-progress") for seed in args.seeds: checkpoint, result = run_paths(args, seed) train = [sys.executable, str(REPO / "train_turn_head.py"), *common, "--output", str(checkpoint), "--epochs", "4", "--batch-size", "512", "--learning-rate", "7e-5", "--weight-decay", "5e-5", "--warmup-epochs", "1.6", "--warmup-factor", "0.3", "--dropout", "0.2", "--loss", "lstr-mce", "--transition-objective", "none", "--seed", str(seed)] infer = [sys.executable, str(REPO / "infer_turn_head.py"), *common, "--checkpoint", str(checkpoint), "--output-dir", str(result)] if args.dataset == "easycom": _, train_sessions, test_sessions = encoder_sample_split("easycom") train += ["--sessions", *train_sessions] infer += ["--sessions", *test_sessions] yield train, infer def main(argv=None, *, dataset="ego4d"): args = parse_args(argv, dataset=dataset) preflight(args) env = os.environ.copy() for name in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"): env.setdefault(name, "4") for commands in experiment_commands(args): for command in commands: print(shlex.join(command), flush=True) if not args.dry_run: subprocess.run(command, cwd=REPO, env=env, check=True) summary = [sys.executable, str(REPO / "summarize_local_baselines.py"), "--root", str(args.result_root), "--dataset", args.dataset] print(shlex.join(summary), flush=True) if not args.dry_run: subprocess.run(summary, cwd=REPO, env=env, check=True) if __name__ == "__main__": main()