icassp2027-experiment/run_ego4d_encoder_only.py

"""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()