diff --git a/README.md b/README.md index ac7f35e3a..d126da420 100644 --- a/README.md +++ b/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 diff --git a/src/olmo_eval/common/constants/infrastructure.py b/src/olmo_eval/common/constants/infrastructure.py index f1c8ded09..118c349d9 100644 --- a/src/olmo_eval/common/constants/infrastructure.py +++ b/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. diff --git a/src/olmo_eval/common/types/literals.py b/src/olmo_eval/common/types/literals.py index 327c1de9c..b2d773f10 100644 --- a/src/olmo_eval/common/types/literals.py +++ b/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"] diff --git a/src/olmo_eval/inference/__init__.py b/src/olmo_eval/inference/__init__.py index a84841a0d..8e3c073ec 100644 --- a/src/olmo_eval/inference/__init__.py +++ b/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 diff --git a/src/olmo_eval/inference/providers/config.py b/src/olmo_eval/inference/providers/config.py index 5fca1553b..8155d8123 100644 --- a/src/olmo_eval/inference/providers/config.py +++ b/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: diff --git a/src/olmo_eval/inference/providers/replay.py b/src/olmo_eval/inference/providers/replay.py new file mode 100644 index 000000000..3c8f812c0 --- /dev/null +++ b/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=") + 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") diff --git a/tests/providers/test_replay_provider.py b/tests/providers/test_replay_provider.py new file mode 100644 index 000000000..e5fb6c72b --- /dev/null +++ b/tests/providers/test_replay_provider.py @@ -0,0 +1,264 @@ +"""Tests for the replay provider.""" + +from pathlib import Path + +import pytest + +from olmo_eval.common.types import ( + CONTEXT_CUT_TOKENS_KEY, + Instance, + LMOutput, + LMRequest, + RequestType, + Response, + SamplingParams, +) +from olmo_eval.evals.tasks.common import get_task +from olmo_eval.inference.errors import request_error +from olmo_eval.inference.providers import ProviderConfig +from olmo_eval.inference.providers.replay import ReplayProvider +from olmo_eval.runners.io.builders import build_predictions +from olmo_eval.runners.io.writers import write_predictions_jsonl, write_requests_jsonl + +CHAT = ({"role": "user", "content": "What is 6 x 7?"},) + + +def _save( + output_dir: Path, + spec: str, + saved: list[tuple[object, list[dict]]], + model: str = "model", + request_type: str = "generate_until", +) -> None: + """Write a run's files for (context, model_output) pairs, as the runner does.""" + requests = [ + {"doc_id": i, "request_type": request_type, "request": {"context": c}} + for i, (c, _) in enumerate(saved) + ] + predictions = [{"doc_id": i, "model_output": out} for i, (_, out) in enumerate(saved)] + write_requests_jsonl(str(output_dir), spec, requests, model) + write_predictions_jsonl(str(output_dir), spec, predictions, model) + + +def _completion(prompt: str) -> LMRequest: + return LMRequest(request_type=RequestType.COMPLETION, prompt=prompt) + + +def _chat(messages: tuple[dict, ...]) -> LMRequest: + return LMRequest(request_type=RequestType.CHAT, messages=messages) + + +def test_replays_saved_generations_by_prompt(tmp_path: Path) -> None: + _save( + tmp_path, + "task", + [ + ("Q: one", [{"text": "1", "finish_reason": "stop", "extracted_answer": "x"}]), + (list(CHAT), [{"text": "42", "original_text": "hm42"}]), + ], + ) + provider = ReplayProvider(predictions=str(tmp_path)) + + chat, completion = provider.generate([_chat(CHAT), _completion("Q: one")]) + + # The raw generation comes back, so the task strips and extracts again. + assert [o.text for o in chat] == ["hm42"] + assert completion[0].text == "1" + assert completion[0].metadata == {"finish_reason": "stop"} + assert completion[0].extracted_answer is None + + +def test_reads_a_single_predictions_file(tmp_path: Path) -> None: + _save(tmp_path, "a", [("Q: a", [{"text": "A"}])]) + _save(tmp_path, "b", [("Q: b", [{"text": "B"}])]) + path = next(tmp_path.rglob("a*-predictions.jsonl")) + + provider = ReplayProvider(predictions=str(path)) + + assert request_error(provider.generate([_completion("Q: b")])[0]) is not None + assert provider.generate([_completion("Q: a")])[0][0].text == "A" + + +def test_unsaved_prompt_and_sample_count_mismatch_fail_the_request(tmp_path: Path) -> None: + _save(tmp_path, "task", [("Q: one", [{"text": "1"}, {"text": "2"}])]) + provider = ReplayProvider(predictions=str(tmp_path)) + + unsaved, wrong_count = provider.generate( + [_completion("Q: other"), _completion("Q: one")], SamplingParams(num_samples=1) + ) + + assert "no saved generation" in str(request_error(unsaved)) + assert "saved 2 samples" in str(request_error(wrong_count)) + two = provider.generate([_completion("Q: one")], SamplingParams(num_samples=2)) + assert [o.text for o in two[0]] == ["1", "2"] + + +def test_a_repeated_prompt_with_one_generation_is_served_once_per_save(tmp_path: Path) -> None: + _save(tmp_path, "task", [("Q", [{"text": "same"}]), ("Q", [{"text": "same"}])]) + provider = ReplayProvider(predictions=str(tmp_path)) + + outputs = provider.generate([_completion("Q"), _completion("Q"), _completion("Q")]) + + assert [o[0].text for o in outputs[:2]] == ["same", "same"] + assert request_error(outputs[2]) is not None + + +def test_a_repeated_prompt_with_different_generations_fails(tmp_path: Path) -> None: + # The runner shuffles requests, so arrival order would hand instances each other's output. + _save(tmp_path, "task", [("Q", [{"text": "first"}]), ("Q", [{"text": "second"}])]) + provider = ReplayProvider(predictions=str(tmp_path)) + + outputs = provider.generate([_completion("Q"), _completion("Q")]) + + assert all("different generations" in str(request_error(o)) for o in outputs) + + +def test_a_repeated_prompt_scored_differently_still_replays(tmp_path: Path) -> None: + # Instances with different gold answers score the same generation differently. + _save( + tmp_path, + "task", + [ + ("Q", [{"text": "same", "sample_metrics": {"acc": 1.0}}]), + ("Q", [{"text": "same", "sample_metrics": {"acc": 0.0}}]), + ], + ) + provider = ReplayProvider(predictions=str(tmp_path)) + + outputs = provider.generate([_completion("Q"), _completion("Q")]) + + assert [o[0].text for o in outputs] == ["same", "same"] + + +def test_a_loglikelihood_run_is_not_replayed_as_generations(tmp_path: Path) -> None: + # Its saved text is the gold continuation it scored. + _save(tmp_path, "bpb", [("Q", [{"text": "gold"}])], request_type="loglikelihood") + _save(tmp_path, "gen", [("P", [{"text": "generated"}])]) + provider = ReplayProvider(predictions=str(tmp_path)) + + gold, generated = provider.generate([_completion("Q"), _completion("P")]) + + assert "no saved generation" in str(request_error(gold)) + assert generated[0].text == "generated" + + +def test_provider_config_passes_the_predictions_path(tmp_path: Path) -> None: + _save(tmp_path, "task", [("Q", [{"text": "A"}])]) + config = ProviderConfig.from_dict( + {"kind": "replay", "model": "replay", "predictions": str(tmp_path)} + ) + + provider = config.create_provider() + + assert isinstance(provider, ReplayProvider) + assert not config.requires_gpu + + +def test_missing_predictions_are_reported(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="provider.predictions"): + ReplayProvider() + with pytest.raises(FileNotFoundError): + ReplayProvider(predictions=str(tmp_path)) + + +def test_rejects_a_directory_with_several_models(tmp_path: Path) -> None: + _save(tmp_path, "task", [("Q", [{"text": "A"}])], model="a") + _save(tmp_path, "task", [("Q", [{"text": "B"}])], model="b") + + with pytest.raises(ValueError, match="several models"): + ReplayProvider(predictions=str(tmp_path)) + provider = ReplayProvider(predictions=str(tmp_path / "predictions" / "b")) + assert provider.generate([_completion("Q")])[0][0].text == "B" + + +def test_image_and_tool_requests_fail(tmp_path: Path) -> None: + _save(tmp_path, "task", [(list(CHAT), [{"text": "A"}])]) + provider = ReplayProvider(predictions=str(tmp_path)) + image = LMRequest(request_type=RequestType.CHAT, messages=CHAT, images=("img",)) + tool = LMRequest(request_type=RequestType.CHAT, messages=CHAT, tools=(object(),)) + + for outputs in provider.generate([image, tool]): + assert "images or tools" in str(request_error(outputs)) + + +def test_rejects_a_prompt_saved_by_two_task_configurations(tmp_path: Path) -> None: + _save(tmp_path, "task", [("Q", [{"text": "A"}])]) + _save(tmp_path, "task:variant", [("Q", [{"text": "B"}, {"text": "C"}])]) + + with pytest.raises(ValueError, match="saved the same prompt"): + ReplayProvider(predictions=str(tmp_path)) + path = next(tmp_path.rglob("task_variant*-predictions.jsonl")) + outputs = ReplayProvider(predictions=str(path)).generate( + [_completion("Q")], SamplingParams(num_samples=2) + ) + assert [o.text for o in outputs[0]] == ["B", "C"] + + +def _harmony_parse(output: LMOutput) -> Response: + task = get_task("ifeval", {"response_parser": "harmony"}) + instance = Instance(question="Q") + response = Response(instance=instance, request=task.format_request(instance), outputs=[output]) + task.strip_thinking_traces([response]) + return response + + +def test_harmony_parsed_output_replays_to_the_same_answer(tmp_path: Path) -> None: + # vLLM's default decode drops the Harmony delimiters; the raw decode keeps them. + raw = ( + "<|channel|>analysis<|message|>wrong<|end|>" + "<|start|>assistant<|channel|>final<|message|>good<|return|>" + ) + live = LMOutput( + text="analysiswrongassistantfinalgood", metadata={"text_with_special_tokens": raw} + ) + saved = build_predictions([_harmony_parse(live)]) + assert live.text == "good" + _save(tmp_path, "task", [("Q", saved[0]["model_output"])]) + + replayed = ReplayProvider(predictions=str(tmp_path)).generate([_completion("Q")])[0][0] + _harmony_parse(replayed) + + assert replayed.text == "good" + + +def test_response_parser_error_carries_over(tmp_path: Path) -> None: + error = "Cannot align Harmony token decode with returned completion text" + _save(tmp_path, "task", [("Q", [{"text": "A", "response_parser_error": error}])]) + + replayed = ReplayProvider(predictions=str(tmp_path)).generate([_completion("Q")])[0][0] + + assert replayed.metadata["response_parser_error"] == error + with pytest.raises(ValueError, match="Cannot align"): + _harmony_parse(replayed) + + +def test_context_cut_tokens_survive_replay(tmp_path: Path) -> None: + _save(tmp_path, "task", [("Q", [{"text": "A", CONTEXT_CUT_TOKENS_KEY: 512}])]) + + replayed = ReplayProvider(predictions=str(tmp_path)).generate([_completion("Q")])[0][0] + + assert replayed.metadata[CONTEXT_CUT_TOKENS_KEY] == 512 + + +def test_an_empty_original_generation_replays_empty(tmp_path: Path) -> None: + # HELMET's summarization :base variants save an empty original and score "Summary:". + _save(tmp_path, "task", [("Q", [{"text": "Summary:", "original_text": ""}])]) + + replayed = ReplayProvider(predictions=str(tmp_path)).generate([_completion("Q")])[0][0] + + assert replayed.text == "" + + +def test_a_completion_request_is_keyed_by_its_prompt_even_with_messages(tmp_path: Path) -> None: + # A harness-level system prompt fills request.messages on a completion request, but the + # requests file saves its prompt. + _save(tmp_path, "task", [("Q", [{"text": "A"}])]) + request = LMRequest( + request_type=RequestType.COMPLETION, + prompt="Q", + messages=({"role": "system", "content": "Be brief."},), + ) + + replayed = ReplayProvider(predictions=str(tmp_path)).generate([request])[0][0] + + assert replayed.text == "A"