scripts/build_submission.py

#!/usr/bin/env python
# build_submission.py — ONE command from the test set to a validated submission.
#
#   python scripts/build_submission.py --out submissions/tryanderror2.jsonl
#
# Everything is config-driven: the METHOD lives in config.json's "ensemble" block (v5+v3+v8 @ recent_k=40,
# the highest-score config = holdout 0.8018) + "dedup"/"gate"/"theta"; secrets (OPENROUTER_API_KEY for the
# VLM, VERTEX_KEY for the LLM dedup) live in .env. The script:
#   1. for EACH ensemble prompt, ensures a per-prompt decision cache exists — BUILDING it from the test
#      frames (data/phase1/data) if missing/incomplete (resumable, skips cached frames);
#   2. REPLAYS the caches through the SAME merge + entailment-dedup as the live policy (APIPolicy.step) —
#      identical output to `run_inference.py --policy api` (live ≡ replay: each prompt per frame, said=[],
#      gate-off, then dedup);
#   3. writes the JSONL and VALIDATES it (500 ids, times ×0.5).
# First run builds the caches (PAID, ~1× VLM call per prompt per frame); subsequent runs are ~free.
# The only manual step is naming the upload file exactly your Team ID.
import argparse
import glob
import json
import os
import sys
import threading
from concurrent.futures import ThreadPoolExecutor

import humomni.core.emb_cache as emb_cache
import humomni.core.judges as judges
from humomni.tuning.cache import VLMCache, frames_of
from humomni.phase1.dedup import EntailmentGuard
from humomni.core.validate_submission import validate, gold_ids_from_dir

_save_lock = threading.Lock()
_progress = [0]


def _checkpoint(every=20):
    """Persist the (LLM-dedup) judge cache every `every` samples so a Vertex 429 mid-build never
    re-pays for finished work — a resubmit resumes from the disk cache. The original one-off never
    saved these, which is why a fresh build floods Vertex with uncached entailment calls."""
    with _save_lock:
        _progress[0] += 1
        if _progress[0] % every == 0:
            judges.jc_save()


