File size: 3,639 Bytes
7a2335a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """What does the model want to output after the prompt? DEV DATA ONLY (never the test set).
python -m sinmmlu.debug --model configs/models/gemma3_4b_it.yaml --n 5 [--prompt paper_subject_si] [--chat off] [--generate 60]
Per question: the tail of the exact scored text, the top-k next tokens, the digit probabilities and mass, the gold
answer, and optionally a greedy continuation.
"""
import argparse
from . import prompts as P
from .data import load_jsonl
from .evaluate import load_model_cfg
from .model import LM, label_token_ids, load_lm
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--data", default="dev/archive/easy_full.jsonl")
ap.add_argument("--prompt", default="paper_nosubject", choices=sorted(P.PROMPTS))
ap.add_argument("--chat", default="auto", choices=["auto", "on", "off"])
ap.add_argument("--n", type=int, default=5)
ap.add_argument("--start", type=int, default=0)
ap.add_argument("--topk", type=int, default=10)
ap.add_argument("--generate", type=int, default=0, help="greedy continuation tokens (0 = off)")
ap.add_argument("--suffix", default="", help="text appended after the rendered prompt")
ap.add_argument("--no-prefill", action="store_true", help="score without the answer pre-fill (as before 2026-10-24)")
ap.add_argument("--quiet", action="store_true", help="only the per-question label mass and a summary")
a = ap.parse_args()
assert a.data != "test", "debug output is for labelled dev data only"
import torch
cfg = load_model_cfg(a.model)
chat = cfg["chat_capable"] if a.chat == "auto" else a.chat == "on"
model, tok, info = load_lm(cfg["model"], cfg["revision"], cfg["dtype"], cfg["quant"])
lm = LM(model, tok, chat=chat, batch_size=1, context=info["context_length"])
labels = label_token_ids(tok, P.PROMPTS[a.prompt].labels)
dev = next(model.parameters()).device
print("label token ids:", dict(zip(P.PROMPTS[a.prompt].labels, labels)))
masses, correct = [], 0
for r in load_jsonl(a.data)[a.start:a.start + a.n]:
word = None if a.no_prefill else P.PROMPTS[a.prompt].answer_word
text = lm.render(P.build(r, a.prompt), word) + a.suffix
ids = lm.encode([text])[0]
x = torch.tensor([ids], device=dev)
with torch.no_grad():
logits = model(input_ids=x).logits[0, -1].float()
lp = torch.log_softmax(logits, -1)
top = torch.topk(lp, a.topk)
digit = [torch.logsumexp(lp[v], 0).exp().item() for v in labels]
masses.append(sum(digit))
correct += max(range(4), key=lambda j: digit[j]) + 1 == int(r["answer"])
if a.quiet:
continue
print(f"\n=== {r['uid']} gold={r['answer']} tokens={len(ids)}")
print("text tail:", repr(text[-160:]))
print("last tokens:", [tok.decode([i]) for i in ids[-4:]])
print("top next tokens:", [(tok.decode([i]), round(v.exp().item(), 3)) for v, i in zip(top.values, top.indices.tolist())])
print("labels:", [round(p, 3) for p in digit], "mass:", round(sum(digit), 3))
if a.generate:
with torch.no_grad():
g = model.generate(input_ids=x, attention_mask=torch.ones_like(x), max_new_tokens=a.generate, do_sample=False)
print("greedy:", repr(tok.decode(g[0, len(ids):], skip_special_tokens=True)))
print(f"SUMMARY prompt={a.prompt} chat={chat} suffix={a.suffix!r} n={len(masses)} "
f"mean_label_mass={sum(masses) / len(masses):.3f} min={min(masses):.3f} acc={correct / len(masses):.3f}")
if __name__ == "__main__":
main()
|