Blame for code/run_model.py

#!/usr/bin/env python3
"""
Run the MSQ probe suite against one model and write raw results.

Usage:
    python code/run_model.py --model minicpm5-2b --out results/raw/minicpm5-2b.jsonl

Probes
------
ft_default      bare question, neutral system prompt            (free text)
ft_specified    context variable (school / stance / jurisdiction) made explicit
ft_two_schools  asked to enumerate and attribute the competing positions
mcq_neutral     multiple choice over the competing positions, unlabelled
mcq_specified   multiple choice over the competing positions, labelled by holder

For every prompt we record the model's own answer together with the mean
token log-probability of that answer, which is our confidence proxy.
"""
from __future__ import annotations

import argparse
import json
import pathlib
import sys
import time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
import msq  # noqa: E402

ROOT = pathlib.Path(__file__).resolve().parent.parent

MODELS = {
    "minicpm5-2b": "openbmb/MiniCPM5-2B",
    "minicpm5-2b-base": "openbmb/MiniCPM5-2B-Base",
    "minicpm5-2b-sft": "openbmb/MiniCPM5-2B-SFT",
    "qwen3-4b": "Qwen/Qwen3-4B-Instruct-2507",
    "smollm3-3b": "HuggingFaceTB/SmolLM3-3B",
    "phi4-mini": "microsoft/Phi-4-mini-instruct",
    "tiny-test": "hf-internal-testing/tiny-random-LlamaForCausalLM",
}


def load_model(name: str, load_in_4bit: bool, max_len: int):
    repo = MODELS[name]
    tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    tok.padding_side = "left"

    kwargs = dict(trust_remote_code=True, device_map="auto",
                  dtype=torch.bfloat16,
                  max_memory={0: "6GiB", "cpu": "8GiB"})
    if load_in_4bit:
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
        )
    model = AutoModelForCausalLM.from_pretrained(repo, **kwargs)
    model.eval()
    model.config.use_cache = True
    return tok, model, repo


def build_chat(tok, system: str, user: str) -> str:
    """Render a prompt.

    Instruct/chat models get their chat template (with reasoning disabled
    where the template supports it). Base models have no chat template, so they
    get a plain instruction prefix -- which is the fairest way to compare them,
    and is noted in the README.
    """
    has_template = bool(getattr(tok, "chat_template", None))
    if not has_template:
        prefix = (f"{system}\n\n" if system else "")
        return f"{prefix}### Instruction:\n{user}\n\n### Response:\n"
    msgs = []
    if system:
        msgs.append({"role": "system", "content": system})
    msgs.append({"role": "user", "content": user})
    try:
        return tok.apply_chat_template(msgs, tokenize=False,
                                       add_generation_prompt=True,
                                       enable_thinking=False)
    except Exception:
        return tok.apply_chat_template(msgs, tokenize=False,
                                       add_generation_prompt=True)


@torch.no_grad()
def generate_batch(tok, model, chats: list[str], max_new_tokens: int,
                   batch_size: int):
    outs = []
    for i in range(0, len(chats), batch_size):
        chunk = chats[i:i + batch_size]
        enc = tok(chunk, return_tensors="pt", padding=True, add_special_tokens=False)
        enc = {k: v.to(model.device) for k, v in enc.items()}
        gen = model.generate(**enc, max_new_tokens=max_new_tokens,
                             do_sample=False, temperature=None, top_p=None,
                             top_k=None, pad_token_id=tok.pad_token_id)
        new = gen[:, enc["input_ids"].shape[1]:]
        for row in new:
            ids = [t for t in row.tolist() if t != tok.pad_token_id]
            outs.append(tok.decode(ids, skip_special_tokens=True).strip())
        print(f"    gen {min(i + batch_size, len(chats))}/{len(chats)}", flush=True)
    return outs


