#!/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()