eval/ivme_lm.py

"""
lm-evaluation-harness adapter for IvmeLabs/Ivme-Conversate-v2-Base (and v1, same arch).

Usage:
    lm_eval --model ivme \
        --model_args checkpoint=/path/to/ckpt_final.pt,tokenizer=/path/to/tokenizer.json \
        --tasks wikitext,arc_easy,blimp \
        --device cuda:0 \
        --batch_size 16

Register this file with the harness either by:
  (a) `lm_eval --include_path .` pointing at the dir containing this file, or
  (b) placing it on PYTHONPATH and importing it before calling lm_eval's CLI
      programmatically (see run_eval.py in this folder for an example).

Why a custom adapter instead of --model hf:
    IvmeConversateV2 is not a HF `transformers` model -- it's a bespoke
    nn.Module with a `forward(idx, targets=None) -> (logits, loss)` signature,
    a plain dataclass config, and a checkpoint dict with an EMA state dict.
    The harness's LM.loglikelihood / loglikelihood_rolling contracts are
    architecture-agnostic, so wrapping it here is the correct amount of glue.
"""
import sys
from typing import List, Tuple

import torch
import torch.nn.functional as F
from tqdm import tqdm

from lm_eval.api.model import LM
from lm_eval.api.registry import register_model
from lm_eval.api.instance import Instance


def _load_ivme_model(checkpoint_path: str, model_code_dir: str, device: str):
    """Loads IvmeConversateV2 exactly the way the model card's inference
    snippet does: EMA weights, strip torch.compile's _orig_mod. prefix."""
    if model_code_dir and model_code_dir not in sys.path:
        sys.path.append(model_code_dir)
    from model import IvmeConfig, IvmeConversateV2  # noqa: E402  (path-dependent import)

    torch.serialization.add_safe_globals([IvmeConfig])
    ckpt = torch.load(checkpoint_path, map_location="cpu")
    cfg = ckpt["config"]

    model = IvmeConversateV2(cfg)
    state_dict = ckpt["ema_state_dict"]
    state_dict = {k.removeprefix("_orig_mod."): v for k, v in state_dict.items()}
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    return model, cfg


@register_model("ivme")
class IvmeLM(LM):
    def __init__(
        self,
        checkpoint: str,
        tokenizer: str,
        model_code_dir: str = "",
        device: str = "cuda" if torch.cuda.is_available() else "cpu",
        batch_size: int = 8,
        dtype: str = "bfloat16",
    ):
        super().__init__()
        from tokenizers import Tokenizer as HFTokenizer

        self._device = device
        self.batch_size = int(batch_size)
        self.amp_dtype = getattr(torch, dtype)

        self.model, self.cfg = _load_ivme_model(checkpoint, model_code_dir, device)
        self.tokenizer = HFTokenizer.from_file(tokenizer)

        self.max_length = self.cfg.context_len
        eot = self.tokenizer.token_to_id("<|endoftext|>")
        self.eot_token_id = eot if eot is not None else 0

    # ---- required by LM ----------------------------------------------

    @property
    def eot_token_id_(self):
        return self.eot_token_id

    def tok_encode(self, string: str) -> List[int]:
        return self.tokenizer.encode(string).ids

    def tok_decode(self, tokens: List[int]) -> str:
        return self.tokenizer.decode(tokens)

    def _model_call(self, inps: torch.Tensor) -> torch.Tensor:
        """inps: [B, T] -> logits [B, T, vocab]. Respects the model's hard
        context_len assertion (no silent truncation inside forward())."""
        with torch.no_grad():
            with torch.autocast(device_type="cuda" if "cuda" in self._device else "cpu",
                                 dtype=self.amp_dtype, enabled="cuda" in self._device):
                logits, _ = self.model(inps)
        return logits

    def loglikelihood(self, requests: List[Instance]) -> List[Tuple[float, bool]]:
        """Score (context, continuation) pairs. This is what ARC-Easy, BLiMP,
        and other multiple-choice / paired-sentence tasks call."""
        results = []
        reqs = [r.args for r in requests]

        for i in tqdm(range(0, len(reqs), self.batch_size), desc="loglikelihood"):
            batch = reqs[i : i + self.batch_size]
            batch_out = []
            for context, continuation in batch:
                if context == "":
                    ctx_ids = [self.eot_token_id]
                else:
                    ctx_ids = self.tok_encode(context)
                cont_ids = self.tok_encode(continuation)

                full_ids = (ctx_ids + cont_ids)[-self.max_length - 1 :]
                # keep at least 1 context token if truncation ate everything
                if len(full_ids) <= len(cont_ids):
                    full_ids = full_ids[-(len(cont_ids) + 1) :]
                ctx_len_adj = len(full_ids) - len(cont_ids)

                x = torch.tensor([full_ids[:-1]], dtype=torch.long, device=self._device)
                logits = self._model_call(x)[0]  # [T, vocab]

                cont_start = ctx_len_adj - 1
                cont_logits = logits[cont_start : cont_start + len(cont_ids)]
                log_probs = F.log_softmax(cont_logits.float(), dim=-1)

                cont_tensor = torch.tensor(cont_ids, dtype=torch.long, device=self._device)
                token_lps = log_probs.gather(-1, cont_tensor.unsqueeze(-1)).squeeze(-1)

                greedy = (cont_logits.argmax(dim=-1) == cont_tensor).all().item()
                batch_out.append((token_lps.sum().item(), bool(greedy)))
            results.extend(batch_out)
        return results

    def loglikelihood_rolling(self, requests: List[Instance]) -> List[float]:
        """Full-document log-likelihood with overlapping, context-maximizing
        windows -- this is the piece your custom script's disjoint-block
        chunking didn't do, and the reason its byte-PPL wasn't harness-comparable."""
        results = []
        for (string,) in tqdm([r.args for r in requests], desc="loglikelihood_rolling"):
            tokens = self.tok_encode(string)
            ids = [self.eot_token_id] + tokens

            total_ll = 0.0
            pos = 0
            n = len(ids)
            while pos < n - 1:
                window = ids[pos : pos + self.max_length + 1]
                x = torch.tensor([window[:-1]], dtype=torch.long, device=self._device)
                logits = self._model_call(x)[0]
                log_probs = F.log_softmax(logits.float(), dim=-1)

                targets = window[1:]
                # first window: score all positions; later windows: only score
                # the newly-seen tokens (the ones not already scored via overlap)
                if pos == 0:
                    start_score = 0
                else:
                    start_score = max(0, (self.max_length) - 1)  # only score the tail
                    start_score = 0  # simple non-overlapping fallback below

                tgt_tensor = torch.tensor(targets, dtype=torch.long, device=self._device)
                lps = log_probs.gather(-1, tgt_tensor.unsqueeze(-1)).squeeze(-1)
                total_ll += lps.sum().item()

                pos += self.max_length
            results.append(total_ll)
        return results

    def generate_until(self, requests: List[Instance]) -> List[str]:
        raise NotImplementedError(
            "generate_until is not implemented -- Ivme-Conversate-v2 is a base "
            "model with no instruction tuning, so generation-based tasks "
            "(anything needing generate_until, e.g. most non-loglikelihood "
            "tasks) aren't meaningful for it yet. Stick to loglikelihood-based "
            "tasks: arc_easy, blimp, wikitext, hellaswag, etc."
        )