File size: 7,903 Bytes
840c3a4 | 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 203 204 205 206 | """Rescore existing unconditional samples with independent causal language models.
This leaves the canonical GPT-2 Large metrics untouched. Run from the repository
root with the dlb-eval environment and one free GPU.
"""
from __future__ import annotations
import argparse
from dataclasses import asdict
import json
import math
from pathlib import Path
from dlb.io import atomic_json_write, sha256_file, validate_samples
from evaluation.generative_perplexity import compute_gen_ppl
ROOT = Path(__file__).resolve().parents[1]
MODELS = {
"pythia-1b": (
"EleutherAI/pythia-1b",
"f73d7dcc545c8bd326d8559c8ef84ffe92fea6b2",
["*.json", "model.safetensors"],
),
"opt-1.3b": (
"facebook/opt-1.3b",
"3f5c25d0bc631cb57ac65913f76e22c2dfb61d62",
["*.json", "pytorch_model.bin", "merges.txt", "vocab.json"],
),
}
DATASETS = ("lm1b", "owt")
GENERATORS = ("fmlm", "duo_dcd", "duo_di4c", "mdlm_sdtt", "mdlm_di4c")
STEPS = (1, 2, 4)
SAMPLE_COUNT = 1024
CONTEXT_LENGTH = 1024
def scorer_class():
import torch
import torch.nn.functional as functional
class Float32LossScorer:
def __init__(self, model):
self.model = model
def eval(self):
self.model.eval()
return self
def score_batch(self, input_ids, attention_mask):
ids = torch.as_tensor(input_ids, dtype=torch.long, device="cuda")
mask = torch.as_tensor(attention_mask, dtype=torch.bool, device="cuda")
with torch.inference_mode():
logits = self.model(input_ids=ids, attention_mask=mask).logits
labels = ids[:, 1:].clone()
valid = mask[:, 1:]
labels.masked_fill_(~valid, -100)
loss = functional.cross_entropy(
logits[:, :-1, :].float().reshape(-1, logits.size(-1)),
labels.reshape(-1),
ignore_index=-100,
reduction="sum",
)
return float(loss.item()), int(valid.sum().item())
return Float32LossScorer
def load_model(model_key: str):
import torch
from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer
repo_id, revision, patterns = MODELS[model_key]
snapshot = Path(
snapshot_download(repo_id=repo_id, revision=revision, allow_patterns=patterns)
)
tokenizer = AutoTokenizer.from_pretrained(snapshot, local_files_only=True)
tokenizer.padding_side = "right"
tokenizer.truncation_side = "right"
if tokenizer.pad_token_id is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
snapshot, local_files_only=True, dtype=torch.float32
).to("cuda").eval()
return scorer_class()(model), tokenizer, snapshot
def sample_path(dataset: str, generator: str, steps: int) -> Path:
return ROOT / "results" / "samples" / dataset / generator / f"steps_{steps}" / "samples.jsonl"
def output_path(output_root: Path, model_key: str, dataset: str, generator: str, steps: int) -> Path:
return output_root / model_key / dataset / generator / f"steps_{steps}.json"
def valid_existing(path: Path, sample_sha: str, model_key: str) -> bool:
if not path.is_file():
return False
try:
row = json.loads(path.read_text(encoding="utf-8"))
ppl = row["generative_perplexity"]
return (
row["samples_sha256"] == sample_sha
and row["sample_count"] == SAMPLE_COUNT
and ppl["model_id"] == MODELS[model_key][0]
and ppl["model_revision"] == MODELS[model_key][1]
and ppl["forward_dtype"] == "float32"
and ppl["valid_token_count"] > 0
and math.isfinite(ppl["perplexity"])
)
except (KeyError, TypeError, ValueError, OSError):
return False
def main() -> None:
import torch
import transformers
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", choices=(*MODELS, "all"), default="all")
parser.add_argument("--output-root", type=Path, default=ROOT / "results" / "ppl_rescore_20260923")
parser.add_argument("--batch-size", type=int, default=8)
parser.add_argument("--smoke-samples", type=int, default=0)
args = parser.parse_args()
if args.batch_size < 1 or args.smoke_samples < 0:
parser.error("batch-size must be positive and smoke-samples must be nonnegative")
# Validate the entire input matrix before loading a scorer or writing results.
inputs = {}
for dataset in DATASETS:
for generator in GENERATORS:
for steps in STEPS:
path = sample_path(dataset, generator, steps)
validate_samples(path, expected=SAMPLE_COUNT)
inputs[(dataset, generator, steps)] = (path, sha256_file(path))
print(f"Validated {len(inputs)} sample files x {SAMPLE_COUNT} samples", flush=True)
model_keys = tuple(MODELS) if args.model == "all" else (args.model,)
for model_key in model_keys:
pending = [
(dataset, generator, steps)
for dataset in DATASETS
for generator in GENERATORS
for steps in STEPS
if args.smoke_samples or not valid_existing(
output_path(args.output_root, model_key, dataset, generator, steps),
inputs[(dataset, generator, steps)][1], model_key,
)
]
if not pending:
print(f"{model_key}: all {len(inputs)} outputs already validated", flush=True)
continue
scorer, tokenizer, snapshot = load_model(model_key)
print(f"{model_key}: loaded {snapshot}; pending={len(pending)}", flush=True)
for index, (dataset, generator, steps) in enumerate(pending, 1):
path, sample_sha = inputs[(dataset, generator, steps)]
with path.open(encoding="utf-8") as handle:
texts = [json.loads(line)["text"] for line in handle]
if args.smoke_samples:
texts = texts[:args.smoke_samples]
result = compute_gen_ppl(
texts, scorer, tokenizer, batch_size=args.batch_size,
max_length=CONTEXT_LENGTH,
model_revision=MODELS[model_key][1],
tokenizer_revision=MODELS[model_key][1],
)
metric = {
**asdict(result),
"model_id": MODELS[model_key][0],
"tokenizer_id": MODELS[model_key][0],
"model_snapshot": str(snapshot),
"tokenizer_snapshot": str(snapshot),
"forward_dtype": "float32",
"loss_dtype": "float32",
"torch_version": torch.__version__,
"transformers_version": transformers.__version__,
"offline_scoring": True,
}
print(
f"{model_key} {index}/{len(pending)} {dataset}/{generator}/{steps} "
f"samples={len(texts)} tokens={result.valid_token_count} "
f"ppl={result.perplexity:.6f}", flush=True,
)
if args.smoke_samples:
break
atomic_json_write(
output_path(args.output_root, model_key, dataset, generator, steps),
{
"schema_version": 1,
"evaluation": "independent_unconditional_ppl_rescore",
"samples": str(path),
"samples_sha256": sample_sha,
"sample_count": SAMPLE_COUNT,
"dataset": dataset,
"generator": generator,
"steps": steps,
"generative_perplexity": metric,
},
)
if __name__ == "__main__":
main()
|