run_lmeval.py

"""run_lmeval.py — reproduce min-spark's published benchmarks.

A LoglikelihoodLM (TemplateLM subclass) over MinSparkForCausalLM with the
repo's scoring rules. CLI: python run_lmeval.py --effort medium --tasks arc_easy,wikitext [--limit N]
"""
import argparse
import json
from pathlib import Path

import torch

from transformers import AutoModelForCausalLM, AutoTokenizer

from lm_eval.api.model import TemplateLM
from lm_eval.evaluator import simple_evaluate
from lm_eval.utils import get_rolling_token_windows

HERE = Path(__file__).resolve().parent
EFFORT_MAP = {"low": 2, "medium": 3, "high": 4}
MAX_LENGTH = 512


class MinSparkLM(TemplateLM):
    def __init__(self, effort: str, batch_size: int = 32, device: str = "cpu", max_length: int = MAX_LENGTH):
        super().__init__()
        self._model = AutoModelForCausalLM.from_pretrained(str(HERE), trust_remote_code=True).to(device).eval()
        # trust_remote_code=True even though the tokenizer is stock: config.json's
        # auto_map marks the repo as custom code, so without it the interactive
        # prompt pollutes stdout and breaks --json parsing under subprocess capture.
        self._tok = AutoTokenizer.from_pretrained(str(HERE), trust_remote_code=True)
        self._effort = effort
        self._batch_size = batch_size
        self._device = device
        self._max_length = max_length

    @property
    def eot_token_id(self) -> int:
        return 2  # EOS

    @property
    def max_length(self) -> int:
        return self._max_length

    @property
    def device(self): return self._device
    @property
    def rank(self): return 0
    @property
    def world_size(self): return 1
    @property
    def batch_size(self): return self._batch_size

    def tok_encode(self, string, left_truncate_len=None, add_special_tokens=None):
        ids = self._tok(string, add_special_tokens=False).input_ids
        if left_truncate_len is not None:
            ids = ids[-left_truncate_len:]
        return ids

    def tok_decode(self, tokens, skip_special_tokens=True):
        return self._tok.decode(tokens, skip_special_tokens=skip_special_tokens)

    def _loglikelihood_tokens(self, requests, disable_tqdm=False, override_use_cache=None):
        # Abstract in TemplateLM; unused (we override loglikelihood directly).
        raise NotImplementedError("MinSparkLM overrides loglikelihood directly")

    def _logits(self, inp: torch.Tensor) -> torch.Tensor:
        """Run the wrapper forward under inference_mode, on device, uncapped,
        float32."""
        with torch.inference_mode():
            out = self._model(inp, effort=self._effort).logits
        return out.float()

    def loglikelihood(self, requests, disable_tqdm: bool = False):
        items = [self._prepare_pair(*req.args) for req in requests]
        scored = self._score_batched(items)
        out = []
        for token_lp, target, argmax in scored:
            is_greedy = bool((argmax == target).all().item())
            out.append((float(token_lp.sum().item()), is_greedy))
        return out

    def _prepare_pair(self, context: str, continuation: str):
        """(context, continuation) -> (model_in ids, target ids) exactly as the
        unbatched scorer fed them: cross-merge split, EOS prefix iff empty ctx,
        tail-truncated to max_length+1."""
        context_enc = self.tok_encode(context)          # [] when context == ""
        whole_enc = self.tok_encode(context + continuation)
        continuation_enc = whole_enc[len(context_enc):]  # cross-merge-correct split
        if context == "":
            context_enc = [self.eot_token_id]            # prefix exactly once, empty ctx only
        full = context_enc + continuation_enc
        inp_ids = full[-(self._max_length + 1):]        # keep <= max_length+1
        return inp_ids[:-1], inp_ids[-len(continuation_enc):]

    def loglikelihood_rolling(self, requests, disable_tqdm: bool = False):
        items, owner = [], []                            # window -> owning request
        for ri, req in enumerate(requests):
            (string,) = req.args
            toks = self.tok_encode(string)
            for input_win, target_win in get_rolling_token_windows(
                token_list=toks,
                prefix_token=self.eot_token_id,        # prepended to first window only
                max_seq_len=self._max_length,
                context_len=1,
            ):
                items.append((input_win, target_win))
                owner.append(ri)
        window_lps = [None] * len(items)
        for wi, (token_lp, _, _) in enumerate(self._score_batched(items)):
            window_lps[wi] = float(token_lp.sum().item())
        totals = [0.0] * len(requests)
        for ri, lp in zip(owner, window_lps):            # original window order per doc
            totals[ri] += lp
        return totals

    def _score_batched(self, items):
        """items: list of (model_in ids, target ids); target positions are the
        LAST len(target) rows of each sequence's logits. Sequences are scored in
        length-sorted right-padded batches of self._batch_size — right padding
        cannot reach real positions (attention is causal). Returns, in items
        order: (per-token logprobs, target tensor, argmax over scored rows)."""
        order = sorted(range(len(items)), key=lambda i: len(items[i][0]), reverse=True)
        results = [None] * len(items)
        for start in range(0, len(order), self._batch_size):
            chunk = order[start:start + self._batch_size]
            lens = [len(items[i][0]) for i in chunk]
            assert max(lens, default=0) <= self._max_length
            batch = torch.zeros(len(chunk), max(lens), dtype=torch.long)
            for row, i in enumerate(chunk):
                batch[row, :lens[row]] = torch.tensor(items[i][0], dtype=torch.long)
            logprobs = torch.log_softmax(self._logits(batch.to(self._device)), dim=-1)
            for row, i in enumerate(chunk):
                target = torch.tensor(items[i][1], dtype=torch.long, device=logprobs.device)
                scored = logprobs[row, lens[row] - len(target):lens[row]]
                token_lp = scored.gather(-1, target.unsqueeze(-1)).squeeze(-1)
                results[i] = (token_lp, target, scored.argmax(-1))
        return results

    def generate_until(self, requests, disable_tqdm: bool = False):
        return [""] * len(requests)  # not used by the reproduction tasks


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--effort", choices=sorted(EFFORT_MAP), default="medium")
    ap.add_argument("--tasks", default="blimp,arc_easy,arc_challenge,hellaswag,piqa,wikitext")
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--batch-size", type=int, default=32)
    ap.add_argument("--json", action="store_true")
    args = ap.parse_args()

    lm = MinSparkLM(args.effort, batch_size=args.batch_size)
    results = simple_evaluate(
        model=lm, tasks=args.tasks.split(","), limit=args.limit,
    )
    if args.json:
        print(json.dumps(results["results"]))
    else:
        from lm_eval.utils import make_table
        print(make_table(results))


if __name__ == "__main__":
    main()