diff --git a/eval/robust_api.py b/eval/robust_api.py index 9b71a52d..fc7ef706 100644 --- a/eval/robust_api.py +++ b/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() diff --git a/tests/uncheatable_eval/test_rolling_loglikelihood.py b/tests/uncheatable_eval/test_rolling_loglikelihood.py new file mode 100644 index 00000000..0806a748 --- /dev/null +++ b/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 + }