#!/usr/bin/env python """ Subtitles-only paradigm. Usage: python run.py python run.py --qa-backbones gemini gpt claude qwen3vl glm4v """ from __future__ import annotations import argparse import os import sys import time from pathlib import Path _PROJECT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(_PROJECT)) from util.backbones import call_backbone, init_backbone from util.benchmark import load_benchmark from util.config import ( ANTHROPIC_API_KEY_ENV, DEFAULT_MODELS, GOOGLE_CLOUD_PROJECT_ENV, OPENAI_API_KEY_ENV, ) from util.eval_loop import run_sweep from util.parsers import parse_mcq_json from util.paths import ( BENCHMARK_PATH, RESULTS_DIR, ensure_dir, results_filename, ) from util.prompts import HEADER_SUBTITLES_ONLY, mcq_prompt_json from util.subtitles import load_subtitles, list_movies_with_subtitles, subtitles_coverage EXPERIMENT_NAME = 'subtitles_only' SUPPORTED_BACKBONES = ('gemini', 'gpt', 'claude', 'qwen3vl', 'qwen3', 'glm4v', 'glm4') DEFAULT_BACKBONES = ('gemini', 'gpt', 'claude', 'qwen3vl', 'glm4v') TEXT_ONLY_VARIANTS = ('qwen3', 'glm4') def _init_backbone_guard(backbone: str, model_override: str | None = None) -> None: """Initialize a backbone; secrets are loaded from .env automatically by util.backbones.""" from util.secrets import require if backbone == 'gpt': require(OPENAI_API_KEY_ENV) if backbone == 'claude': require(ANTHROPIC_API_KEY_ENV) kwargs = {} if model_override: kwargs['model_id'] = model_override if backbone == 'gemini': init_backbone('gemini', project_id=os.environ.get(GOOGLE_CLOUD_PROJECT_ENV), **kwargs) else: init_backbone(backbone, **kwargs) def _print_subtitle_inventory(bench_vids: list[str]) -> None: """Report subtitle coverage against the benchmark's vids.""" have, total = subtitles_coverage(bench_vids) on_disk = list_movies_with_subtitles() print(f'Subtitle inventory:') print(f' files on disk under data/subtitles/ : {len(on_disk)}') print(f' benchmark movies with a subtitle file : {have}/{total}') if have < total: missing = sorted(set(bench_vids) - set(on_disk)) print(f' {len(missing)} benchmark movies have NO subtitles — ' f'their questions will be skipped for this run.') for v in missing: print(f' {v}') def _dry_run(df, bench_vids: list[str], *, vid: str | None, n_questions: int) -> None: """Render the actual prompt(s) for a sample movie + question and exit. No backbone init, no API call. Use this to verify: - Subtitle preprocessing works (LANGS header dropped, text col extracted) - Prompt template is well-formed (no missing fields, sensible joining) - Length of subtitle text per movie (for context-window planning) """ SYSTEM_INSTRUCTION = ( "You answer movie MCQs based on the movie's subtitles only. " "Answer ONLY from the dialog text provided. Output valid JSON." ) # Pick the sample vid: caller-specified, else first bench movie with subtitles. if vid: target_vid = vid print(f'\n=== DRY RUN — vid={target_vid!r} (caller-specified) ===') else: on_disk = set(list_movies_with_subtitles()) candidates = [v for v in bench_vids if v in on_disk] if not candidates: print('\n=== DRY RUN — NO benchmark movies have subtitles on disk ===') print(f'Looked under: {SUBTITLES_DIR_HINT}') return target_vid = candidates[0] print(f'\n=== DRY RUN — vid={target_vid!r} (first bench movie with subs) ===') text = load_subtitles(target_vid) if not text: print(f'\n No subtitles loaded for {target_vid}. ' f'Either the file is missing or the format is unexpected.') return print(f'\nSubtitle text:') print(f' total chars = {len(text):,}') print(f' total words = {len(text.split()):,}') print(f' rough tokens ~ {len(text) // 4:,} (chars/4 heuristic; English)') preview_n = 200 print(f' first {preview_n} chars: {text[:preview_n]!r}') print(f' last {preview_n} chars: {text[-preview_n:]!r}') # Pick the first N questions for this movie. movie_qs = df[df['video_file_id'].astype(str).str.strip() == target_vid] if len(movie_qs) == 0: print(f'\n WARNING: vid {target_vid} has subtitles but no questions in benchmark.') return print(f'\n{len(movie_qs)} questions for this movie; rendering ' f'{min(n_questions, len(movie_qs))} prompt(s):') header = HEADER_SUBTITLES_ONLY.format(subtitles=text) for i, (_, row) in enumerate(movie_qs.head(n_questions).iterrows()): prompt = mcq_prompt_json(row.to_dict(), header=header) print() print('=' * 80) print(f'QUESTION {i + 1}/{n_questions} — qid={row.get("qid")!r}, ' f'gold={row.get("answer", row.get("gold", "?"))!r}') print('=' * 80) print(f'\n[system_instruction]') print(SYSTEM_INSTRUCTION) print(f'\n[user prompt]') print(prompt) print() print(f' prompt total chars = {len(prompt):,}') print(f' prompt rough tokens ~ {len(prompt) // 4:,}') print() print('Dry run complete. Pass --dry-run-n-questions N to render more, ' 'or --dry-run-vid VID to target a specific movie.') # Used by _dry_run to print a hint about where subtitles are looked up from util.paths import SUBTITLES_DIR as SUBTITLES_DIR_HINT def main() -> None: ap = argparse.ArgumentParser() ap.add_argument('--qa-backbones', nargs='+', default=list(DEFAULT_BACKBONES), choices=SUPPORTED_BACKBONES, help='One or more backbones to run QA with. ' f'Default: {" ".join(DEFAULT_BACKBONES)} ') ap.add_argument('--include-text-only-variants', action='store_true') ap.add_argument('--model-override', nargs='+', metavar='backbone=model_id', default=[], help='Override default model per backbone. ' 'e.g. --model-override claude=claude-haiku-4-5 gpt=gpt-5-mini.') ap.add_argument('--benchmark', default=str(BENCHMARK_PATH)) ap.add_argument('--max-output-tokens', type=int, default=24576) ap.add_argument('--temperature', type=float, default=0.0) ap.add_argument('--checkpoint-every', type=int, default=50) ap.add_argument('--dry-run', action='store_true', help='Build and print the prompt(s) for a sample movie + question, ' 'then exit. Skips all backbone init and API calls. Useful for ' 'verifying subtitle preprocessing + prompt formatting before ' 'kicking off something real.') ap.add_argument('--dry-run-vid', default=None, help='With --dry-run: which video_file_id to sample. Default: ' 'first benchmark movie with subtitles on disk.') ap.add_argument('--dry-run-n-questions', type=int, default=1, help='With --dry-run: number of questions to render prompts for.') args = ap.parse_args() if args.include_text_only_variants: for v in TEXT_ONLY_VARIANTS: if v not in args.qa_backbones: args.qa_backbones.append(v) print(f' --include-text-only-variants: extended to ' f'{" ".join(args.qa_backbones)}') df = load_benchmark(args.benchmark) bench_vids = sorted(df['video_file_id'].astype(str).str.strip().unique().tolist()) bench_vids = [v for v in bench_vids if v] print(f'\nBenchmark: {len(df)} questions across {len(bench_vids)} movies') _print_subtitle_inventory(bench_vids) if args.dry_run: _dry_run(df, bench_vids, vid=args.dry_run_vid, n_questions=args.dry_run_n_questions) return out_dir = ensure_dir(RESULTS_DIR / EXPERIMENT_NAME) print(f'\nSweeping: {len(args.qa_backbones)} backbones') overrides = {} for kv in args.model_override: if '=' not in kv: sys.exit(f'--model-override entries must be backbone=model_id, got: {kv!r}') bb, mid = kv.split('=', 1) overrides[bb.strip()] = mid.strip() if 'claude' in args.qa_backbones and 'claude' not in overrides: overrides['claude'] = DEFAULT_MODELS['claude'] print(f' (auto-selecting {DEFAULT_MODELS["claude"]}; pass ' '--model-override claude=claude-sonnet-4-5 to force Sonnet.)') args._overrides = overrides timings: list[tuple[str, str, float, int]] = [] overall_t0 = time.perf_counter() for backbone in args.qa_backbones: # Lazy init: never holds two local backbones in GPU memory at once, # and a single bad backbone doesn't block the rest of the sweep. print(f'--- Initializing backbone: {backbone} ---') try: _init_backbone_guard(backbone, overrides.get(backbone)) except Exception as e: print(f' init failed for {backbone}: {type(e).__name__}: {e!r}') print(f' skipping {backbone}; other backbones will continue.') timings.append((backbone, _model_id_for(backbone), 0.0, 0)) continue bb_t0 = time.perf_counter() model_id = _model_id_for(backbone) results_path = out_dir / results_filename(backbone=backbone, model_id=model_id) n_before = _count_results_csv(results_path) _run_one(df, backbone, args, out_dir) bb_dt = time.perf_counter() - bb_t0 n_after = _count_results_csv(results_path) timings.append((backbone, model_id, bb_dt, n_after - n_before)) overall_dt = time.perf_counter() - overall_t0 print() print('Backbone wall time (this invocation, excludes model-init):') for backbone, model_id, dt, n_new in timings: per_q = (dt / n_new) if n_new else float('nan') label = f'{backbone}/{model_id}' if len(label) > 42: label = label[:39] + '...' rate_str = f'{per_q:.2f}s/q' if n_new else 'no new rows' print(f' {label:42s} {dt:8.1f}s new={n_new:>4} {rate_str}') print(f" {'TOTAL':42s} {overall_dt:8.1f}s") _persist_run_summary(out_dir, args.qa_backbones, timings, overall_dt) def _count_results_csv(path: Path) -> int: """Row count in an existing results CSV, or 0 if missing/empty.""" if not path.exists() or path.stat().st_size == 0: return 0 try: import pandas as pd return len(pd.read_csv(path)) except Exception: return 0 def _persist_run_summary(out_dir: Path, backbones: list[str], timings: list[tuple[str, str, float, int]], overall_dt: float) -> None: """Append one row per backbone-run to runs.csv (append-only).""" import csv from datetime import datetime summary_path = out_dir / 'runs.csv' started_at = datetime.now().isoformat(timespec='seconds') new_rows = not summary_path.exists() with summary_path.open('a', newline='') as fh: w = csv.writer(fh) if new_rows: w.writerow(['run_started_at', 'backbone', 'model_id', 'wall_s', 'new_rows', 'overall_wall_s']) for backbone, model_id, dt, n_new in timings: w.writerow([started_at, backbone, model_id, f'{dt:.2f}', n_new, f'{overall_dt:.2f}']) print(f' Run summary appended -> {summary_path}') def _run_one(df, backbone: str, args, out_dir: Path) -> None: print(f'\n=== {EXPERIMENT_NAME}: backbone={backbone} ===') def build_movie_ctx(vid, movie_df): text = load_subtitles(vid) if not text: print(f' SKIP {vid}: no subtitles at data/subtitles/{vid}.speech.txt') return {} return {'subtitles': text} def call_fn(row, ctx): text = ctx.get('subtitles') if not text: return '' header = HEADER_SUBTITLES_ONLY.format(subtitles=text) prompt = mcq_prompt_json(row, header=header) raw = call_backbone( backbone, prompt, images=None, system_instruction=('You answer movie MCQs based on the movie\'s subtitles only. ' 'Answer ONLY from the dialog text provided. Output valid JSON.'), json_mode=True, max_output_tokens=args.max_output_tokens, temperature=args.temperature, ) return parse_mcq_json(raw, qid=row['qid']) model_id = _model_id_for(backbone) results_path = out_dir / results_filename(backbone=backbone, model_id=model_id) run_sweep( df, call_fn=call_fn, build_movie_ctx=build_movie_ctx, results_path=results_path, experiment_name=f'{EXPERIMENT_NAME} {backbone}', extra_row_fields={ 'backbone': backbone, 'model': _model_id_for(backbone), 'subtitles_source': 'whisper', }, checkpoint_every=args.checkpoint_every, ) def _model_id_for(backbone: str) -> str: from util.backbones import DEFAULT_MODELS, _configs return (_configs.get(backbone, {}) or {}).get('model_id') or DEFAULT_MODELS.get(backbone, backbone) if __name__ == '__main__': main()