eval.py

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