eval/scripts/eval_coverage.py

#!/usr/bin/env python3
"""Full-coverage capability battery — FULL BUDGET, saved generations, per-type scorers.

One battery jsonl, every capability type, short + long items, no token caps. Each item:
  {"id","type","subtype":"short|long","prompt","scorer", ...scorer-args}
Scorers:
  math_exact      -> expected (final number match, ####/cue/lastnum extraction)
  code_exec       -> tests (+ entry_point); execute extracted function against unit tests
  sql_shape       -> expected_terms (all must appear, case-insens) — cheap structural check
  factual_terms   -> expected_terms (any must appear) — held-out topics only
  rubric          -> min_len + must_not_degenerate; saved for gatekeeper read (ok = passes mechanical rubric)
  bleed           -> expected_terms (instruction satisfied) AND no cross-family jargon (BLEED_RE)
Loader supports raw .pt (--config + build_model) or an HF model dir. KV-cached fill-context generation.
Usage: --ckpt --tokenizer --out [--config ...] [--battery data/eval_fixed/coverage/all.jsonl] [--max-new 0] [--device cuda]
"""
import sys, json, argparse, re, subprocess, tempfile, os
sys.path.insert(0, "scripts")

BLEED_RE = re.compile(
    r"\b(scope|scoped|scoping|dependency order|validation|stop condition|acceptance criteria|"
    r"tool use|tool-use|artifact|deliverable|implementation plan|agentic|route handler)\b", re.I)
