| """Run the Phase 3 notebook's complete ArithMark-3.0 evaluation protocol.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import time |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
| from datasets import load_dataset |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model", required=True, help="Local model directory or Hub id") |
| parser.add_argument("--output", required=True, type=Path) |
| parser.add_argument("--max-context", type=int, default=1024) |
| parser.add_argument("--progress-every", type=int, default=25) |
| parser.add_argument( |
| "--dtype", |
| choices=("float32", "bfloat16"), |
| default="float32", |
| help=( |
| "Forward precision. bfloat16 reproduces the Phase 3 notebook's " |
| "CUDA autocast protocol while keeping the stored model FP32." |
| ), |
| ) |
| return parser.parse_args() |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| if not torch.cuda.is_available(): |
| raise RuntimeError("This final benchmark run requires CUDA.") |
|
|
| device = torch.device("cuda") |
| tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained( |
| args.model, |
| trust_remote_code=True, |
| dtype=torch.float32, |
| ).to(device) |
| model.eval() |
|
|
| try: |
| dataset = load_dataset("AxiomicLabs/Arithmark-3.0", split="train") |
| except Exception: |
| dataset = load_dataset( |
| "AxiomicLabs/Arithmark-3.0", |
| data_files="arithmark-3.jsonl", |
| split="train", |
| ) |
|
|
| @torch.no_grad() |
| def score_continuation(context: str, ending: str) -> tuple[float, int]: |
| |
| |
| continuation_ids = tokenizer( |
| ending, add_special_tokens=False |
| )["input_ids"] |
| context_ids = tokenizer( |
| context, add_special_tokens=False |
| )["input_ids"] |
| context_ids = context_ids[-(args.max_context - len(continuation_ids)) :] |
|
|
| input_ids = torch.tensor( |
| [context_ids + continuation_ids], device=device, dtype=torch.long |
| ) |
| with torch.autocast( |
| device_type="cuda", |
| dtype=torch.bfloat16, |
| enabled=args.dtype == "bfloat16", |
| ): |
| output = model(input_ids) |
|
|
| logits = output["logits"] if isinstance(output, dict) else output.logits |
| if isinstance(logits, (tuple, list)): |
| logits = logits[0] |
| log_probs = F.log_softmax( |
| logits[0, len(context_ids) - 1 : -1].float(), dim=-1 |
| ) |
| targets = torch.tensor(continuation_ids, device=device, dtype=torch.long) |
| score = log_probs.gather(-1, targets[:, None]).sum().item() |
| return score, len(continuation_ids) |
|
|
| started = time.perf_counter() |
| raw_correct = 0 |
| normalized_correct = 0 |
| categories: dict[str, list[int]] = defaultdict(lambda: [0, 0, 0]) |
| samples: list[dict[str, object]] = [] |
|
|
| for index, example in enumerate(dataset): |
| label = int(example["label"]) |
| scored = [ |
| score_continuation(example["ctx"], ending) |
| for ending in example["endings"] |
| ] |
| raw_scores = [score for score, _ in scored] |
| normalized_scores = [score / max(length, 1) for score, length in scored] |
| raw_prediction = max(range(len(raw_scores)), key=raw_scores.__getitem__) |
| normalized_prediction = max( |
| range(len(normalized_scores)), key=normalized_scores.__getitem__ |
| ) |
| raw_hit = int(raw_prediction == label) |
| normalized_hit = int(normalized_prediction == label) |
| raw_correct += raw_hit |
| normalized_correct += normalized_hit |
|
|
| category = example["activity_label"] |
| categories[category][0] += raw_hit |
| categories[category][1] += normalized_hit |
| categories[category][2] += 1 |
| samples.append( |
| { |
| "index": index, |
| "label": label, |
| "raw_prediction": raw_prediction, |
| "normalized_prediction": normalized_prediction, |
| "raw_scores": raw_scores, |
| "normalized_scores": normalized_scores, |
| "continuation_lengths": [length for _, length in scored], |
| "activity_label": category, |
| } |
| ) |
|
|
| complete = index + 1 |
| if complete % args.progress_every == 0 or complete == len(dataset): |
| elapsed = time.perf_counter() - started |
| print( |
| f"ArithMark-3 {complete}/{len(dataset)} " |
| f"acc={100 * raw_correct / complete:.2f}% " |
| f"acc_norm={100 * normalized_correct / complete:.2f}% " |
| f"elapsed={elapsed:.1f}s", |
| flush=True, |
| ) |
|
|
| total = len(dataset) |
| elapsed = time.perf_counter() - started |
| result = { |
| "benchmark": "ArithMark-3.0", |
| "dataset": "AxiomicLabs/Arithmark-3.0", |
| "split": "train", |
| "model": args.model, |
| "stored_model_dtype": "float32", |
| "dtype": args.dtype, |
| "forward_precision": ( |
| "CUDA bfloat16 autocast" if args.dtype == "bfloat16" else "float32" |
| ), |
| "log_softmax_dtype": "float32", |
| "max_context": args.max_context, |
| "num_examples": total, |
| "limit": None, |
| "acc": raw_correct / total, |
| "acc_norm": normalized_correct / total, |
| "raw_correct": raw_correct, |
| "normalized_correct": normalized_correct, |
| "elapsed_seconds": elapsed, |
| "categories": { |
| category: { |
| "num_examples": counts[2], |
| "acc": counts[0] / counts[2], |
| "acc_norm": counts[1] / counts[2], |
| } |
| for category, counts in sorted(categories.items()) |
| }, |
| "samples": samples, |
| } |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| args.output.write_text(json.dumps(result, indent=2), encoding="utf-8") |
| print(json.dumps({key: result[key] for key in ( |
| "num_examples", "acc", "acc_norm", "elapsed_seconds" |
| )}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