@torch.no_grad()
def score_own_answers(tok, model, chats: list[str], answers: list[str],
                      batch_size: int):
    """Mean token log-probability of the model's own answer -> confidence proxy.

    The vocabulary is 130k wide, so materialising float logits for the whole
    sequence is what blows up the GPU. We therefore ask the model for only the
    logits we need (`logits_to_keep`) and convert one row at a time.
    """
    means = []
    for i in range(0, len(chats), batch_size):
        chunk_c = chats[i:i + batch_size]
        chunk_a = answers[i:i + batch_size]
        seqs, plens, glens = [], [], []
        for c, a in zip(chunk_c, chunk_a):
            pid = tok(c, add_special_tokens=False)["input_ids"]
            aid = tok(a, add_special_tokens=False)["input_ids"] if a else [tok.eos_token_id]
            seqs.append(pid + aid)
            plens.append(len(pid))
            glens.append(len(aid))
        maxlen = max(len(s) for s in seqs)
        min_plen = min(plens)
        pad = tok.pad_token_id
        inp = torch.full((len(seqs), maxlen), pad, dtype=torch.long)
        att = torch.zeros((len(seqs), maxlen), dtype=torch.long)
        for j, s in enumerate(seqs):
            inp[j, :len(s)] = torch.tensor(s)
            att[j, :len(s)] = 1
        inp, att = inp.to(model.device), att.to(model.device)
        keep = maxlen - min_plen + 1
        logits = model(input_ids=inp, attention_mask=att,
                       logits_to_keep=keep).logits
        for j, s in enumerate(seqs):
            tot, n = 0.0, 0
            for k in range(plens[j], plens[j] + glens[j]):
                row = logits[j, k - min_plen].float()
                lp = row.log_softmax(-1)[s[k]].item()
                tot += lp
                n += 1
            means.append(tot / max(n, 1))
        del logits
        print(f"    score {min(i + batch_size, len(chats))}/{len(chats)}", flush=True)
    return means


@torch.no_grad()
def letter_logits(tok, model, chats: list[str], n_opts: list[int],
                  batch_size: int):
    """Distribution over the option letters at the first generated position."""
    letter_ids = []
    for L in msq.LETTERS:
        cand = []
        for form in (L, " " + L):
            ids = tok(form, add_special_tokens=False)["input_ids"]
            if len(ids) == 1:
                cand.append(ids[0])
        letter_ids.append(sorted(set(cand)))
    out = []
    for i in range(0, len(chats), batch_size):
        chunk = chats[i:i + batch_size]
        enc = tok(chunk, return_tensors="pt", padding=True, add_special_tokens=False)
        enc = {k: v.to(model.device) for k, v in enc.items()}
        logits = model(**enc, logits_to_keep=1).logits[:, -1, :].float()
        for j in range(len(chunk)):
            k = n_opts[i + j]
            per = []
            for li in range(k):
                ids = letter_ids[li]
                per.append(float(torch.logsumexp(logits[j, ids], dim=0).item())
                           if ids else -1e9)
            mx = max(per)
            exps = [pow(2.718281828, v - mx) for v in per]
            s = sum(exps)
            out.append([e / s for e in exps])
        del logits
        print(f"    logit {min(i + batch_size, len(chats))}/{len(chats)}", flush=True)
    return out


