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