| """ |
| 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) |
| _, 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() |
|
|
