Upload folder using huggingface_hub

eval/ivme_lm.py

@@ -0,0 +1,180 @@
+"""
+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."
+        )

eval/run_eval.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+"""
+Runs lm-evaluation-harness against Ivme-Conversate-v2 using the ivme_lm.py adapter.
+
+Example:
+    python run_eval.py \
+        --checkpoint /path/to/ckpt_final.pt \
+        --tokenizer /path/to/tokenizer.json \
+        --model_code_dir /path/to/model_folder \
+        --tasks wikitext,arc_easy,blimp \
+        --device cuda:0 \
+        --batch_size 16
+
+Notes:
+- `--model_code_dir` should point at the folder containing the model/ package
+  (config.py, transformer.py, etc.) -- i.e. what snapshot_download gave you,
+  or wherever you cloned the repo.
+- `blimp` here means the actual harness BLiMP group task, which covers the
+  real 67 paradigms with their real, correct config names on nyu-mll/blimp.
+  This replaces the paradigm list in the old custom script, which had several
+  fabricated/misspelled task names.
+- generate_until (free-form generation tasks) is not implemented in the
+  adapter since this is a non-instruction-tuned base model -- stick to
+  loglikelihood-based tasks (wikitext, arc_easy, blimp, hellaswag, piqa, etc.)
+"""
+import argparse
+
+import ivme_lm  # noqa: F401  (registers the "ivme" model with lm_eval on import)
+import lm_eval
+from lm_eval.utils import make_table
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", type=str, required=True)
+    parser.add_argument("--tokenizer", type=str, required=True)
+    parser.add_argument("--model_code_dir", type=str, default="")
+    parser.add_argument("--tasks", type=str, default="wikitext,arc_easy,blimp")
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--batch_size", type=int, default=16)
+    parser.add_argument("--limit", type=float, default=None,
+                         help="Optional: cap number of docs per task, for a quick sanity run first.")
+    args = parser.parse_args()
+
+    model_args = (
+        f"checkpoint={args.checkpoint},"
+        f"tokenizer={args.tokenizer},"
+        f"model_code_dir={args.model_code_dir},"
+        f"device={args.device},"
+        f"batch_size={args.batch_size}"
+    )
+
+    results = lm_eval.simple_evaluate(
+        model="ivme",
+        model_args=model_args,
+        tasks=args.tasks.split(","),
+        limit=args.limit,
+    )
+
+    print(make_table(results))
+    if "groups" in results and results["groups"]:
+        print(make_table(results, "groups"))
+
+
+if __name__ == "__main__":
+    main()