| import argparse |
| import json |
| import zipfile |
| from pathlib import Path |
|
|
| import torch |
| from tqdm import tqdm |
| from transformers import AutoTokenizer, Gemma4ForCausalLM |
|
|
| MAX_PARAMS = 8_000_000_000 |
| MAX_BATCH_TOKENS = 36864 |
| LABELS = ["1", "2", "3", "4"] |
| ANSWER_CUE = "පිළිතුර:" |
| THOUGHT_OPEN = "<|channel>thought\n" |
| THOUGHT_CLOSE = "\n<channel|>" |
| STOP_TOKENS = ["<channel|>", "<eos>", "<turn|>"] |
|
|
|
|
| def normalise(question): |
| metadata = question.get("metadata", {}) |
| question = { |
| **question, |
| "subject_original": metadata.get("subject_original", question.get("subject_original")), |
| "level": metadata.get("difficulty", question.get("level")), |
| } |
| answer = question.get("answer") |
| if answer is not None and not (isinstance(answer, int) and 1 <= answer <= len(question["choices"])): |
| texts = [choice.strip() for choice in question["choices"]] |
| question["answer"] = texts.index(str(answer).strip()) + 1 |
| return question |
|
|
|
|
| def load_questions(data_dir): |
| questions = [] |
| for path in sorted(Path(data_dir).rglob("*.json")): |
| if any(part.startswith(".") for part in path.relative_to(data_dir).parts): |
| continue |
| questions.extend(normalise(q) for q in json.loads(path.read_text(encoding="utf-8"))) |
| return questions |
|
|
|
|
| def user_prompt(question): |
| lines = [ |
| "පහත බහුවරණ ප්රශ්නයට නිවැරදි පිළිතුරේ අංකය පමණක් ලියන්න.", |
| "", |
| f"විෂයය: {question['subject_original']}", |
| f"ප්රශ්නය: {question['question'].strip()}", |
| *(f"{i}. {choice}" for i, choice in enumerate(question["choices"], 1)), |
| ANSWER_CUE, |
| ] |
| return "\n".join(lines) |
|
|
|
|
| def render(tokenizer, prompt, think): |
| if tokenizer.chat_template is None: |
| return tokenizer.bos_token + prompt + " " |
| messages = [{"role": "user", "content": prompt}] |
| if think: |
| text = tokenizer.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True, enable_thinking=True |
| ) |
| return text + THOUGHT_OPEN |
| return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
|
|
|
|
| def thinking_text(tokenizer, prompt, thought): |
| return render(tokenizer, prompt, True) + thought + THOUGHT_CLOSE + ANSWER_CUE |
|
|
|
|
| DTYPES = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32} |
|
|
|
|
| def placement(dtype_name="auto"): |
| if not torch.cuda.is_available(): |
| return torch.float32, "cpu" |
| if dtype_name != "auto": |
| return DTYPES[dtype_name], "auto" |
| return (torch.float16 if torch.cuda.get_device_capability()[0] < 8 else torch.bfloat16), "auto" |
|
|
|
|
| def load_model(name, adapter=None, revision=None, dtype_name="auto"): |
| tokenizer = AutoTokenizer.from_pretrained(name, revision=revision, padding_side="left") |
| dtype, device_map = placement(dtype_name) |
| model, info = Gemma4ForCausalLM.from_pretrained( |
| name, |
| revision=revision, |
| dtype=dtype, |
| device_map=device_map, |
| key_mapping={r"^model\.language_model\.": "model."}, |
| output_loading_info=True, |
| ) |
| assert not info["missing_keys"], f"weights missing from checkpoint: {sorted(info['missing_keys'])[:5]}" |
| n_adapter = 0 |
| if adapter is not None: |
| from peft import PeftModel |
|
|
| model = PeftModel.from_pretrained(model, adapter) |
| n_adapter = sum(p.numel() for name, p in model.named_parameters() if "lora_" in name) |
| assert n_adapter > 0, f"adapter {adapter} has no LoRA parameters" |
| model = model.merge_and_unload() |
| model.eval() |
| n_params = sum(p.numel() for p in model.parameters()) + n_adapter |
| print( |
| f"Loaded {name} (text-only) on {model.device} in {model.dtype}: {n_params:,} parameters" |
| f" ({n_adapter:,} adapter, merged)" |
| ) |
| assert n_params <= MAX_PARAMS, f"{n_params:,} parameters exceeds the 8B limit" |
| return tokenizer, model |
|
|
|
|
| def label_ids(tokenizer): |
| ids = [tokenizer.encode(label, add_special_tokens=False) for label in LABELS] |
| assert all(len(i) == 1 for i in ids), f"answer labels are not single tokens: {ids}" |
| return [i[0] for i in ids] |
|
|
|
|
| def batches(lengths, budget, batch_size): |
| batch = [] |
| for i in sorted(range(len(lengths)), key=lambda i: -lengths[i]): |
| if batch and ( |
| len(batch) == batch_size or (len(batch) + 1) * (lengths[batch[0]] + budget) > MAX_BATCH_TOKENS |
| ): |
| yield batch |
| batch = [] |
| batch.append(i) |
| if batch: |
| yield batch |
|
|
|
|
| def encode(tokenizer, texts, device): |
| return tokenizer(texts, return_tensors="pt", padding=True, add_special_tokens=False).to(device) |
|
|
|
|
| def score(model, tokenizer, texts, n_choices, batch_size): |
| labels = label_ids(tokenizer) |
| lengths = [len(tokenizer.encode(t, add_special_tokens=False)) for t in texts] |
| probs = [None] * len(texts) |
| for batch in tqdm(list(batches(lengths, 0, batch_size)), desc="score"): |
| inputs = encode(tokenizer, [texts[i] for i in batch], model.device) |
| position_ids = (inputs["attention_mask"].cumsum(-1) - 1).clamp(min=0) |
| with torch.inference_mode(): |
| logits = model(**inputs, position_ids=position_ids, logits_to_keep=1).logits[:, -1, labels].float() |
| if not torch.isfinite(logits).all(): |
| raise RuntimeError("non-finite answer logits") |
| counts = torch.tensor([n_choices[i] for i in batch], device=logits.device) |
| valid = torch.arange(len(labels), device=logits.device)[None, :] < counts[:, None] |
| for i, row in zip(batch, logits.masked_fill(~valid, float("-inf")).softmax(-1).tolist()): |
| probs[i] = row[: n_choices[i]] |
| return probs |
|
|
|
|
| def think(model, tokenizer, contexts, budget, batch_size): |
| stops = tokenizer.convert_tokens_to_ids(STOP_TOKENS) |
| lengths = [len(tokenizer.encode(t, add_special_tokens=False)) for t in contexts] |
| thoughts = [None] * len(contexts) |
| for batch in tqdm(list(batches(lengths, budget, batch_size)), desc="think"): |
| inputs = encode(tokenizer, [contexts[i] for i in batch], model.device) |
| with torch.inference_mode(): |
| out = model.generate( |
| **inputs, |
| max_new_tokens=budget, |
| do_sample=False, |
| eos_token_id=stops, |
| pad_token_id=tokenizer.pad_token_id, |
| ) |
| for i, row in zip(batch, out[:, inputs["input_ids"].shape[1]:].tolist()): |
| end = next((k for k, token in enumerate(row) if token in stops), len(row)) |
| thoughts[i] = { |
| "thought": tokenizer.decode(row[:end], skip_special_tokens=True).strip(), |
| "finished": end < len(row) and row[end] == stops[0], |
| "thought_tokens": end, |
| } |
| return thoughts |
|
|
|
|
| def argmax(values): |
| return max(range(len(values)), key=values.__getitem__) |
|
|
|
|
| def answer(model, tokenizer, prompts, n_choices, budget, batch_size, score_batch_size=None): |
| if budget == 0: |
| texts = [render(tokenizer, p, False) for p in prompts] |
| thoughts = [{} for _ in prompts] |
| else: |
| contexts = [render(tokenizer, p, True) for p in prompts] |
| thoughts = think(model, tokenizer, contexts, budget, batch_size) |
| texts = [thinking_text(tokenizer, p, t["thought"]) for p, t in zip(prompts, thoughts)] |
| probs = score(model, tokenizer, texts, n_choices, score_batch_size or batch_size) |
| return [{**t, "probs": p, "pred": argmax(p) + 1} for t, p in zip(thoughts, probs)] |
|
|
|
|
| def order(n, rotation, reverse): |
| positions = list(range(n))[::-1] if reverse else list(range(n)) |
| k = rotation % n |
| return positions[k:] + positions[:k] |
|
|
|
|
| def rotate(question, rotation, reverse=False): |
| return {**question, "choices": [question["choices"][i] for i in order(len(question["choices"]), rotation, reverse)]} |
|
|
|
|
| def unrotate(probs, rotation, reverse=False): |
| original = [0.0] * len(probs) |
| for i, p in zip(order(len(probs), rotation, reverse), probs): |
| original[i] = p |
| return original |
|
|
|
|
| def vote(model, tokenizer, questions, rotations, budget, batch_size, score_batch_size=None, reverse=False): |
| prompts = [user_prompt(rotate(q, r, reverse)) for q in questions for r in rotations] |
| n_choices = [len(q["choices"]) for q in questions for _ in rotations] |
| runs = answer(model, tokenizer, prompts, n_choices, budget, batch_size, score_batch_size) |
| results = [] |
| for start in range(0, len(runs), len(rotations)): |
| group = [ |
| {**run, "rotation": r, "reverse": reverse, "probs": unrotate(run["probs"], r, reverse)} |
| for run, r in zip(runs[start : start + len(rotations)], rotations) |
| ] |
| group = [{**run, "pred": argmax(run["probs"]) + 1} for run in group] |
| probs = [sum(p) / len(group) for p in zip(*(run["probs"] for run in group))] |
| results.append({"pred": argmax(probs) + 1, "probs": probs, "runs": group}) |
| return results |
|
|
|
|
| def metrics(questions, results): |
| levels = {} |
| for question, result in zip(questions, results): |
| counts = levels.setdefault(question["level"], [0, 0]) |
| counts[0] += result["pred"] == question["answer"] |
| counts[1] += 1 |
| correct = sum(c for c, _ in levels.values()) |
| return { |
| "accuracy": correct / len(questions), |
| "correct": correct, |
| "total": len(questions), |
| "levels": {k: {"accuracy": c / t, "correct": c, "total": t} for k, (c, t) in sorted(levels.items())}, |
| } |
|
|
|
|
| def write_jsonl(path, rows): |
| path.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows), encoding="utf-8") |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model", default="google/gemma-4-E4B-it") |
| parser.add_argument("--revision", help="model commit to load; the latest one when not given") |
| parser.add_argument("--dtype", choices=["auto", *DTYPES], default="auto") |
| parser.add_argument("--data_dir", default="data/test") |
| parser.add_argument("--out_dir", default="outputs") |
| parser.add_argument("--think", type=int, default=512) |
| parser.add_argument("--adapter") |
| parser.add_argument("--limit", type=int) |
| parser.add_argument("--level", help="answer only the questions of this level, for example hard") |
| parser.add_argument("--batch_size", type=int, default=32) |
| parser.add_argument( |
| "--score_batch_size", type=int, |
| help="batch size for scoring the answer after the thought; --batch_size when not given, lower it for long thoughts", |
| ) |
| parser.add_argument( |
| "--rotations", type=int, nargs="+", default=[0], |
| help="answer once per rotation of the choices and average the answer probabilities; 0 2 votes over two orders", |
| ) |
| parser.add_argument( |
| "--reverse", action="store_true", |
| help="reverse the choices before rotating them, which gives 4 more orders of 4 choices", |
| ) |
| args = parser.parse_args() |
|
|
| questions = load_questions(args.data_dir) |
| if args.level: |
| questions = [q for q in questions if q["level"] == args.level] |
| questions = questions[: args.limit] |
| print(f"Loaded {len(questions)} questions from {args.data_dir}") |
| out_dir = Path(args.out_dir) |
| if not questions: |
| out_dir.mkdir(parents=True, exist_ok=True) |
| write_jsonl(out_dir / "predictions.jsonl", []) |
| write_jsonl(out_dir / "thoughts.jsonl", []) |
| return |
|
|
| tokenizer, model = load_model(args.model, args.adapter, args.revision, args.dtype) |
| if args.think and tokenizer.chat_template is None: |
| parser.error(f"{args.model} has no chat template, so it has no thinking mode; pass --think 0") |
|
|
| results = vote( |
| model, tokenizer, questions, args.rotations, args.think, args.batch_size, args.score_batch_size, args.reverse |
| ) |
| predictions = [ |
| { |
| "q_no": question["q_no"], |
| "subject": question["subject"], |
| "category": question["category"], |
| "level": question["level"], |
| "answer": LABELS[result["pred"] - 1], |
| } |
| for question, result in zip(questions, results) |
| ] |
|
|
| out_dir.mkdir(parents=True, exist_ok=True) |
| jsonl_path = out_dir / "predictions.jsonl" |
| write_jsonl(jsonl_path, predictions) |
| with zipfile.ZipFile(out_dir / "predictions.zip", "w", zipfile.ZIP_DEFLATED) as zf: |
| zf.write(jsonl_path, "predictions.jsonl") |
| print(f"Wrote {len(predictions)} predictions to {out_dir / 'predictions.zip'}") |
|
|
| runs = [(p, run) for p, r in zip(predictions, results) for run in r["runs"]] |
| thoughts = [ |
| {"q_no": p["q_no"], "subject": p["subject"], "level": p["level"]} |
| | {k: run[k] for k in ("rotation", "reverse", "pred", "probs", "finished", "thought_tokens", "thought") if k in run} |
| for p, run in runs |
| ] |
| write_jsonl(out_dir / "thoughts.jsonl", thoughts) |
| if args.think: |
| finished = sum(run["finished"] for _, run in runs) |
| print(f"Thoughts finished within {args.think} tokens: {finished}/{len(runs)}") |
|
|
| if questions and all(q.get("answer") is not None for q in questions): |
| report = metrics(questions, results) | {"think": args.think, "adapter": args.adapter} |
| if len(args.rotations) > 1: |
| report["rotations"] = { |
| str(r): metrics(questions, [res["runs"][i] for res in results])["accuracy"] |
| for i, r in enumerate(args.rotations) |
| } |
| if args.think: |
| report["finished"] = finished |
| (out_dir / "metrics.json").write_text(json.dumps(report, ensure_ascii=False, indent=1) + "\n", encoding="utf-8") |
| levels = ", ".join(f"{k} {v['accuracy']:.4f}" for k, v in report["levels"].items()) |
| print(f"Accuracy {report['accuracy']:.4f} ({report['correct']}/{report['total']}); {levels}") |
| if "rotations" in report: |
| print("Accuracy of each rotation alone:", report["rotations"]) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
