| |
| """ |
| 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 |
|
|
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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() |
|
|