def build(data_dir, out, cfg, cache_tpl, workers, limit=0, build_workers=8):
    # Fail LOUDLY on a wrong/empty data_dir: an empty gold set would otherwise sail through
    # validation ("wrote 0 samples ... VALID") and produce an uploadable-looking empty file.
    if not any(os.path.isdir(x) for x in glob.glob(os.path.join(data_dir, "*"))):
        sys.exit(f"no sample dirs found under {data_dir!r} — expected the test set at "
                 f"data/phase1/data/<id>/ (frames 0.5.jpg, 1.0.jpg, ... + question.json), "
                 f"or pass --data_dir")
    ens = cfg.get("ensemble")
    if not ens:
        sys.exit("config.json has no 'ensemble' block — the best method is the v5+v3 ensemble; add it.")
    prompts = ens["prompts"]
    recent_k = ens.get("recent_k", cfg.get("dedup", {}).get("recent_k", 4))
    theta = cfg.get("theta", 0.6)
    d = cfg.get("dedup", {})

    # Optional perception gate (config.json "gate", use_gate=true): replayed BIT-IDENTICALLY to the live
    # phase1.perception_gate (same grayscale-novelty feature, same fire rule, runs BEFORE the VLM lookup).
    # Novelty per frame is computed once on CPU (free) and cached; gated-off frames emit nothing.
    g = cfg.get("gate", {})
    novelty = None
    if g.get("use_gate"):
        from humomni.tuning.cache import build_novelty
        nov_path = "cache/phase1_novelty.json"
        if not os.path.exists(nov_path):
            print(f"[gate] computing per-frame novelty for {data_dir} -> {nov_path} (CPU, one-off) …", flush=True)
            build_novelty(data_dir=data_dir, out=nov_path)   # ALWAYS full — a --limit-partial cache would poison later runs
        novelty = json.load(open(nov_path, encoding="utf-8"))
        print(f"[gate] ON — novelty>={g['novelty']}, max_silence={g['max_silence_s']}s, warmup={g['warmup_s']}s", flush=True)

    # Ensure a per-prompt decision cache for EACH ensemble prompt, BUILDING it from the test frames if
    # missing/incomplete (resumable — skips already-cached frames). This makes the whole pipeline a SINGLE
    # command from the test set to the submission; the config lives in config.json + secrets in .env.
    # Building calls the VLM (PAID, ~1× per prompt per frame); a complete cache is a no-op replay.
    from humomni.tuning.cache import build_web_cache
    from humomni.phase1.vlm_client import prompt_for
    caches = {}
    for p in prompts:
        path = cache_tpl.format(p)
        print(f"[cache] ensuring {path} (prompt {p}) from {data_dir} …", flush=True)
        build_web_cache(data_dir=data_dir, cache_path=path, limit=limit, workers=build_workers, prompt=prompt_for(p))
        caches[p] = VLMCache(path)

    emb_cache.load_disk()                              # prewarm draft embeddings (cosine dedup prefilter)
    judges.jc_load()                                   # reuse any saved LLM-dedup results (resume / re-run = free)
    warm = [v["draft"].strip() for c in caches.values() for v in c.mem.values() if v.get("draft", "").strip()]
    if warm:
        emb_cache.embed_many(warm)
    dirs = sorted(x for x in glob.glob(os.path.join(data_dir, "*")) if os.path.isdir(x))
    if limit:
        dirs = dirs[:limit]                            # quick correctness check (first N samples, same sort order)

    # Sample-level RESUME: completed records are appended to a `<out>.partial` sidecar as they finish,
    # so a killed build resumes instead of restarting (the dedup judge cache is already persisted
    # incrementally). The final `out` is written — and the sidecar dropped — only once all ids are done.
    partial = out + ".partial"
    done = {}
    if os.path.exists(partial):
        for line in open(partial, encoding="utf-8"):
            try:
                done[json.loads(line)["question_id"]] = line.rstrip("\n")
            except Exception:
                pass
    if done:
        print(f"[resume] {len(done)} samples already done in {partial}", flush=True)
    pw = open(partial, "a", encoding="utf-8")
    pwlock = threading.Lock()

    def one(sd):
        qid = json.load(open(os.path.join(sd, "question.json"), encoding="utf-8"))["question_id"]
        if qid in done:                                # resumed — skip, keep its earlier result
            return
        # one shared guard per sample == APIPolicy.step's dedup (sim band + LLM entailment, recent_k).
        guard = EntailmentGuard(sim_high=d.get("sim_high", 0.97), sim_low=d.get("sim_low", 0.55),
                                recent_k=recent_k, use_llm=d.get("use_llm", True))
        said, res, last_fire = [], [], -1e9
        for t, _ in frames_of(qid, data_dir):          # ASCENDING time — causal (frames ≤ t only)
            if novelty is not None:                    # gate replay == PerceptionGate.should_decide
                nov = novelty[qid][f"{t:.1f}"]
                fire = (t <= g["warmup_s"]) or (nov >= g["novelty"]) or (t - last_fire >= g["max_silence_s"])
                if not fire:
                    continue
                last_fire = t
            for p in prompts:                          # merge order = prompt order (v5 first, then v3)
                dec = caches[p].get(qid, t)
                if dec and dec["decision"] == "NEW_ANSWER" and dec["confidence"] >= theta:
                    draft = dec["draft"].strip()
                    if draft and not guard.is_duplicate(draft, said):
                        said.append(draft)
                        res.append({"time": round(t, 3), "content": draft})
        line = json.dumps({"question_id": qid, "model_response_list": res}, ensure_ascii=False)
        with pwlock:                                   # append + fsync the completed sample immediately
            pw.write(line + "\n"); pw.flush(); os.fsync(pw.fileno())
            done[qid] = line
        _checkpoint()

    try:
        with ThreadPoolExecutor(max_workers=workers) as ex:
            list(ex.map(one, dirs))
    finally:
        pw.close()
        judges.jc_save()                               # persist dedup results even on a mid-build 429
        emb_cache.save_disk()
    order = [json.load(open(os.path.join(sd, "question.json"), encoding="utf-8"))["question_id"] for sd in dirs]
    records = [json.loads(done[q]) for q in order if q in done]   # final output in dir order
    with open(out, "w", encoding="utf-8") as w:
        for r in records:
            w.write(json.dumps(r, ensure_ascii=False) + "\n")
    if len(records) == len(dirs) and os.path.exists(partial):
        os.remove(partial)                             # complete -> drop the sidecar
    return records


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Build + validate the best Phase-1 submission from per-prompt caches.")
    ap.add_argument("--data_dir", default="data/phase1/data")
    ap.add_argument("--out", default="submissions/submission.jsonl")
    ap.add_argument("--config", default="config.json")
    ap.add_argument("--cache_tpl", default="cache/phase1_{}.jsonl", help="per-prompt cache path ({} = prompt name)")
    ap.add_argument("--workers", type=int, default=1,
                    help="parallel samples. Keep at 1 on a COLD judge cache (the LLM dedup floods Vertex → 429); "
                         "raise to 3-4 once cached (re-run is then free).")
    ap.add_argument("--limit", type=int, default=0, help="only the first N samples (quick correctness check)")
    ap.add_argument("--build_workers", type=int, default=8, help="parallelism for BUILDING missing per-prompt VLM caches")
    a = ap.parse_args()
    cfg = json.load(open(a.config, encoding="utf-8"))
    os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)

    records = build(a.data_dir, a.out, cfg, a.cache_tpl, a.workers, a.limit, a.build_workers)
    avg = sum(len(r["model_response_list"]) for r in records) / max(1, len(records))
    ens = cfg["ensemble"]
    print(f"wrote {a.out}: {len(records)} samples, avg {avg:.1f} resp/sample "
          f"(ensemble {ens['prompts']} @ recent_k={ens.get('recent_k')})")

    errs = validate(a.out, gold_ids_from_dir(a.data_dir))   # same checks as validate_submission.py
    if errs:
        print(f"INVALID — {len(errs)} problem(s):")
        for e in errs[:20]:
            print("  -", e)
        sys.exit(1)
    print(f"VALID — all ids present, times are multiples of 0.5. Rename {a.out} to your Team ID to upload.")