code/experiments/exp9_kcn_focused_unified_encoder.py

#!/usr/bin/env python
"""Focused KSL+CASL+NSL wrapper for the Exp8-v2 unified encoder.

The base Exp8 collector can see every local stream. This wrapper restricts a run
to the paper-facing focused setting:

  - languages: KSL, CASL, NSL
  - modes: pose-only, RGB/image-only, or multimodal
  - no GSL Health sentence task unless explicitly added through base args

It preserves the Exp8-v2 model, losses, metrics, and result JSON format.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Sequence

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from experiments import exp8_unified_mixed_encoder as exp8  # noqa: E402
from experiments import exp8_v2_strong_unified_encoder as exp8_v2  # noqa: E402


def parse_focus_args(argv: Sequence[str]) -> tuple[argparse.Namespace, list[str]]:
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--focus-mode", choices=["pose", "rgb", "multimodal"], default="multimodal")
    parser.add_argument("--focus-languages", nargs="+", default=["casl", "ksl", "nsi"])
    parser.add_argument(
        "--require-eval-split",
        action="store_true",
        default=True,
        help="Drop train-only focused tasks from evaluation-oriented runs.",
    )
    parser.add_argument("--allow-train-only", dest="require_eval_split", action="store_false")
    return parser.parse_known_args(argv)


def main() -> None:
    focus_args, remaining = parse_focus_args(sys.argv[1:])
    sys.argv = [sys.argv[0], *remaining]

    wanted_langs = {lang.lower() for lang in focus_args.focus_languages}
    mode = focus_args.focus_mode
    original_collect_tasks = exp8.collect_tasks

    def focused_collect_tasks(args: argparse.Namespace) -> list[exp8.TaskSpec]:
        tasks = original_collect_tasks(args)
        focused: list[exp8.TaskSpec] = []
        for task in tasks:
            lang = task.language_code.lower()
            if lang not in wanted_langs:
                continue
            if mode == "pose" and task.modality != "pose":
                continue
            if mode == "rgb" and task.modality != "rgb":
                continue
            if mode == "multimodal" and task.modality not in {"pose", "rgb"}:
                continue
            if focus_args.require_eval_split and not (task.val_rows or task.test_rows):
                continue
            focused.append(task)
        if not focused:
            raise SystemExit(
                f"No tasks left after focus filtering: mode={mode}, languages={sorted(wanted_langs)}. "
                "Check manifests and local data."
            )
        return focused

    exp8.collect_tasks = focused_collect_tasks
    exp8_v2.main()


if __name__ == "__main__":
    main()