"""run_lmeval.py — reproduce min-spark's published benchmarks (ADR-0015/0025). Port of engine/eval/lm_eval_model.py's LoglikelihoodLM (which subclasses TemplateLM) over MinSparkForCausalLM. See that file's docstring for the exact 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 # --- required TemplateLM/LM surface --- @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) # --- scoring execution context (adapter-owned) --- 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()