208 lines
7.2 KiB
Python
Executable File
208 lines
7.2 KiB
Python
Executable File
#!/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()
|