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