Add a replay provider that rescores saved generations

README.md

@@ -1203,6 +1203,26 @@ uv run olmo-eval run -m allenai/OLMo-2-1124-7B -t mmlu --force-download-model
 uv run olmo-eval run -m llama3.1-8b -t simpleqa:judge --harness dr_tulu
 ```
 
+### Rescoring Saved Generations
+
+The `replay` provider answers each request with the generation an earlier run saved
+for the same prompt, so scorer and metric changes can be measured without a model or
+GPU. Point it at a predictions file or a results directory holding one model's
+predictions (each file is read with the requests file saved beside it; a prompt saved by
+two task configurations needs the file), and repeat the original run's model and task
+options so the prompts match:
+
+```bash
+uv run olmo-eval run -m mock -t gsm8k -o limit=100 \
+    -H default -o provider.kind=replay -o provider.predictions=/tmp/results/ \
+    -O /tmp/rescored/
+```
+
+A prompt with no saved generation, or saved with a different number of samples, fails
+like a generation error and counts against the task's failure budget. Only text
+generations are replayed: logprob tasks, image or tool requests and multi-turn
+rollouts are not.
+
 ### Resuming Interrupted Runs
 
 With `--resume`, a run records each instance as it finishes. Restarted on the same

src/olmo_eval/common/constants/infrastructure.py

@@ -226,6 +226,7 @@ def cluster_has_weka(cluster: str) -> bool:
     "olmo_core_vlm": "olmo_core_vlm",
     "litellm": "litellm",
     "mock": None,
+    "replay": None,
 }
 """Mapping of backend types to their pyproject.toml optional dependency group names.
 

src/olmo_eval/common/types/literals.py

@@ -14,11 +14,20 @@ class ProviderKind(StrEnum):
     OLMO_CORE = "olmo_core"
     OLMO_CORE_VLM = "olmo_core_vlm"
     MOCK = "mock"
+    REPLAY = "replay"
     LITELLM = "litellm"
 
 
 ProviderLiteral = Literal[
-    "sglang", "vllm", "vllm_server", "hf", "olmo_core", "olmo_core_vlm", "mock", "litellm"
+    "sglang",
+    "vllm",
+    "vllm_server",
+    "hf",
+    "olmo_core",
+    "olmo_core_vlm",
+    "mock",
+    "replay",
+    "litellm",
 ]
 DtypeLiteral = Literal["auto", "float16", "bfloat16", "float32"]
 PriorityLiteral = Literal["low", "normal", "high", "urgent"]

src/olmo_eval/inference/__init__.py

@@ -104,6 +104,10 @@ def create_provider(
             from .providers.mock import MockProvider
 
             return MockProvider(model_name)
+        case "replay":
+            from .providers.replay import ReplayProvider
+
+            return ReplayProvider(model_name, predictions=kwargs.get("predictions", ""))
         case "hf":
             from .providers.huggingface import HuggingFaceProvider
 

src/olmo_eval/inference/providers/config.py

@@ -23,7 +23,7 @@ class ProviderConfig:
     a provider via the create_provider factory function.
 
     Attributes:
-        kind: Provider type (sglang, vllm, vllm_server, hf, olmo_core, litellm, mock).
+        kind: Provider type (sglang, vllm, vllm_server, hf, olmo_core, litellm, mock, replay).
         model: Model name or path (HuggingFace ID, API model name, or local path).
         alias: Short display name for the model (used in DB and S3 paths).
         base_url: Base URL for API-based providers (vllm_server, litellm).
@@ -147,6 +147,7 @@ def requires_local_gpu(self) -> bool:
             "max_model_len",
         ),
         "mock": (),
+        "replay": (),
     }
 
     def create_provider(self) -> InferenceProvider:

src/olmo_eval/inference/providers/replay.py

@@ -0,0 +1,168 @@
+"""Replay provider: serves the generations a previous run saved, so a task can be rescored."""
+
+from __future__ import annotations
+
+import json
+from collections import defaultdict, deque
+from pathlib import Path
+from typing import Any
+
+from olmo_eval.common.types import (
+    CONTEXT_CUT_TOKENS_KEY,
+    LMOutput,
+    LMRequest,
+    RequestType,
+    SamplingParams,
+)
+from olmo_eval.inference.base import InferenceProvider
+from olmo_eval.inference.errors import REQUEST_ERROR_KEY
+from olmo_eval.runners.common.types import PREDICTIONS_SUFFIX, REQUESTS_SUFFIX
+from olmo_eval.runners.io.writers import _lazy_loader_placeholder
+
+#: Saved per-output fields the replayed output's metadata carries over: the raw Harmony
+#: decode and its alignment error, so response parsing reproduces, and the context-cut
+#: marker, so overflow instances still score zero.
+_METADATA_FIELDS = (
+    "finish_reason",
+    "completion_tokens",
+    "prompt_tokens",
+    "text_with_special_tokens",
+    "response_parser_error",
+    CONTEXT_CUT_TOKENS_KEY,
+)
+
+
+def _prompt_key(context: Any, system_prompt: str | None) -> str:
+    """The key a saved request and a live request share when the model saw the same prompt."""
+    return json.dumps([context, system_prompt], sort_keys=True, default=_lazy_loader_placeholder)
+
+
+def _request_key(request: LMRequest) -> str:
+    # The context the requests writer saves for this request type.
+    if request.request_type == RequestType.CHAT:
+        context: Any = list(request.messages or ())
+    else:
+        context = request.prompt
+    return _prompt_key(context, request.system_prompt)
+
+
+def _requests_path(predictions_path: Path) -> Path:
+    """The requests file the writers place beside a predictions file."""
+    name = predictions_path.name.removesuffix(PREDICTIONS_SUFFIX) + REQUESTS_SUFFIX
+    return predictions_path.parent.parent.parent / "requests" / predictions_path.parent.name / name
+
+
+def _read_jsonl(path: Path) -> list[dict[str, Any]]:
+    with open(path) as f:
+        return [json.loads(line) for line in f if line.strip()]
+
+
+def _replayed_output(saved: dict[str, Any]) -> LMOutput:
+    # The raw generation, so the task strips reasoning and extracts answers again.
+    # An empty original is still the generation; only a missing one falls back to text.
+    original = saved.get("original_text")
+    text = original if original is not None else saved.get("text") or ""
+    metadata = {key: saved[key] for key in _METADATA_FIELDS if key in saved}
+    return LMOutput(text=text, metadata=metadata)
+
+
+def _generation_key(save: list[dict[str, Any]]) -> str:
+    """What replay serves for a saved prediction, leaving out the fields it scored."""
+    outputs = [_replayed_output(saved) for saved in save]
+    return json.dumps([[o.text, o.metadata] for o in outputs], sort_keys=True, default=str)
+
+
+class ReplayProvider(InferenceProvider):
+    """Answers each request with the generations a previous run saved for the same prompt.
+
+    ``predictions`` is a predictions JSONL file or a results directory searched for
+    them; each file is paired with the requests file written beside it, which holds the
+    prompts. A request whose prompt was not saved, or was saved with a different number
+    of samples, fails like a generation error, so the task's failure budget catches a
+    replay against the wrong run. Loglikelihood runs are skipped: their saved text is the
+    scored continuation, not a generation. A request does not say which instance it is, so a
+    prompt saved for several instances replays only when every save holds the same
+    generations; a prompt saved by two files (two configurations of a task) is rejected.
+    Requests with images or tools fail too: the saved files keep neither image
+    identities nor native tool calls.
+    """
+
+    def __init__(self, model_name: str = "replay", predictions: str = "") -> None:
+        super().__init__(model_name)
+        if not predictions:
+            raise ValueError("The replay provider needs provider.predictions=<file or directory>")
+        root = Path(predictions)
+        files = [root] if root.is_file() else sorted(root.rglob(f"*{PREDICTIONS_SUFFIX}"))
+        if not files:
+            raise FileNotFoundError(f"No *{PREDICTIONS_SUFFIX} files found under {root}")
+        models = sorted({path.parent.name for path in files})
+        if len(models) > 1:
+            raise ValueError(
+                f"{root} holds predictions for several models ({', '.join(models)}); "
+                "point provider.predictions at one model's predictions directory or file"
+            )
+
+        self._saved: dict[str, deque[list[dict[str, Any]]]] = defaultdict(deque)
+        sources: dict[str, Path] = {}
+        for path in files:
+            requests = {r["doc_id"]: r for r in _read_jsonl(_requests_path(path))}
+            for prediction in _read_jsonl(path):
+                saved_request = requests[prediction["doc_id"]]
+                # A loglikelihood run's output text is the continuation it scored, not a generation.
+                if saved_request.get("request_type") == "loglikelihood":
+                    continue
+                request = saved_request["request"]
+                key = _prompt_key(request.get("context"), request.get("system_prompt"))
+                if sources.setdefault(key, path) != path:
+                    raise ValueError(
+                        f"{sources[key].name} and {path.name} saved the same prompt; "
+                        "point provider.predictions at the predictions file to replay"
+                    )
+                self._saved[key].append(prediction["model_output"])
+        self._ambiguous = {
+            key
+            for key, saves in self._saved.items()
+            if len({_generation_key(save) for save in saves}) > 1
+        }
+
+    def _replay(self, request: LMRequest, num_samples: int) -> list[LMOutput]:
+        key = _request_key(request)
+        queue = self._saved.get(key)
+        if request.images or request.tools:
+            error = "requests with images or tools cannot be replayed"
+        elif not queue:
+            error = "no saved generation matches this prompt"
+        elif key in self._ambiguous:
+            error = (
+                f"{len(queue)} instances saved this prompt with different generations, "
+                "and a request does not say which instance it is"
+            )
+        elif len(queue[0]) != num_samples:
+            error = (
+                f"saved {len(queue[0])} samples for this prompt, the task asks for {num_samples}"
+            )
+        else:
+            return [_replayed_output(saved) for saved in queue.popleft()]
+        return [LMOutput(text="", metadata={REQUEST_ERROR_KEY: f"replay: {error}"})]
+
+    def generate(
+        self,
+        requests: list[LMRequest],
+        sampling_params: SamplingParams | None = None,
+    ) -> list[list[LMOutput]]:
+        num_samples = self._default_sampling_params(sampling_params).num_samples
+        return [self._replay(request, num_samples) for request in requests]
+
+    async def agenerate(
+        self,
+        requests: list[LMRequest],
+        sampling_params: SamplingParams | None = None,
+    ) -> list[list[LMOutput]]:
+        return self.generate(requests, sampling_params)
+
+    def logprobs(
+        self,
+        requests: list[LMRequest],
+        sampling_params: SamplingParams | None = None,
+    ) -> list[list[LMOutput]]:
+        raise NotImplementedError("The replay provider replays generations, not logprobs")