File size: 6,774 Bytes
d6f5237 | 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 | #!/usr/bin/env python3
"""eval.py: lm-evaluation-harness wrapper for MetaDiffusion-600M.
Scoring: single-step diffusion (mask the continuation, forward once at
t=1.0, log-prob of the true tokens at masked positions).
Usage (needs lm-eval in the environment):
python eval.py --checkpoint checkpoints/step_30000.pt \
--tasks hellaswag,arc_easy,arc_challenge,piqa \
--tokenizer data/tokenizer
"""
import argparse
import sys
from pathlib import Path
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer
from lm_eval.api.model import LM
from lm_eval.api.registry import register_model
sys.path.insert(0, str(Path(__file__).resolve().parent))
from model import MetaDiffusionLM, MetaDiffusionConfig # noqa: E402
import logging
logger = logging.getLogger(__name__)
@register_model("metadiffusion_600m")
class MetaDiffusion600MWrapper(LM):
def __init__(self, checkpoint: str, dtype: str = "float32",
device: str = "cuda", tokenizer_name: str = "Qwen/Qwen3-0.6B",
max_length: int = 1024, batch_size: int = 4, **kwargs):
super().__init__()
self._device = torch.device(device)
self._max_length = max_length
self._batch_size = batch_size
dtype_map = {"float32": torch.float32, "float16": torch.float16,
"bfloat16": torch.bfloat16}
self._dtype = dtype_map.get(dtype, torch.float32)
logger.info(f"Loading checkpoint: {checkpoint}")
ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
config = MetaDiffusionConfig(
**{k: v for k, v in ckpt.get("config", ckpt).items()
if k in MetaDiffusionConfig.__dataclass_fields__})
self.model = MetaDiffusionLM(config)
sd = ckpt.get("model_state_dict", ckpt)
sd = {k.replace("_orig_mod.", "", 1) if isinstance(k, str) and k.startswith("_orig_mod.") else k: v
for k, v in sd.items()}
self.model.load_state_dict(sd, strict=True)
self.model = self.model.to(device=self._device, dtype=self._dtype).eval()
self._mask_token_id = config.mask_token_id
self._pad_token_id = config.pad_token_id
self._tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
if self._tokenizer.pad_token_id is None:
self._tokenizer.pad_token_id = self._tokenizer.eos_token_id
self._eos_token_id = self._tokenizer.eos_token_id
logger.info(f"Loaded {config.num_hidden_layers}L x {config.hidden_size}W, "
f"vocab={config.mask_vocab_size}, {self._dtype}")
def _score_pair(self, context_tokens, continuation_tokens):
full_ids = context_tokens + continuation_tokens
if len(full_ids) > self._max_length:
excess = len(full_ids) - self._max_length
context_tokens = context_tokens[excess:] if len(context_tokens) > excess else []
full_ids = full_ids[excess:]
ctx_len = len(context_tokens)
seq_len = len(full_ids)
input_ids = torch.tensor([full_ids], device=self._device)
for i in range(ctx_len, seq_len):
input_ids[0, i] = self._mask_token_id
t = torch.tensor([1.0], device=self._device)
with torch.no_grad():
logits = self.model(input_ids, t)
log_probs = F.log_softmax(logits[0], dim=-1)
total = 0.0
is_greedy = True
for pos in range(ctx_len, seq_len):
true_token = full_ids[pos]
total += log_probs[pos, true_token].item()
if log_probs[pos].argmax().item() != true_token:
is_greedy = False
return total, is_greedy
def loglikelihood(self, requests, disable_tqdm=False):
results = []
for request in requests:
context, continuation = request.arguments
ctx = self._tokenizer.encode(context, add_special_tokens=False)
cont = self._tokenizer.encode(continuation, add_special_tokens=False)
if not cont:
cont = [self._eos_token_id]
results.append(self._score_pair(ctx, cont))
return results
def loglikelihood_rolling(self, requests, disable_tqdm=False):
results = []
for request in requests:
tokens = self._tokenizer.encode(request.arguments[0], add_special_tokens=False)
if len(tokens) <= 1:
results.append(0.0)
continue
lp, _ = self._score_pair(tokens[:1], tokens[1:])
results.append(lp)
return results
def generate_until(self, requests, disable_tqdm=False):
from chat import generate_response
results = []
for request in requests:
prompt = request.arguments[0]
prompt_ids = torch.tensor(
[self._tokenizer.encode(prompt, add_special_tokens=False)],
device=self._device)
x = generate_response(self.model, self._tokenizer, prompt_ids,
gen_len=128, num_steps=32, temperature=0.2,
repetition_penalty=1.2,
device=self._device, stop_on_end=True,
min_p=0.1)
out = x[0, prompt_ids.shape[1]:].cpu().tolist()
results.append(self._tokenizer.decode(out, skip_special_tokens=True))
return results
def main():
p = argparse.ArgumentParser(description="Evaluate MetaDiffusion-600M with lm-eval")
p.add_argument("--checkpoint", required=True)
p.add_argument("--tasks", default="hellaswag,arc_easy,arc_challenge,piqa")
p.add_argument("--tokenizer", default="data/tokenizer")
p.add_argument("--device", default="cuda:0")
p.add_argument("--dtype", default="bfloat16")
p.add_argument("--limit", default=None, help="Sample limit (smoke test)")
p.add_argument("--output", default=None, help="Result JSON path")
args = p.parse_args()
from lm_eval import simple_evaluate
results = simple_evaluate(
model="metadiffusion_600m",
model_args=f"checkpoint={args.checkpoint},dtype={args.dtype},"
f"device={args.device},tokenizer_name={args.tokenizer}",
tasks=args.tasks.split(","),
limit=float(args.limit) if args.limit is not None else None,
)
for task, res in results["results"].items():
acc = res.get("acc_norm,none") or res.get("acc,none")
print(f"{task}: {acc:.4f}" if acc is not None else f"{task}: {res}")
if args.output:
import json
with open(args.output, "w") as f:
json.dump(results["results"], f, indent=2)
print(f"Wrote {args.output}")
if __name__ == "__main__":
main()
|