File size: 9,706 Bytes
784f572 | 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 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | #!/usr/bin/env python3
"""Evaluate public ELF OWT-prefix samples with the legacy 64+512 contract.
The ELF checkpoint still samples its native 1024-position T5 canvas. This
evaluator defines the comparable text-level contract: exactly 64 GPT-2 prompt
tokens and the first 512 GPT-2 response tokens per completion.
"""
from __future__ import annotations
import argparse
import csv
import json
from pathlib import Path
import statistics
import sys
from elf_common import ROOT, asset, manifest, offline, read_jsonl, require_server, sha256, write_json
from evaluate_elf import validate_records
def _load_prompt_rows(root: Path):
rows = read_jsonl(root / "data/conditional/owt-c64/prompts.jsonl")
prepared = read_jsonl(root / "data/elf/owt-prefix.jsonl")
by_id = {row["prompt_id"]: row for row in rows}
prepared_by_id = {row["id"]: row for row in prepared}
expected_ids = list(range(1024))
if len(rows) != 1024 or sorted(by_id) != expected_ids:
raise ValueError("canonical OWT C64 prompt file must contain 1024 unique prompt IDs")
if len(prepared) != 1024 or sorted(prepared_by_id) != expected_ids:
raise ValueError("prepared ELF OWT-prefix input must contain 1024 unique prompt IDs")
for prompt_id, row in by_id.items():
if len(row["prefix_token_ids"]) != 64:
raise ValueError(f"prompt_id={prompt_id} does not contain 64 canonical GPT-2 tokens")
row["prepared_input"] = prepared_by_id[prompt_id]["input"]
return by_id
def _tokenize_contract(rows, prompts, tokenizer, response_tokens: int):
"""Return (prompt IDs, response IDs, short IDs) for every saved completion."""
prompt_ids, response_ids, short_ids = [], [], []
for row in rows:
expected = prompts[row["prompt_id"]]
if row["input"] != expected["prepared_input"]:
raise ValueError(f"prompt_id={row['prompt_id']} input text differs from the prepared ELF input")
# Use the canonical source IDs for GPT-2 scoring. GPT-2 byte-BPE
# decode->encode is not guaranteed to preserve token boundaries, even
# when the decoded input text is unchanged.
prompt = list(expected["prefix_token_ids"])
response = tokenizer(
row["generated"], add_special_tokens=False,
truncation=True, max_length=response_tokens,
)["input_ids"]
prompt_ids.append(prompt)
response_ids.append(response)
if len(response) < response_tokens:
short_ids.append(row["id"])
return prompt_ids, response_ids, short_ids
def _score_ppl(prompt_ids, response_ids, model, tokenizer, batch_size, response_tokens):
from evaluation.conditional_perplexity import PromptExcludedCausalLMScorer
from evaluation.generative_perplexity import aggregate_nll
scorer = PromptExcludedCausalLMScorer(model, device="cuda").eval()
parts = []
for start in range(0, len(prompt_ids), batch_size):
p_batch = prompt_ids[start:start + batch_size]
r_batch = [row[:response_tokens] for row in response_ids[start:start + batch_size]]
combined = [p + r for p, r in zip(p_batch, r_batch, strict=True)]
width = max(map(len, combined))
ids = [row + [tokenizer.pad_token_id] * (width - len(row)) for row in combined]
attention = [[1] * len(row) + [0] * (width - len(row)) for row in combined]
targets = []
expected = 0
for p, r in zip(p_batch, r_batch, strict=True):
mask = [0] * len(p) + [1] * len(r)
mask += [0] * (width - len(mask))
targets.append(mask)
expected += len(r)
nll, count = scorer.score_prompt_excluded_batch(ids, attention, targets)
if count != expected:
raise ValueError(f"response target count mismatch: {count} != {expected}")
parts.append((nll, count))
return aggregate_nll(parts)
def main(argv=None):
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--root", type=Path, default=ROOT)
parser.add_argument("--input-root", type=Path, required=True,
help="quality root containing owt-prefix/c64-text-t5/steps_N/... runs")
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--steps", nargs="+", type=int, default=[1, 2, 4, 8, 16, 32, 1024])
parser.add_argument("--response-tokens", type=int, default=512)
parser.add_argument("--batch-size", type=int, default=8)
parser.add_argument("--allow-short", action="store_true",
help="score observed shorter responses; result is not strict 64+512")
args = parser.parse_args(argv)
if args.response_tokens < 2 or args.batch_size < 1:
parser.error("response-tokens must be >=2 and batch-size must be positive")
root = args.root.resolve()
input_root = args.input_root.resolve()
output = args.output.resolve()
sys.path[:0] = [str(root), str(root / "src")]
require_server()
offline()
from transformers import AutoModelForCausalLM, AutoTokenizer
from evaluation.unigram_entropy import unigram_entropy
prompts = _load_prompt_rows(root)
assets = manifest(root)["assets"]
gpt2 = AutoTokenizer.from_pretrained(
str(asset(root, "gpt2")), local_files_only=True,
clean_up_tokenization_spaces=False,
)
gpt2.pad_token = gpt2.eos_token
gpt2.padding_side = "right"
model = AutoModelForCausalLM.from_pretrained(
str(asset(root, "gpt2-large")), local_files_only=True
).cuda().eval()
output.mkdir(parents=True, exist_ok=True)
summary = []
for steps in args.steps:
run = input_root / "quality" / "owt-prefix" / "c64-text-t5" / f"steps_{steps}" / "seed_42" / "compile_0"
samples_path = run / "samples.jsonl"
if not samples_path.exists():
raise FileNotFoundError(f"missing generated samples: {samples_path}")
rows = read_jsonl(samples_path)
validate_records(rows, {"sample_count": len(rows), "prompt_count": 1024, "diversity": True})
if len(rows) != 2048:
raise ValueError(f"expected 2048 completions at {run}, found {len(rows)}")
prompt_ids, response_ids, short_ids = _tokenize_contract(rows, prompts, gpt2, args.response_tokens)
invalid = [] if args.allow_short else short_ids
result = {
"schema": "dlb-elf-owt-c64r512-v1",
"protocol": "owt_c64_r512_observed_v1" if args.allow_short else "owt_c64_r512_strict_v1",
"steps": steps,
"sample_count": len(rows),
"prompt_count": 1024,
"completion_schedule": "1024 primary + 256 prompts x 4 extra completions",
"prompt_tokenizer": "gpt2",
"prompt_tokens": 64,
"response_tokenizer": "gpt2",
"response_tokens_requested": args.response_tokens,
"response_short_count": len(short_ids),
"short_sample_ids": short_ids,
"source_run": str(run),
"samples_sha256": sha256(samples_path),
"generation_manifest_sha256": sha256(run / "generation.json"),
"assets_manifest_sha256": sha256(root / "data/elf/assets.json"),
"metrics": {},
"valid": not invalid,
"policy": "all 2048 completions; first 512 GPT-2 response tokens; prompt loss excluded",
"native_generation_note": "ELF samples a 1024-position T5 canvas before GPT-2 text retokenization",
"length_diagnostics": {
"response_gpt2_capped_mean": statistics.fmean(map(len, response_ids)),
"response_gpt2_capped_min": min(map(len, response_ids)),
"response_gpt2_capped_max": max(map(len, response_ids)),
},
}
if invalid:
result["reason"] = "strict 64+512 contract rejected short responses; no samples were dropped or padded"
else:
effective = [row[:args.response_tokens] for row in response_ids]
result["metrics"]["entropy_gpt2_nats"] = statistics.fmean(
unigram_entropy(row) for row in effective
)
ppl = _score_ppl(prompt_ids, response_ids, model, gpt2, args.batch_size, args.response_tokens)
result["metrics"]["conditional_ppl"] = {
"perplexity": ppl.perplexity,
"total_nll": ppl.total_nll,
"valid_token_count": ppl.valid_token_count,
"sample_count": len(rows),
"context_length": 64 + args.response_tokens,
"model_id": "gpt2-large",
"model_revision": assets["gpt2-large"]["revision"],
"tokenizer_id": "gpt2",
"tokenizer_revision": assets["gpt2"]["revision"],
}
destination = output / f"steps_{steps}" / "metrics.json"
write_json(destination, result)
summary.append({
"steps": steps, "sample_count": len(rows), "valid": result["valid"],
"short_count": len(short_ids),
"ppl": result["metrics"].get("conditional_ppl", {}).get("perplexity"),
"entropy": result["metrics"].get("entropy_gpt2_nats"),
"metrics_path": str(destination),
})
print(json.dumps(summary[-1], ensure_ascii=False), flush=True)
with (output / "summary.csv").open("w", newline="") as handle:
writer = csv.DictWriter(handle, fieldnames=list(summary[0]))
writer.writeheader()
writer.writerows(summary)
print(f"SUMMARY: {output / 'summary.csv'}", flush=True)
if any(not row["valid"] for row in summary):
raise SystemExit(2)
if __name__ == "__main__":
main()
|