ROLE_STOPS = ["\nUser:", "\nSystem:", "\nTool:", "\nAssistant:"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--tokenizer", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--config", default="config.json")
    ap.add_argument("--battery", default="data/eval_fixed/coverage/all.jsonl")
    ap.add_argument("--max-new", type=int, default=0, help="0 = fill context (default). >0 only for smoke.")
    ap.add_argument("--temperature", type=float, default=0.0, help="0 = greedy; >0 = sample")
    ap.add_argument("--top-p", type=float, default=1.0)
    ap.add_argument("--repetition-penalty", type=float, default=1.0, help=">1 discourages repeats (kills greedy loops)")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--device", default="cuda")
    a = ap.parse_args()
    import torch
    torch.manual_seed(a.seed)
    from transformers import AutoTokenizer
    hf = os.path.isdir(a.ckpt)
    if hf:
        from transformers import AutoModelForCausalLM
        model = AutoModelForCausalLM.from_pretrained(a.ckpt, torch_dtype=torch.float32).to(a.device).eval()
        ctx = int(json.load(open(os.path.join(a.ckpt, "config.json"))).get("max_position_embeddings", 4096))
    else:
        from pretrain_corpus import build_model, PretrainConfig
        cfg_d = json.load(open(a.config))
        cfg = PretrainConfig(**{k: v for k, v in cfg_d.items() if k in PretrainConfig.__dataclass_fields__})
        model = build_model(cfg, a.device)
        sd = torch.load(a.ckpt, map_location=a.device)
        model.load_state_dict(sd["model"] if "model" in sd else sd)
        model.eval(); model.to(a.device)
        ctx = int(cfg_d.get("block_size", 4096))
    tok = AutoTokenizer.from_pretrained(a.tokenizer)

    def gen(prompt):
        pids = tok(f"User:\n{prompt}\nAssistant:\n").input_ids
        budget = (ctx - len(pids) - 1) if a.max_new <= 0 else min(a.max_new, ctx - len(pids) - 1)
        ids = torch.tensor([pids]).to(a.device)
        gi = []; eos = False
        with torch.no_grad():
            out = model(input_ids=ids, use_cache=True, return_dict=True); past = out.past_key_values
            for _ in range(max(1, budget)):
                logits = out.logits[0, -1].float()
                if a.repetition_penalty != 1.0 and gi:
                    idx = torch.tensor(sorted(set(gi)), device=logits.device)
                    lg = logits[idx]
                    logits[idx] = torch.where(lg > 0, lg / a.repetition_penalty, lg * a.repetition_penalty)
                if a.temperature <= 0:
                    tid = int(torch.argmax(logits))
                else:
                    probs = torch.softmax(logits / a.temperature, -1)
                    if a.top_p < 1.0:
                        sp, si = torch.sort(probs, descending=True)
                        keep = (torch.cumsum(sp, -1) - sp) <= a.top_p
                        probs = torch.zeros_like(probs).scatter(0, si, sp * keep)
                        probs = probs / probs.sum()
                    tid = int(torch.multinomial(probs, 1))
                nxt = torch.tensor([[tid]], device=ids.device)
                if tid == tok.eos_token_id: eos = True; break
                gi.append(tid)
                if len(gi) % 8 == 0:
                    if any(s in tok.decode(gi[-24:], skip_special_tokens=True) for s in ROLE_STOPS): break
                    # loop-detector: stop on an exact-repeat cycle (genuine terminal degeneration, not a length cap)
                    if len(gi) >= 48 and gi[-24:] == gi[-48:-24]: break
                out = model(input_ids=nxt, past_key_values=past, use_cache=True, return_dict=True); past = out.past_key_values
        txt = tok.decode(gi, skip_special_tokens=True)
        for s in ROLE_STOPS:
            txt = txt.split(s, 1)[0]
        return txt.strip(), eos, len(gi)

    # ---- scorers ----
    def s_math(it, g):
        m = re.search(r"####\s*([-\d,\.]+)", g)
        if not m:
            cu = re.findall(r"(?:answer|final|total|result|equals?|is)\D{0,15}(-?\d[\d,]*\.?\d*)", g, re.I)
            pred = cu[-1] if cu else (re.findall(r"-?\d[\d,]*\.?\d*", g.replace(",", "")) or [None])[-1]
        else:
            pred = m.group(1)
        try: return abs(float(str(pred).replace(",", "")) - float(str(it["expected"]))) < 1e-6, pred
        except Exception: return False, pred

    def s_code(it, g):
        m = re.search(r"```(?:python)?\n(.*?)```", g, re.S)
        body = m.group(1) if m else g
        pc = it.get("prompt_code", "")
        prefix_lines = []
        for ln in pc.splitlines():
            stripped = ln.strip()
            if stripped.startswith(("import ", "from ")) or (
                stripped and not ln[:1].isspace() and "=" in stripped and "def " not in stripped
            ):
                prefix_lines.append(ln)
                continue
            if stripped.startswith(("def ", "async def ", "class ")):
                break
            if stripped:
                break
        prompt_prefix = ("\n".join(prefix_lines).rstrip() + "\n\n") if prefix_lines else ""
        cands = []
        starts_full = body.lstrip().startswith(("def ", "async def ", "class ", "import ", "from "))
        if "def " in body and starts_full:
            cands.append(body)                                   # model wrote the whole function
            if prompt_prefix and body.lstrip().startswith(("def ", "async def ", "class ")):
                cands.append(prompt_prefix + body)                # keep imports/type aliases from the prompt only
        else:
            cands.append(pc + body)                              # raw completion append
            indented = "\n".join(("    " + ln if ln.strip() and not ln[:1].isspace() else ln)
                                 for ln in body.splitlines())
            cands.append(pc + indented)                          # give the model its fairest shot on indentation
        wrap = (f"\ncheck({it['entry_point']})\n" if it.get("entry_point") else "\n")
        last = ""
        for cand in cands:
            prog = cand + "\n" + it["tests"] + wrap
            try:
                with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
                    f.write(prog); p = f.name
                r = subprocess.run([sys.executable, p], capture_output=True, timeout=12, text=True)
                os.unlink(p)
                if r.returncode == 0:
                    return True, "ok"
                last = (r.stderr or "")[-200:]
            except Exception as e:
                last = str(e)[-200:]
        return False, last

    def s_terms_all(it, g):
        lo = g.lower(); return all(t.lower() in lo for t in it["expected_terms"]), None

    def s_terms_any(it, g):
        lo = g.lower(); return any(t.lower() in lo for t in it["expected_terms"]), None

    def degenerate(g):
        lines = [x.strip() for x in g.splitlines() if x.strip()]
        if lines and max({l: lines.count(l) for l in set(lines)}.values()) / len(lines) > 0.4 and len(lines) > 4:
            return True
        words = g.split()
        return len(words) > 12 and len(set(words)) / len(words) < 0.35

    def s_rubric(it, g):
        ok = len(re.findall(r"[A-Za-z']+", g)) >= it.get("min_len", 30) and not degenerate(g)
        if it.get("expected_terms"):
            ok = ok and any(t.lower() in g.lower() for t in it["expected_terms"])
        return ok, None

    def s_bleed(it, g):
        bleed = sorted(set(m.group(0).lower() for m in BLEED_RE.finditer(g)))
        term_ok = (not it.get("expected_terms")) or any(t.lower() in g.lower() for t in it["expected_terms"])
        return (term_ok and not bleed), (",".join(bleed) if bleed else None)

    SC = {"math_exact": s_math, "code_exec": s_code, "sql_shape": s_terms_all,
          "factual_terms": s_terms_any, "rubric": s_rubric, "bleed": s_bleed}

    items = [json.loads(l) for l in open(a.battery) if l.strip()]
    rows = []; agg = {}
    for i, it in enumerate(items):
        g, eos, nt = gen(it["prompt"])
        ok, info = SC[it["scorer"]](it, g)
        t = it["type"]
        agg.setdefault(t, {"ok": 0, "n": 0, "trunc": 0})
        agg[t]["ok"] += int(ok); agg[t]["n"] += 1; agg[t]["trunc"] += int(not eos)
        rows.append({"id": it["id"], "type": t, "subtype": it.get("subtype", "short"),
                     "scorer": it["scorer"], "ok": ok, "info": info, "truncated": not eos,
                     "ntok": nt, "prompt": it["prompt"][:400], "gen": g})
        if (i + 1) % 20 == 0:
            print(f"  {i+1}/{len(items)}", flush=True)
    summary = {t: {"ok": v["ok"], "n": v["n"], "rate": round(v["ok"]/v["n"], 3),
                   "trunc": round(v["trunc"]/v["n"], 2)} for t, v in sorted(agg.items())}
    tot_ok = sum(v["ok"] for v in agg.values()); tot_n = sum(v["n"] for v in agg.values())
    os.makedirs(os.path.dirname(a.out), exist_ok=True)
    json.dump({"ckpt": a.ckpt, "ctx": ctx, "overall": {"ok": tot_ok, "n": tot_n},
               "by_type": summary, "rows": rows}, open(a.out, "w"), indent=1)
    print(f"\nOVERALL {tot_ok}/{tot_n}")
    for t, v in summary.items():
        print(f"  {t:26s} {v['ok']:3d}/{v['n']:<3d} = {100*v['rate']:5.1f}%  trunc={v['trunc']}")
    print("->", a.out)


if __name__ == "__main__":
    main()