| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | """MRCR long-context evaluation for vLLM's OpenAI-compatible server. |
| 5 | |
| 6 | Streams samples from `openai/mrcr` on HuggingFace, sends chat completions to |
| 7 | the server, and scores each response with a prefix-gated SequenceMatcher ratio |
| 8 | against the reference answer. |
| 9 | """ |
| 10 | |
| 11 | import argparse |
| 12 | import asyncio |
| 13 | import json |
| 14 | import time |
| 15 | from difflib import SequenceMatcher |
| 16 | |
| 17 | import aiohttp |
| 18 | import numpy as np |
| 19 | import requests |
| 20 | from tqdm.asyncio import tqdm |
| 21 | |
| 22 | DATASET_REPO = "openai/mrcr" |
| 23 | NEEDLE_SHARDS = { |
| 24 | 2: "2needle/2needle_0.parquet", |
| 25 | 4: "4needle/4needle_0.parquet", |
| 26 | 8: "8needle/8needle_0.parquet", |
| 27 | } |
| 28 | # Reserve headroom for chat-template tokens on top of the messages. |
| 29 | PROMPT_SAFETY_BUFFER = 256 |
| 30 | # Pre-filter heuristic before the authoritative /tokenize check. |
| 31 | CHARS_PER_TOKEN = 4 |
| 32 | # Skip chain-of-thought on reasoning models; ignored by non-reasoning templates. |
| 33 | DEFAULT_EXTRA_BODY: dict = {"chat_template_kwargs": {"enable_thinking": False}} |
| 34 | |
| 35 | |
| 36 | def discover_server_model(base_url: str) -> tuple[str, int | None]: |
| 37 | """Return (model_id, max_model_len) from /v1/models.""" |
| 38 | resp = requests.get(f"{base_url}/v1/models", timeout=30) |
| 39 | resp.raise_for_status() |
| 40 | data = resp.json().get("data", []) |
| 41 | if not data: |
| 42 | raise RuntimeError(f"No models advertised at {base_url}/v1/models") |
| 43 | entry = data[0] |
| 44 | return entry["id"], entry.get("max_model_len") |
| 45 | |
| 46 | |
| 47 | def count_chat_tokens(base_url: str, model: str, messages: list[dict]) -> int: |
| 48 | """Return the chat-template-rendered token count via /tokenize.""" |
| 49 | resp = requests.post( |
| 50 | f"{base_url}/tokenize", |
| 51 | json={"model": model, "messages": messages, "add_generation_prompt": True}, |
| 52 | timeout=120, |
| 53 | ) |
| 54 | resp.raise_for_status() |
| 55 | return int(resp.json()["count"]) |
| 56 | |
| 57 | |
| 58 | def _load_mrcr_samples( |
| 59 | needles: list[int], |
| 60 | max_prompt_tokens: int, |
| 61 | num_samples: int, |
| 62 | seed: int, |
| 63 | base_url: str, |
| 64 | model_name: str, |
| 65 | ) -> list[dict]: |
| 66 | """Stream MRCR samples balanced across needle buckets, token-verified.""" |
| 67 | try: |
| 68 | from datasets import load_dataset |
| 69 | except ImportError as e: |
| 70 | raise ImportError( |
| 71 | "MRCR eval requires `datasets`. Install with: uv pip install datasets" |
| 72 | ) from e |
| 73 | |
| 74 | max_chars = max_prompt_tokens * CHARS_PER_TOKEN |
| 75 | per_bucket = num_samples // len(needles) |
| 76 | leftover = num_samples - per_bucket * len(needles) |
| 77 | |
| 78 | samples: list[dict] = [] |
| 79 | for idx, n in enumerate(needles): |
| 80 | if n not in NEEDLE_SHARDS: |
| 81 | raise ValueError(f"Unsupported needle count {n}") |
| 82 | target = per_bucket + (1 if idx < leftover else 0) |
| 83 | if target == 0: |
| 84 | continue |
| 85 | |
| 86 | ds = load_dataset( |
| 87 | DATASET_REPO, |
| 88 | data_files=NEEDLE_SHARDS[n], |
| 89 | split="train", |
| 90 | streaming=True, |
| 91 | ).shuffle(seed=seed + n, buffer_size=16) |
| 92 | |
| 93 | taken = 0 |
| 94 | for row in ds: |
| 95 | if int(row.get("n_chars", 0)) > max_chars: |
| 96 | continue |
| 97 | prompt = row["prompt"] |
| 98 | messages = json.loads(prompt) if isinstance(prompt, str) else list(prompt) |
| 99 | n_tokens = count_chat_tokens(base_url, model_name, messages) |
| 100 | if n_tokens > max_prompt_tokens: |
| 101 | continue |
| 102 | samples.append( |
| 103 | { |
| 104 | "messages": messages, |
| 105 | "answer": row["answer"], |
| 106 | "random_string_to_prepend": row["random_string_to_prepend"], |
| 107 | "n_needles": int(row["n_needles"]), |
| 108 | "n_tokens": n_tokens, |
| 109 | } |
| 110 | ) |
| 111 | taken += 1 |
| 112 | if taken >= target: |
| 113 | break |
| 114 | |
| 115 | if taken < target: |
| 116 | print(f"Warning: only {taken}/{target} samples for n_needles={n}") |
| 117 | |
| 118 | if not samples: |
| 119 | raise RuntimeError("No MRCR samples fit; loosen max_prompt_tokens.") |
| 120 | return samples |
| 121 | |
| 122 | |
| 123 | def score_mrcr(response: str, answer: str, random_prefix: str) -> float: |
| 124 | """Prefix-gated SequenceMatcher ratio; 0 if the prefix is missing.""" |
| 125 | if not response.startswith(random_prefix): |
| 126 | return 0.0 |
| 127 | stripped = response[len(random_prefix) :] |
| 128 | return SequenceMatcher(a=answer, b=stripped, autojunk=False).ratio() |
| 129 | |
| 130 | |
| 131 | async def _call_chat( |
| 132 | session: aiohttp.ClientSession, |
| 133 | url: str, |
| 134 | model: str, |
| 135 | messages: list[dict], |
| 136 | max_tokens: int, |
| 137 | temperature: float, |
| 138 | seed: int | None, |
| 139 | extra_body: dict, |
| 140 | ) -> tuple[str, int]: |
| 141 | data = { |
| 142 | "model": model, |
| 143 | "messages": messages, |
| 144 | "temperature": temperature, |
| 145 | "max_tokens": max_tokens, |
| 146 | **extra_body, |
| 147 | } |
| 148 | if seed is not None: |
| 149 | data["seed"] = seed |
| 150 | try: |
| 151 | async with session.post(f"{url}/v1/chat/completions", json=data) as resp: |
| 152 | resp.raise_for_status() |
| 153 | result = await resp.json() |
| 154 | text = result["choices"][0]["message"]["content"] or "" |
| 155 | return text, result.get("usage", {}).get("completion_tokens", 0) |
| 156 | except Exception as e: |
| 157 | print(f"chat request failed: {e}") |
| 158 | return "", 0 |
| 159 | |
| 160 | |
| 161 | def evaluate_mrcr( |
| 162 | model_name: str | None = None, |
| 163 | num_samples: int = 40, |
| 164 | needles: list[int] | None = None, |
| 165 | max_prompt_tokens: int | None = None, |
| 166 | max_tokens: int = 2048, |
| 167 | host: str = "http://127.0.0.1", |
| 168 | port: int = 8000, |
| 169 | temperature: float = 0.0, |
| 170 | seed: int | None = 42, |
| 171 | concurrency: int = 8, |
| 172 | extra_body: dict | None = None, |
| 173 | ) -> dict: |
| 174 | """Run MRCR against a vLLM server; auto-discovers model and context.""" |
| 175 | needles = needles or [2, 4, 8] |
| 176 | extra_body = DEFAULT_EXTRA_BODY if extra_body is None else extra_body |
| 177 | base_url = f"{host}:{port}" |
| 178 | |
| 179 | discovered_model, server_max_len = discover_server_model(base_url) |
| 180 | if model_name is None: |
| 181 | model_name = discovered_model |
| 182 | if max_prompt_tokens is None: |
| 183 | if server_max_len is None: |
| 184 | raise RuntimeError( |
| 185 | "Server did not advertise max_model_len; pass --max-prompt-tokens." |
| 186 | ) |
| 187 | max_prompt_tokens = max(512, server_max_len - max_tokens - PROMPT_SAFETY_BUFFER) |
| 188 | print( |
| 189 | f"Model: {model_name} | max_prompt_tokens={max_prompt_tokens} " |
| 190 | f"(server max_model_len={server_max_len}, max_tokens={max_tokens})" |
| 191 | ) |
| 192 | |
| 193 | samples = _load_mrcr_samples( |
| 194 | needles=needles, |
| 195 | max_prompt_tokens=max_prompt_tokens, |
| 196 | num_samples=num_samples, |
| 197 | seed=seed or 0, |
| 198 | base_url=base_url, |
| 199 | model_name=model_name, |
| 200 | ) |
| 201 | tok_counts = [s["n_tokens"] for s in samples] |
| 202 | print( |
| 203 | f"Loaded {len(samples)} samples (needles={needles}, " |
| 204 | f"tokens={min(tok_counts)}-{max(tok_counts)})" |
| 205 | ) |
| 206 | |
| 207 | async def run(): |
| 208 | sem = asyncio.Semaphore(concurrency) |
| 209 | responses = [""] * len(samples) |
| 210 | out_tokens = [0] * len(samples) |
| 211 | |
| 212 | async def one(session, i): |
| 213 | async with sem: |
| 214 | text, toks = await _call_chat( |
| 215 | session=session, |
| 216 | url=base_url, |
| 217 | model=model_name, |
| 218 | messages=samples[i]["messages"], |
| 219 | max_tokens=max_tokens, |
| 220 | temperature=temperature, |
| 221 | seed=seed, |
| 222 | extra_body=extra_body, |
| 223 | ) |
| 224 | responses[i] = text |
| 225 | out_tokens[i] = toks |
| 226 | |
| 227 | timeout = aiohttp.ClientTimeout(total=1800) |
| 228 | async with aiohttp.ClientSession(timeout=timeout) as session: |
| 229 | await tqdm.gather( |
| 230 | *[one(session, i) for i in range(len(samples))], desc="MRCR" |
| 231 | ) |
| 232 | return responses, out_tokens |
| 233 | |
| 234 | tic = time.perf_counter() |
| 235 | responses, out_tokens = asyncio.run(run()) |
| 236 | latency = time.perf_counter() - tic |
| 237 | |
| 238 | scores = np.array( |
| 239 | [ |
| 240 | score_mrcr(r, s["answer"], s["random_string_to_prepend"]) |
| 241 | for r, s in zip(responses, samples) |
| 242 | ] |
| 243 | ) |
| 244 | prefix_hits = np.array( |
| 245 | [ |
| 246 | r.startswith(s["random_string_to_prepend"]) |
| 247 | for r, s in zip(responses, samples) |
| 248 | ] |
| 249 | ) |
| 250 | per_needle = { |
| 251 | f"match_ratio_n{n}": float( |
| 252 | scores[np.array([s["n_needles"] == n for s in samples])].mean() |
| 253 | ) |
| 254 | for n in needles |
| 255 | if any(s["n_needles"] == n for s in samples) |
| 256 | } |
| 257 | |
| 258 | total_out = int(sum(out_tokens)) |
| 259 | return { |
| 260 | "model": model_name, |
| 261 | "match_ratio": float(scores.mean()), |
| 262 | "prefix_hit_rate": float(prefix_hits.mean()), |
| 263 | "per_needle": per_needle, |
| 264 | "num_samples": len(samples), |
| 265 | "latency": latency, |
| 266 | "total_output_tokens": total_out, |
| 267 | "tokens_per_second": total_out / latency if latency > 0 else 0.0, |
| 268 | "max_tokens": max_tokens, |
| 269 | "needles": needles, |
| 270 | "max_prompt_tokens": max_prompt_tokens, |
| 271 | } |
| 272 | |
| 273 | |
| 274 | def main() -> None: |
| 275 | p = argparse.ArgumentParser(description="MRCR evaluation for vLLM serve") |
| 276 | p.add_argument("--model", default=None, help="Default: discovered from /v1/models") |
| 277 | p.add_argument("--num-samples", type=int, default=40) |
| 278 | p.add_argument( |
| 279 | "--needles", type=int, nargs="+", default=[2, 4, 8], choices=[2, 4, 8] |
| 280 | ) |
| 281 | p.add_argument( |
| 282 | "--max-prompt-tokens", |
| 283 | type=int, |
| 284 | default=None, |
| 285 | help="Default: server max_model_len - max_tokens - buffer", |
| 286 | ) |
| 287 | p.add_argument("--max-tokens", type=int, default=2048) |
| 288 | p.add_argument("--host", default="http://127.0.0.1") |
| 289 | p.add_argument("--port", type=int, default=8000) |
| 290 | p.add_argument("--temperature", type=float, default=0.0) |
| 291 | p.add_argument("--seed", type=int, default=42) |
| 292 | p.add_argument("--concurrency", type=int, default=8) |
| 293 | p.add_argument( |
| 294 | "--extra-body", |
| 295 | default=None, |
| 296 | help="JSON merged into each request. " |
| 297 | "Pass '{}' to disable the default enable_thinking=false.", |
| 298 | ) |
| 299 | p.add_argument("--save-results", default=None) |
| 300 | args = p.parse_args() |
| 301 | |
| 302 | extra_body = json.loads(args.extra_body) if args.extra_body else None |
| 303 | |
| 304 | result = evaluate_mrcr( |
| 305 | model_name=args.model, |
| 306 | num_samples=args.num_samples, |
| 307 | needles=args.needles, |
| 308 | max_prompt_tokens=args.max_prompt_tokens, |
| 309 | max_tokens=args.max_tokens, |
| 310 | host=args.host, |
| 311 | port=args.port, |
| 312 | temperature=args.temperature, |
| 313 | seed=args.seed, |
| 314 | concurrency=args.concurrency, |
| 315 | extra_body=extra_body, |
| 316 | ) |
| 317 | |
| 318 | print("\nResults:") |
| 319 | print(f" match_ratio: {result['match_ratio']:.4f}") |
| 320 | print(f" prefix_hit_rate: {result['prefix_hit_rate']:.4f}") |
| 321 | for k, v in result["per_needle"].items(): |
| 322 | print(f" {k}: {v:.4f}") |
| 323 | print(f" samples: {result['num_samples']}") |
| 324 | print(f" latency: {result['latency']:.1f}s") |
| 325 | print(f" output tok/s: {result['tokens_per_second']:.1f}") |
| 326 | |
| 327 | if args.save_results: |
| 328 | with open(args.save_results, "w") as f: |
| 329 | json.dump(result, f, indent=2) |
| 330 | |
| 331 | |
| 332 | if __name__ == "__main__": |
| 333 | main() |