| """Score multiple-choice medical questions with a small open model. |
| |
| Confidence = softmax over the option letters' next-token logits (no temperature). |
| Each question is scored in its original and in a shuffled option order. |
| |
| python scripts/score_models.py --model Qwen/Qwen2.5-1.5B-Instruct \ |
| --questions data/medqa.jsonl --out scores/qwen1.5b_medqa.npz [--nf4] |
| |
| questions.jsonl, one per line: {"question": str, "options": [str, ...], "answer": int} |
| Output .npz: conf, correct, pred, conf_shuf, correct_shuf, pred_shuf_in_original_order. |
| Needs: torch, transformers (+ bitsandbytes for --nf4). Runs on a single 16 GB GPU. |
| """ |
| import argparse |
| import json |
|
|
| import numpy as np |
|
|
| LETTERS = "ABCDEFGHIJ" |
|
|
|
|
| def prompt(q, options): |
| lines = [f"Question: {q}"] + [f"{LETTERS[i]}. {o}" for i, o in enumerate(options)] |
| return "\n".join(lines) + "\nAnswer:" |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--questions", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--nf4", action="store_true", help="load in 4-bit NF4") |
| ap.add_argument("--seed", type=int, default=0) |
| args = ap.parse_args() |
|
|
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| kw = dict(device_map="auto", torch_dtype=torch.float16) |
| if args.nf4: |
| from transformers import BitsAndBytesConfig |
| kw["quantization_config"] = BitsAndBytesConfig( |
| load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16) |
| tok = AutoTokenizer.from_pretrained(args.model) |
| model = AutoModelForCausalLM.from_pretrained(args.model, **kw).eval() |
|
|
| rows = [json.loads(l) for l in open(args.questions)] |
| rng = np.random.default_rng(args.seed) |
|
|
| def letter_ids(k): |
| return [tok.encode(" " + LETTERS[i], add_special_tokens=False)[-1] for i in range(k)] |
|
|
| @torch.no_grad() |
| def score(q, opts): |
| ids = tok(prompt(q, opts), return_tensors="pt").to(model.device) |
| logits = model(**ids).logits[0, -1].float() |
| p = torch.softmax(logits[letter_ids(len(opts))], -1).cpu().numpy() |
| return int(p.argmax()), float(p.max()) |
|
|
| out = {k: [] for k in ["pred", "conf", "correct", "pred_shuf", "conf_shuf", "correct_shuf"]} |
| for r in rows: |
| pr, c = score(r["question"], r["options"]) |
| out["pred"].append(pr); out["conf"].append(c); out["correct"].append(pr == r["answer"]) |
| perm = rng.permutation(len(r["options"])) |
| pr2, c2 = score(r["question"], [r["options"][j] for j in perm]) |
| orig = int(perm[pr2]) |
| out["pred_shuf"].append(orig); out["conf_shuf"].append(c2) |
| out["correct_shuf"].append(orig == r["answer"]) |
| np.savez(args.out, **{k: np.array(v) for k, v in out.items()}) |
| acc = np.mean(out["correct"]); agree = np.mean(np.array(out["pred"]) == np.array(out["pred_shuf"])) |
| print(f"{args.model}: accuracy {acc:.3f}, agreement under shuffling {agree:.3f}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
