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