#!/usr/bin/env python3 import argparse import json import math from collections import defaultdict from pathlib import Path import torch import torch.nn.functional as F from transformers import AutoModelForCausalLM, AutoTokenizer def load_jsonl(path, max_items=0, num_shards=1, shard_id=0): rows = [] with open(path, "r", encoding="utf-8") as f: for idx, line in enumerate(f): if num_shards > 1 and idx % num_shards != shard_id: continue if max_items and len(rows) >= max_items: break rows.append(json.loads(line)) return rows def encode(tokenizer, text): return tokenizer.encode(text, add_special_tokens=False) @torch.inference_mode() def continuation_nll(model, tokenizer, prompt, continuation, max_length): prompt_ids = encode(tokenizer, prompt) cont_ids = encode(tokenizer, continuation) if not cont_ids: return None if len(cont_ids) >= max_length: cont_ids = cont_ids[: max_length - 1] prompt_ids = [] else: prompt_budget = max_length - len(cont_ids) prompt_ids = prompt_ids[-prompt_budget:] ids = prompt_ids + cont_ids labels = [-100] * len(prompt_ids) + cont_ids if len(ids) < 2: return None input_ids = torch.tensor([ids], device=model.device, dtype=torch.long) label_ids = torch.tensor([labels], device=model.device, dtype=torch.long) logits = model(input_ids).logits shift_logits = logits[:, :-1, :].contiguous() shift_labels = label_ids[:, 1:].contiguous() mask = shift_labels.ne(-100) if not mask.any(): return None log_probs = F.log_softmax(shift_logits, dim=-1) safe_labels = shift_labels.masked_fill(~mask, 0) token_log_probs = log_probs.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1) nll = -token_log_probs[mask].sum().item() tokens = int(mask.sum().item()) return {"nll": nll, "tokens": tokens} def score_ppl(model, tokenizer, text, max_length): ret = continuation_nll(model, tokenizer, "", text, max_length) if ret is None: return None byte_len = max(1, len(text.encode("utf-8"))) ret["nll_per_token"] = ret["nll"] / max(1, ret["tokens"]) ret["ppl"] = math.exp(min(50, ret["nll_per_token"])) ret["nll_per_byte"] = ret["nll"] / byte_len ret["bits_per_byte"] = ret["nll"] / (byte_len * math.log(2)) ret["bytes"] = byte_len return ret def score_mcq(model, tokenizer, prompt, choices, max_length): scores = [] for choice in choices: sep = "" if prompt.endswith((" ", "\n")) else " " ret = continuation_nll(model, tokenizer, prompt + sep, choice, max_length) if ret is None: scores.append({"sum_logprob": -float("inf"), "avg_logprob": -float("inf"), "tokens": 0}) continue sum_logprob = -ret["nll"] avg_logprob = sum_logprob / max(1, ret["tokens"]) scores.append( { "sum_logprob": sum_logprob, "avg_logprob": avg_logprob, "tokens": ret["tokens"], } ) pred_avg = max(range(len(scores)), key=lambda i: scores[i]["avg_logprob"]) pred_sum = max(range(len(scores)), key=lambda i: scores[i]["sum_logprob"]) return scores, pred_avg, pred_sum def mean(xs): xs = [x for x in xs if x is not None] return sum(xs) / len(xs) if xs else None def summarize(results): groups = defaultdict(list) for row in results: groups[row["category"]].append(row) groups["all"] = results out = {} for cat, rows in groups.items(): out[cat] = { "n": len(rows), "mcq_acc_avg_norm": mean([x["mcq_correct_avg_norm"] for x in rows]), "mcq_acc_sum": mean([x["mcq_correct_sum"] for x in rows]), "ppl_token_mean": mean([x["ppl"]["ppl"] for x in rows if x.get("ppl")]), "nll_per_token_mean": mean([x["ppl"]["nll_per_token"] for x in rows if x.get("ppl")]), "bits_per_byte_mean": mean([x["ppl"]["bits_per_byte"] for x in rows if x.get("ppl")]), "choice_tokens_mean": mean( [s["tokens"] for x in rows for s in x.get("mcq_scores", []) if s["tokens"]] ), } return out def main(): ap = argparse.ArgumentParser() ap.add_argument("--model", required=True) ap.add_argument("--benchmark", required=True) ap.add_argument("--out-dir", required=True) ap.add_argument("--model-label", default="") ap.add_argument("--max-items", type=int, default=0) ap.add_argument("--max-length", type=int, default=2048) ap.add_argument("--num-shards", type=int, default=1) ap.add_argument("--shard-id", type=int, default=0) ap.add_argument("--dtype", choices=["auto", "float16", "bfloat16", "float32"], default="bfloat16") args = ap.parse_args() dtype = { "auto": "auto", "float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32, }[args.dtype] out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) label = args.model_label or Path(args.model).name out_label = label if args.num_shards > 1: if args.shard_id < 0 or args.shard_id >= args.num_shards: raise ValueError("--shard-id must be in [0, --num-shards)") out_label = f"{label}.shard{args.shard_id:02d}of{args.num_shards:02d}" tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( args.model, torch_dtype=dtype, device_map={"": 0}, trust_remote_code=True, ) model.eval() rows = load_jsonl(args.benchmark, args.max_items, args.num_shards, args.shard_id) results = [] for i, row in enumerate(rows, 1): ppl = score_ppl(model, tokenizer, row["ppl_text"], args.max_length) mcq_scores, pred_avg, pred_sum = score_mcq( model, tokenizer, row["mcq_prompt"], row["choices"], args.max_length ) result = { "id": row["id"], "category": row["category"], "source": row.get("source", ""), "answer_idx": row["answer_idx"], "pred_avg_norm": pred_avg, "pred_sum": pred_sum, "mcq_correct_avg_norm": int(pred_avg == row["answer_idx"]), "mcq_correct_sum": int(pred_sum == row["answer_idx"]), "ppl": ppl, "mcq_scores": mcq_scores, } results.append(result) if i % 100 == 0: print(f"[{label}] evaluated {i}/{len(rows)}") summary = { "model": args.model, "model_label": label, "benchmark": args.benchmark, "max_items": args.max_items, "max_length": args.max_length, "num_shards": args.num_shards, "shard_id": args.shard_id, "summary": summarize(results), } with (out_dir / f"{out_label}.per_item.jsonl").open("w", encoding="utf-8") as f: for row in results: f.write(json.dumps(row, ensure_ascii=False) + "\n") with (out_dir / f"{out_label}.summary.json").open("w", encoding="utf-8") as f: json.dump(summary, f, ensure_ascii=False, indent=2) print(json.dumps(summary, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()