@torch.no_grad()
def embed_texts(tok, model, texts: list[str], batch_size: int = 16) -> "torch.Tensor":
    """Mean-pooled, normalised last hidden states from the model's own body.

    Used to score how much of a ground-truth position an answer actually
    carries. The model paraphrases heavily ("tayammum is permissible when
    there is no water" for "tayammum is valid on dust, sand and rock"), so a
    lexical overlap measure is far too brittle. Comparing in the model's own
    representation space is robust to paraphrase and is, as a bonus, a direct
    inspection of what the model internally encodes.
    """
    base = getattr(model, "model", model)
    embs = []
    for i in range(0, len(texts), batch_size):
        chunk = [t if t else " " for t in texts[i:i + batch_size]]
        enc = tok(chunk, return_tensors="pt", padding=True, add_special_tokens=False,
                  truncation=True, max_length=512)
        enc = {k: v.to(model.device) for k, v in enc.items()}
        out = base(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"])
        h = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
        mask = enc["attention_mask"].unsqueeze(-1).to(h.dtype)
        pooled = (h * mask).sum(1) / mask.sum(1).clamp(min=1)
        pooled = torch.nn.functional.normalize(pooled.float(), dim=-1)
        embs.append(pooled.cpu())
    return torch.cat(embs)


@torch.no_grad()
def attach_position_sims(tok, model, items, records, batch_size: int = 16):
    """For every free-text record, the cosine similarity of the answer to each
    ground-truth position of its item."""
    pos_texts, pos_index = [], []
    for it in items:
        start = len(pos_texts)
        for p in it["positions"]:
            pos_texts.append(p["txt"])
        pos_index.append((it["id"], start, len(pos_texts)))
    if not pos_texts:
        return
    pos_emb = embed_texts(tok, model, pos_texts, batch_size)
    todo = [(i, r) for i, r in enumerate(records) if r["probe"].startswith("ft_")]
    if not todo:
        return
    ans_emb = embed_texts(tok, model, [r["answer"] or " " for _, r in todo], batch_size)
    pmap = {iid: (s, e) for iid, s, e in pos_index}
    for k, (i, r) in enumerate(todo):
        s, e = pmap[r["item_id"]]
        sims = (ans_emb[k].unsqueeze(0) @ pos_emb[s:e].T).squeeze(0).tolist()
        r["position_sims"] = [round(float(x), 4) for x in sims]
    print(f"    attached position sims for {len(todo)} answers", flush=True)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True, choices=list(MODELS))
    ap.add_argument("--out", required=True)
    ap.add_argument("--load-in-4bit", action="store_true")
    ap.add_argument("--batch-size", type=int, default=8)
    ap.add_argument("--max-new-tokens", type=int, default=280)
    ap.add_argument("--limit", type=int, default=0)
    args = ap.parse_args()

    items = msq.load_items()
    if args.limit:
        items = items[:args.limit]

    tok, model, repo = load_model(args.model, args.load_in_4bit, 4096)
    print(f"loaded {repo} | dtype={next(model.parameters()).dtype} | "
          f"vram={torch.cuda.memory_allocated()/1e9:.2f}GB", flush=True)

    records: list[dict] = []
    t0 = time.time()

    def add(it, probe, sysmsg, user, answer, conf, extra=None):
        rec = dict(model=args.model, repo=repo, probe=probe, item_id=it["id"],
                   cls=it["cls"], domain=it["domain"], sys=sysmsg, user=user,
                   answer=answer, mean_logprob=conf)
        if extra:
            rec.update(extra)
        records.append(rec)

    # ---- free-text probes ------------------------------------------------
    for probe, builder in (("ft_default", msq.ft_default_prompt),
                            ("ft_specified", msq.ft_specified_prompt),
                            ("ft_technical", msq.ft_technical_prompt),
                            ("ft_repair", msq.ft_repair_prompt)):
        pairs = [builder(it) for it in items]
        chats = [build_chat(tok, s, u) for s, u in pairs]
        print(f"[{probe}] {len(chats)} prompts", flush=True)
        answers = generate_batch(tok, model, chats, args.max_new_tokens,
                                 args.batch_size)
        confs = score_own_answers(tok, model, chats, answers, args.batch_size)
        for it, (s, u), a, c in zip(items, pairs, answers, confs):
            add(it, probe, s, u, a, c)

    # ---- two-schools / pair probes (class A only) ------------------------
    a_items = [it for it in items if it["cls"] == "A"]
    if a_items:
        for probe, builder in (("ft_two_schools", msq.ft_two_schools_prompt),
                               ("ft_pair", msq.ft_pair_prompt)):
            pairs = [builder(it) for it in a_items]
            chats = [build_chat(tok, s, u) for s, u in pairs]
            print(f"[{probe}] {len(chats)} prompts", flush=True)
            answers = generate_batch(tok, model, chats, args.max_new_tokens,
                                     args.batch_size)
            confs = score_own_answers(tok, model, chats, answers, args.batch_size)
            for it, (s, u), a, c in zip(a_items, pairs, answers, confs):
                add(it, probe, s, u, a, c)

    # ---- MCQ logit probes -------------------------------------------------
    for probe, builder in (("mcq_neutral", msq.mcq_neutral_prompt),
                            ("mcq_specified", msq.mcq_specified_prompt)):
        built = []
        for it in items:
            if probe == "mcq_specified" and it["cls"] == "D":
                continue
            built.append((it, builder(it)))
        chats = [build_chat(tok, s, u) for _, (s, u, _, _) in built]
        n_opts = [len(o) for _, (_, _, o, _) in built]
        print(f"[{probe}] {len(chats)} prompts", flush=True)
        dists = letter_logits(tok, model, chats, n_opts, args.batch_size)
        for (it, (s, u, opts, gold)), d in zip(built, dists):
            add(it, probe, s, u, None, None,
                {"options": [o["txt"] for o in opts],
                 "gold_index": gold, "probs": d,
                 "divided_index": msq.divided_index(it, opts),
                 "argmax_index": max(range(len(d)), key=lambda i: d[i]),
                 "max_prob": max(d)})

    outp = pathlib.Path(args.out)
    outp.parent.mkdir(parents=True, exist_ok=True)
    attach_position_sims(tok, model, items, records, args.batch_size)
    with outp.open("w", encoding="utf-8") as fh:
        for r in records:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"wrote {len(records)} records -> {outp} in {time.time()-t0:.0f}s")


if __name__ == "__main__":
    main()