File size: 3,365 Bytes
c6f3061 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """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")
# Validation is the default while developing sampling policies.
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)
|