File size: 6,288 Bytes
a2b8b4a | 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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | """
Stage 5 — EVALUATION (lecture: Evaluation, for pretraining AND post-training).
Compare base vs SFT vs DPO on three axes:
1. Perplexity on the pretraining validation set -> pure "language modeling" ability.
Usually rises a bit after SFT/DPO (the "alignment tax").
2. Instruction following on held-out prompts -> the goal of post-training.
Metrics: keyword hit rate (required words used), all-3-words rate, finish rate (story closed with EOS).
3. Reward-hacking detectors — things the reward does NOT see:
- fluency ppl: perplexity of the generated story under the frozen BASE model (lower = more natural text)
- mentions/word: how often each required word appears (~1-2 natural, higher = keyword stuffing)
- repeated trigrams: fraction of word trigrams that repeat (looping text)
Output: results/eval.json + results/eval.md (table + example generations).
Run: python evaluate.py
"""
from __future__ import annotations
import json
import math
import os
import numpy as np
import torch
from common import (autocast_ctx, ensure_dir, get_device, keyword_hit_rate, keyword_mentions, load_model,
repeated_trigram_rate, set_seed)
from config import get_config
from data import TokenDataset
from generation import sample_responses
from pretrain import evaluate_loss
from tokenizer import BPETokenizer
@torch.no_grad()
def fluency_nll(base_model, tok: BPETokenizer, outs: list[dict], device: str, batch_size: int = 32) -> float:
"""Mean per-token NLL of the generated stories under the frozen base model (story only, no prompt)."""
from dpo import seq_logprob
ctx, _ = autocast_ctx(device)
total, n_tok = 0.0, 0
resp = [o["ids"] for o in outs if o["ids"]]
for i in range(0, len(resp), batch_size):
b = resp[i:i + batch_size]
lp = seq_logprob(base_model, [[tok.eos_id]] * len(b), b, tok.eos_id, device, ctx)
total += -lp.sum().item()
n_tok += sum(len(r) for r in b)
return total / max(n_tok, 1)
def evaluate_model(model, tok, cfg, val_prompts, base_model, device, n_samples_kept: int = 3):
e = cfg.eval
val_ds = TokenDataset(os.path.join(cfg.data.data_dir, "val.bin"), model.cfg.context_len)
vl = evaluate_loss(model, val_ds, cfg.pretrain.batch_size, cfg.pretrain.eval_batches, device)
set_seed(cfg.seed) # identical sampling across models -> fair comparison
_, outs = sample_responses(model, tok, val_prompts, 1, e.gen_max_tokens, e.temperature, e.top_k, device)
outs = [o[0] for o in outs]
hits = [keyword_hit_rate(o["text"], w) for o, w in zip(outs, val_prompts)]
metrics = {
"val_loss": vl,
"val_ppl": math.exp(vl),
"keyword_hit_rate": float(np.mean(hits)),
"all_keywords_rate": float(np.mean([h == 1.0 for h in hits])),
"finish_rate": float(np.mean([o["finished"] for o in outs])),
"avg_tokens": float(np.mean([len(o["ids"]) for o in outs])),
"fluency_ppl": math.exp(fluency_nll(base_model, tok, outs, device)),
"mentions_per_word": float(np.mean([keyword_mentions(o["text"], w) for o, w in zip(outs, val_prompts)])),
"repeated_trigrams": float(np.mean([repeated_trigram_rate(o["text"]) for o in outs])),
}
samples = [{"words": w, "text": o["text"]} for w, o in zip(val_prompts[:n_samples_kept], outs[:n_samples_kept])]
return metrics, samples
def format_row(name: str, r: dict) -> str:
return (f"| {name} | {r['val_ppl']:.2f} | {r['keyword_hit_rate']:.1%} | {r['all_keywords_rate']:.1%} | "
f"{r['finish_rate']:.1%} | {r['fluency_ppl']:.2f} | {r['mentions_per_word']:.2f} | "
f"{r['repeated_trigrams']:.1%} | {r['avg_tokens']:.0f} |")
TABLE_HEADER = ["| Model | Val ppl ↓ | Keyword hit ↑ | All 3 words ↑ | Finished ↑ | Fluency ppl ↓ | "
"Mentions/word | Repeated trigrams ↓ | Avg tokens |",
"|---|---|---|---|---|---|---|---|---|"]
def load_val_prompts(cfg) -> list[list[str]]:
with open(os.path.join(cfg.data.data_dir, "sft_val.json")) as f:
return [x["words"] for x in json.load(f)][: cfg.eval.n_prompts]
def main() -> None:
cfg = get_config()
e = cfg.eval
device = get_device()
tok = BPETokenizer.load(os.path.join(cfg.data.data_dir, "tokenizer.json"))
val_prompts = load_val_prompts(cfg)
base_model, _ = load_model(cfg.pretrain.ckpt_path, device)
base_model.eval()
ckpts = {"base": cfg.pretrain.ckpt_path, "sft": cfg.sft.ckpt_path, "dpo": cfg.dpo.ckpt_path}
results, samples = {}, {}
for name, path in ckpts.items():
if not os.path.exists(path):
print(f"Skipping {name}: {path} not found")
continue
model, _ = load_model(path, device)
model.eval()
results[name], samples[name] = evaluate_model(model, tok, cfg, val_prompts, base_model, device)
r = results[name]
print(f"{name:5s} | ppl {r['val_ppl']:6.2f} | keyword hit {r['keyword_hit_rate']:.1%} | "
f"all-3-words {r['all_keywords_rate']:.1%} | finished {r['finish_rate']:.1%} | "
f"fluency ppl {r['fluency_ppl']:.2f} | mentions/word {r['mentions_per_word']:.2f} | "
f"rep-trigrams {r['repeated_trigrams']:.1%} | {r['avg_tokens']:.0f} tok")
del model
if device == "cuda":
torch.cuda.empty_cache()
ensure_dir("results/eval.json")
with open("results/eval.json", "w") as f:
json.dump({"n_prompts": len(val_prompts), "temperature": e.temperature, "top_k": e.top_k,
"metrics": results, "samples": samples}, f, indent=2)
lines = ["# Evaluation results", "",
f"Held-out prompts: {len(val_prompts)} | temperature {e.temperature} | top-k {e.top_k}", "",
*TABLE_HEADER]
lines += [format_row(n, r) for n, r in results.items()]
lines += ["", "## Example generations", ""]
for n, ss in samples.items():
lines.append(f"### {n}")
for s in ss:
lines += [f"**Required words:** {', '.join(s['words'])}", "", f"> {s['text'].strip()[:600]}", ""]
with open("results/eval.md", "w") as f:
f.write("\n".join(lines))
print("\nSaved: results/eval.json, results/eval.md")
if __name__ == "__main__":
main()
|