run_all.py

"""
Run every classifier variant against one frozen backbone, in sequence.

Each variant script already writes its artifacts to
results/<model>/<methodology>/ and upserts its row into
results/<model>/summary.csv (see shared/utils.py), so after a full sweep the
summary holds one row per methodology; a partial or by-hand run just fills
the rows it produced.

The two utterance-level cache builds are shared across the four attention
scripts (identical cache config), and the mean-pool script has its own
transcript-level cache, so a sweep extracts embeddings at most twice.

Usage:
  python run_all.py --model google/gemma-4-E2B-it
  python run_all.py --model meta-llama/Llama-3.1-70B --load-mode 4bit --gpu 1
  # hyperparameter sweep: longer utterances, three seeds
  python run_all.py --model google/gemma-4-E2B-it \
      --max-length-utterance 1024 --seeds 42 43 44

Only the flags you pass are forwarded, so each script keeps its own defaults
(e.g. their different extract batch sizes) unless you override them. Length is
per-featurization: --max-length-utterance goes to the four attention scripts,
--max-length-transcript to mean_pool. With --seeds the whole sweep repeats per
seed; caches are seed-independent, so only the first seed pays extraction.
"""
import os
import sys
import time
import argparse
import subprocess

REPO_ROOT = os.path.dirname(os.path.abspath(__file__))

# mean_pool first: its transcript-level cache build is the cheapest smoke test
# of the backbone; the four attention scripts then share one utterance cache.
# mean_pool truncates a whole transcript, so its length knob is separate from
# the per-utterance one the four attention scripts share.
# Transcript-level scripts (one vector per patient); they take the transcript
# length knob. mean_pool first (its cache is the cheapest backbone smoke test),
# then the last-token pooling variants. last_token runs before last_token_dropout
# so the dropout variant finds the shared last-token cache already built.
TRANSCRIPT_SCRIPTS = [
    "classifier_mean_pool.py",
    "classifier_last_token.py",
    "classifier_last_token_dropout.py",
    "classifier_last_k.py",
]
MEAN_POOL = TRANSCRIPT_SCRIPTS[0]  # back-compat alias
# Utterance-level scripts share one utterance cache and take the utterance knob.
ATTENTION_SCRIPTS = [
    "classifier_attention_pool.py",
    "classifier_attention_pool_multihead.py",
    "classifier_attention_pool_ffn.py",
    "classifier_attention_pool_multihead_ffn.py",
]
SCRIPTS = TRANSCRIPT_SCRIPTS + ATTENTION_SCRIPTS


def main():
    p = argparse.ArgumentParser(description="Run all classifier variants sequentially.")
    p.add_argument("--model", default=None,
                   help="HF model id for the frozen embedder (forwarded to every script)")
    p.add_argument("--load-mode", default=None, choices=["bf16", "8bit", "4bit"],
                   help="how to load the backbone (forwarded)")
    p.add_argument("--extract-batch", type=int, default=None,
                   help="texts per forward pass during extraction (forwarded; "
                        "omit to keep each script's own default)")
    p.add_argument("--gpu", type=int, default=None,
                   help="CUDA device index (forwarded; omit to keep each "
                        "script's own default)")
    p.add_argument("--max-length-utterance", type=int, default=None,
                   help="--max-length forwarded to the four attention scripts "
                        "(per-utterance truncation)")
    p.add_argument("--max-length-transcript", type=int, default=None,
                   help="--max-length forwarded to mean_pool (whole-transcript "
                        "truncation)")
    p.add_argument("--seeds", type=int, nargs="+", default=[None],
                   help="one or more seeds; the whole sweep repeats per seed "
                        "(embedding caches are seed-independent, so only the "
                        "first seed pays extraction)")
    p.add_argument("--keep-going", action="store_true",
                   help="continue with the remaining scripts if one fails")
    args = p.parse_args()

    # flags common to every script (per-script length is added below)
    common = []
    for flag, value in [
        ("--model", args.model),
        ("--load-mode", args.load_mode),
        ("--extract-batch", args.extract_batch),
        ("--gpu", args.gpu),
    ]:
        if value is not None:
            common += [flag, str(value)]

    def length_for(script: str):
        if script in TRANSCRIPT_SCRIPTS:
            return args.max_length_transcript
        return args.max_length_utterance

    outcomes = []  # (label, "ok"/"FAILED", seconds)
    stop = False
    for seed in args.seeds:
        for script in SCRIPTS:
            cmd = [sys.executable, os.path.join(REPO_ROOT, script)] + common
            length = length_for(script)
            if length is not None:
                cmd += ["--max-length", str(length)]
            if seed is not None:
                cmd += ["--seed", str(seed)]
            label = script if seed is None else f"{script} (seed {seed})"
            print(f"\n{'=' * 70}\nrunning: {' '.join(cmd)}\n{'=' * 70}", flush=True)
            start = time.time()
            rc = subprocess.run(cmd, cwd=REPO_ROOT).returncode
            elapsed = time.time() - start
            outcomes.append((label, "ok" if rc == 0 else "FAILED", elapsed))
            if rc != 0 and not args.keep_going:
                print(f"\n{label} exited with code {rc}; stopping "
                      "(pass --keep-going to continue past failures)")
                stop = True
                break
        if stop:
            break

    total = len(SCRIPTS) * len(args.seeds)
    print(f"\n{'=' * 70}\nSWEEP SUMMARY\n{'=' * 70}")
    for label, status, elapsed in outcomes:
        print(f"  {status:>6}  {elapsed:8.1f}s  {label}")
    n_failed = sum(s == "FAILED" for _, s, _ in outcomes)
    n_skipped = total - len(outcomes)
    if n_skipped:
        print(f"  (skipped {n_skipped} run(s) after failure)")
    sys.exit(1 if n_failed else 0)


if __name__ == "__main__":
    main()