| """Sampler ablations on fixed held-out subword masks; no selection on test.""" |
| import json |
| from pathlib import Path |
| import time |
| import numpy as np |
| import torch |
| from .data import load_arrays |
| from .inference import load, denoise, continue_text |
| from .model import PAD, BOS, EOS, corrupt |
|
|
|
|
| def run(args): |
| if args.examples < 1: |
| raise ValueError("examples must be positive") |
| torch.set_num_threads(4) |
| model, tokenizer = load(args.model, args.device) |
| arrays, manifest = load_arrays(args.data) |
| info = json.loads((Path(args.model) / "training_info.json").read_text()) |
| if info["manifest"] != manifest: |
| raise ValueError("Benchmark corpus differs from training") |
| |
| array = arrays[args.split][:args.examples] |
| clean = torch.tensor(np.array(array), dtype=torch.long) |
| noisy, selected, _ = corrupt(clean, torch.Generator().manual_seed(314), mode="random", rate=0.5) |
| schedules = [("one_pass", 1, "confidence"), ("left_to_right", 12, "left_to_right"), |
| ("confidence", 12, "confidence"), ("adaptive", 12, "adaptive")] |
| denoise(model, tokenizer, [BOS, 3, EOS], temperature=0, precision=args.precision) |
| results = [] |
| for name, steps, policy in schedules: |
| correct = total = exact = passes = 0 |
| latencies = [] |
| for truth, source, mask in zip(clean, noisy, selected): |
| size = int(truth.ne(PAD).sum()) |
| result = denoise(model, tokenizer, source[:size].tolist(), steps=steps, temperature=0, |
| policy=policy, precision=args.precision) |
| prediction = torch.tensor(result["tokens"]) |
| correct += int(prediction[mask[:size]].eq(truth[:size][mask[:size]]).sum()) |
| total += int(mask.sum()) |
| exact += int(prediction.eq(truth[:size]).all()) |
| passes += result["forward_passes"] |
| latencies.append(result["latency_ms"]) |
| results.append({"policy": name, "masked_accuracy": correct/total, "exact_match": exact/len(clean), |
| "mean_forward_passes": passes/len(clean), "mean_latency_ms": float(np.mean(latencies)), |
| "p95_latency_ms": float(np.percentile(latencies,95)), "masked_tokens": total}) |
| prompts = ["Once upon a time, a little girl found", "Tom opened the door and saw", "The woman stood on the balcony and", |
| "A small bird was afraid to fly. One day,"] |
| samples = [{"prompt": prompt, **continue_text(model,tokenizer,prompt,max_new_tokens=64,steps=12, |
| temperature=0.7,seed=42,precision=args.precision)} for prompt in prompts] |
| from .baselines import bigram_accuracy |
| baseline = bigram_accuracy(arrays["train"], noisy.numpy(), clean.numpy(), selected.numpy(), model.config.vocab_size) |
| report = {"split": args.split, "examples": len(clean), "mask_seed":314,"mask_rate":0.5,"temperature":0, |
| "device":args.device,"precision":args.precision,"checkpoint_step":info["step"], |
| "tokenizer_sha256":manifest["tokenizer_sha256"],"results":results,"bigram":baseline,"samples":samples} |
| output = Path(args.output) |
| output.parent.mkdir(parents=True,exist_ok=True) |
| output.write_text(json.dumps(report,indent=2,ensure_ascii=False),encoding="utf-8") |
| print(json.dumps(results,indent=2),flush=True) |
|
|
