[eval] Batch rolling likelihood across documents

eval/robust_api.py

@@ -52,10 +52,12 @@
 logger = logging.getLogger("eval.robust_api")
 
 _PATCH_FLAG = "_marin_resilient_batch_patched"
+_ROLLING_PATCH_FLAG = "_marin_rolling_batch_patched"
 _COMPLETION_PATCH_FLAG = "_marin_completion_normalization_patched"
 _OPENAI_PAYLOAD_PATCH_FLAG = "_marin_openai_payload_patched"
 _REQUEST_FAILURE_PREFIX = "[EVALCHEMY_INFRASTRUCTURE_ERROR]"
 _MAX_REQUEST_FAILURE_DETAIL = 512
+_ROLLING_WINDOWS_PER_CONCURRENT_SLOT = 4
 _OPENAI_FIXED_GENERATION_MODEL = re.compile(r"^(?:gpt-5|o[134])(?:$|[-.])", re.IGNORECASE)
 
 
@@ -213,6 +215,71 @@ async def _guarded(message, cache_key, ctxlen, call_kwargs):
     return True
 
 
+def apply_rolling_loglikelihood_batching() -> bool:
+    """Batch rolling-likelihood windows across documents for API models."""
+    try:
+        from tqdm import tqdm
+
+        from lm_eval import utils
+        from lm_eval.models.api_models import TemplateAPI
+    except Exception as exc:  # noqa: BLE001 - never let the patch import break eval startup
+        logger.warning("rolling likelihood batching: could not import lm-eval symbols (%r); patch skipped.", exc)
+        return False
+
+    if getattr(TemplateAPI, _ROLLING_PATCH_FLAG, False):
+        return True
+
+    def loglikelihood_rolling(self, requests, disable_tqdm: bool = False):
+        loglikelihoods = []
+        pending_documents = []
+        pending_window_count = 0
+        target_window_count = max(1, self._concurrent * _ROLLING_WINDOWS_PER_CONCURRENT_SLOT)
+
+        def score_pending_documents():
+            nonlocal pending_documents, pending_window_count
+            if not pending_documents:
+                return
+
+            windows = [window for _, document_windows in pending_documents for window in document_windows]
+            window_scores = self._loglikelihood_tokens(windows, disable_tqdm=True)
+            offset = 0
+            for string, document_windows in pending_documents:
+                next_offset = offset + len(document_windows)
+                string_nll = sum(score for score, _ in window_scores[offset:next_offset])
+                loglikelihoods.append(string_nll)
+                self.cache_hook.add_partial("loglikelihood_rolling", (string,), string_nll)
+                offset = next_offset
+
+            pending_documents = []
+            pending_window_count = 0
+
+        for (string,) in tqdm([request.args for request in requests], disable=disable_tqdm):
+            document_windows = [
+                (None,) + window
+                for window in map(
+                    utils.make_disjoint_window,
+                    utils.get_rolling_token_windows(
+                        token_list=self.tok_encode(string),
+                        prefix_token=self.prefix_token_id,
+                        max_seq_len=self.max_length - 1,
+                        context_len=1,
+                    ),
+                )
+            ]
+            pending_documents.append((string, document_windows))
+            pending_window_count += len(document_windows)
+            if pending_window_count >= target_window_count:
+                score_pending_documents()
+
+        score_pending_documents()
+        return loglikelihoods
+
+    TemplateAPI.loglikelihood_rolling = loglikelihood_rolling
+    setattr(TemplateAPI, _ROLLING_PATCH_FLAG, True)
+    logger.info("rolling likelihood batching: patched TemplateAPI.loglikelihood_rolling.")
+    return True
+
+
 def apply_completion_normalization() -> bool:
     """Preserve reasoning content returned by lm-eval's chat-completions adapter."""
     try:
@@ -343,5 +410,6 @@ def _create_payload(
 
 # Apply on import so `from eval import robust_api` is enough to activate the patch.
 _APPLIED = apply()
+_ROLLING_LOGLIKELIHOOD_BATCHING_APPLIED = apply_rolling_loglikelihood_batching()
 _COMPLETION_NORMALIZATION_APPLIED = apply_completion_normalization()
 _OPENAI_PAYLOAD_CONTROLS_APPLIED = apply_openai_payload_controls()

tests/uncheatable_eval/test_rolling_loglikelihood.py

@@ -0,0 +1,74 @@
+import asyncio
+
+from eval import robust_api  # noqa: F401 - installs the rolling-likelihood patch
+from lm_eval.api.instance import Instance
+from lm_eval.models.api_models import TemplateAPI
+
+
+class _CacheHook:
+    def __init__(self):
+        self.values = {}
+
+    def add_partial(self, request_type, key, value):
+        self.values[(request_type, key)] = value
+
+
+class _ConcurrentRollingAPI(TemplateAPI):
+    def __init__(self):
+        self._batch_size = 1
+        self._concurrent = 3
+        self.max_length = 32
+        self.max_retries = 1
+        self.timeout = 30
+        self.verify_certificate = True
+        self.tokenizer = object()
+        self.cache_hook = _CacheHook()
+        self.active_requests = 0
+        self.max_active_requests = 0
+
+    @property
+    def prefix_token_id(self):
+        return 0
+
+    def tok_encode(self, string, **kwargs):
+        return list(range(1, len(string) + 1))
+
+    async def amodel_call(self, *, sem, messages, ctxlens, **kwargs):
+        await sem.acquire()
+        try:
+            self.active_requests += 1
+            self.max_active_requests = max(self.max_active_requests, self.active_requests)
+            await asyncio.sleep(0)  # Let the other scheduled endpoint calls start.
+            self.active_requests -= 1
+            return [
+                (-float(len(token) - ctxlen), False)
+                for token, ctxlen in zip(messages, ctxlens, strict=True)
+            ]
+        finally:
+            sem.release()
+
+    def parse_logprobs(self, *, outputs, tokens, ctxlens, **kwargs):
+        return [(-float(len(token) - ctxlen), False) for token, ctxlen in zip(tokens, ctxlens, strict=True)]
+
+    def _create_payload(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def parse_generations(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+def test_rolling_loglikelihood_batches_documents_concurrently():
+    model = _ConcurrentRollingAPI()
+    documents = ["one", "three", "seven", "nine", "eleven", "thirteen", "fifteen"]
+    requests = [
+        Instance(request_type="loglikelihood_rolling", doc={}, arguments=(document,), idx=index)
+        for index, document in enumerate(documents)
+    ]
+
+    scores = model.loglikelihood_rolling(requests, disable_tqdm=True)
+
+    assert scores == [-float(len(document)) for document in documents]
+    assert model.max_active_requests == 3
+    assert model.cache_hook.values == {
+        ("loglikelihood_rolling", (document,)): -float(len(document)) for document in documents
+    }