#!/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()