513 lines
21 KiB
Python

"""Async generation runner: model + dataset -> predictions -> scored report.
The async boundary is exactly "waiting on the model". Data loading and
scoring stay synchronous (fast, CPU/disk bound); this coroutine fans out
model calls with a semaphore, streams progress, then hands the collected
raw strings to the sync evaluate().
from evalharness.model import run_eval
report = asyncio.run(run_eval(ds, 'mock', limit=50)) # offline smoke
report = asyncio.run(run_eval(ds, 'openai/http://gpu03:8000/v1?qwen3-8b'))
"""
import asyncio
import os
import time
from typing import Any, Dict, List, Optional, Union
from ..data.dataset import Dataset
from ..data.sample import ChatMessage, Sample
from ..eval.recipe import EvalRecipe
from ..eval.record import EvalReport
from ..eval.runner import evaluate
from .adapter import ModelAdapter, resolve_adapter
from .output import Usage
async def generate_predictions(
adapter: ModelAdapter,
samples: List[Sample],
concurrency: int = 32,
limit: Optional[int] = None,
gen_kwargs: Optional[Dict[str, Any]] = None,
progress: bool = True,
env_factory=None,
system: str = '',
max_turns: int = 8,
max_input_chars: int = 0,
max_input_tokens: int = 0,
tokenizer_path: str = '',
attach_context_keys: tuple = ('passage', 'context'),
dataset_spec=None,
limit_per_task: Optional[int] = None,
checkpoint: Union[bool, str] = False,
dataset_name: str = 'adhoc',
few_shot_num: int = 0,
few_shot_samples: Optional[List[Sample]] = None,
few_shot_text: Optional[str] = None,
prompt_style: str = 'strict_letter',
) -> tuple:
"""Fan out model calls; returns (pred-dicts, total_usage).
MCQ samples are generated with the strict-letter contract ('ANSWER: X',
evalscope parity). few_shot: official exemplar text (few_shot_text) or
dev/train-split samples (few_shot_samples) are prepended.
"""
gen_kwargs = gen_kwargs or {}
def _default_max_tokens() -> int:
return 4096
sem = asyncio.Semaphore(concurrency)
total_usage = Usage()
done_count = 0
t0 = time.time()
def assemble(sample: Sample) -> str:
parts = []
if few_shot_text:
parts.append(few_shot_text.strip()) # official exemplars, verbatim
elif few_shot_num and few_shot_samples:
letters_fs = 'ABCDEFGHIJ'
for fs in few_shot_samples[:few_shot_num]:
line = f'Question: {fs.input_text}'
if fs.choices:
line += '\n' + '\n'.join(f'{letters_fs[j]}. {c}' for j, c in enumerate(fs.choices))
ans = fs.target if not isinstance(fs.target, list) else fs.target[0]
line += f'\nAnswer: {ans}'
parts.append(line)
for key in attach_context_keys:
ctx = (sample.metadata or {}).get(key)
if ctx:
parts.append(str(ctx))
question = sample.input_text
if sample.choices:
if prompt_style in ('strict_letter', 'auto'):
# evalscope/OpenAI-style contract: reply ONLY 'ANSWER: X'
letters = 'ABCDEFGHIJ'
opts = '\n'.join(f'{letters[i]}. {c}' for i, c in enumerate(sample.choices)
if i < len(letters))
question = (f'Answer the following multiple choice question. The entire '
f'content of your response should be of the following format: '
f"'ANSWER: [LETTER]' (without quotes) where [LETTER] is one of "
f'{letters[:len(sample.choices)]}.\n\n{question}\n\n{opts}')
else:
letters = 'ABCDEFGHIJ'
opts = '\n'.join(f'{letters[i]}. {c}' for i, c in enumerate(sample.choices)
if i < len(letters))
question = (f'{question}\n\n{opts}\n\n'
'Answer with the letter of the correct option.')
elif sample.task_type in ('qa',):
# hle OFFICIAL protocol: answer_type-specific system contract
at = (sample.metadata or {}).get('answer_type')
if at == 'exactMatch':
question = (
'Your response should be in the following format:\n'
'Explanation: {your explanation for your final answer}\n'
'Exact Answer: {your succinct, final answer}\n'
'Confidence: {your confidence score between 0% and 100% for your answer}\n\n'
f'{question}')
elif at == 'multipleChoice':
question = (
'Your response should be in the following format:\n'
'Explanation: {your explanation for your answer choice}\n'
'Answer: {your chosen answer}\n'
'Confidence: {your confidence score between 0% and 100% for your answer}\n\n'
f'{question}')
else:
question = (f'{question}\n\n'
'End your reply with the final answer on its own last line '
'in the form "Answer: <answer>".')
ds_spec = dataset_spec
if ds_spec is not None and getattr(ds_spec, 'prompt_suffix', ''):
question = question + ds_spec.prompt_suffix
parts.append(question)
text = '\n\n'.join(parts)
if max_input_tokens:
# reserve room for the OUTPUT budget + safety margin, else the
# server rejects input+max_tokens > context_limit by 1 token
budget = max(1024, max_input_tokens
- int(gen_kwargs.get('max_tokens') or 4096) - 2048)
try:
from .truncation import truncate_middle_tokens, default_tokenizer_path
text = truncate_middle_tokens(text, budget,
tokenizer_path or default_tokenizer_path())
except Exception as e:
# no tokenizer/transformers: degrade to a CHARS budget that
# approximates the token cap (never send the raw 2M-token input)
approx_chars = budget * 3
if len(text) > approx_chars:
keep = approx_chars // 2
text = f'{text[:keep]}\n\n...[truncated {len(text) - 2 * keep} chars]...\n\n{text[-keep:]}'
print(f'truncation degraded to chars ({type(e).__name__})', flush=True)
if max_input_chars and len(text) > max_input_chars:
keep = max_input_chars // 2
head = text[:keep]
tail = text[-keep:]
text = f'{head}\n\n...[truncated {len(text) - 2 * keep} chars]...\n\n{tail}'
return text
async def one(sample: Sample) -> Dict[str, Any]:
nonlocal done_count, total_usage
if env_factory is not None:
from ..agent import drive, trajectory_to_prediction
from ..agent.loop import Environment, Usage as _U # noqa: F401
async with sem:
env = env_factory()
if type(env).run_task is not Environment.run_task:
# self-running env (official engine bundles: tau2/swe)
pred = await env.run_task(adapter, sample,
max_turns=max_turns, system=system)
if pred is None:
traj = await drive(adapter, sample, env=env,
max_turns=max_turns, system=system)
pred = trajectory_to_prediction(traj)
else:
traj = await drive(adapter, sample, env=env,
max_turns=max_turns, system=system)
pred = trajectory_to_prediction(traj)
if not pred.get('usage'):
pred['usage'] = traj.usage.model_dump() if 'traj' in dir() else {}
pred.setdefault('group_key', str(sample.metadata.get('test_category')
or sample.metadata.get('category')
or sample.metadata.get('domain')
or sample.metadata.get('id') or sample.id or ''))
u = pred.get('usage') or {}
total_usage = total_usage + Usage(
input_tokens=int(u.get('input_tokens', 0) or 0),
output_tokens=int(u.get('output_tokens', 0) or 0),
total_tokens=int(u.get('total_tokens', 0) or 0),
latency_s=float(u.get('latency_s', 0) or 0))
done_count += 1
_progress(progress, done_count, len(samples), t0, total_usage)
return pred
messages = ([ChatMessage(role='user', content=assemble(sample))]
if isinstance(sample.input, str) else list(sample.input))
tools = None
if sample.tools:
tools = [{'name': t.name, 'description': t.description or '',
'parameters': t.parameters} for t in sample.tools]
if getattr(adapter, 'name', '') == 'mock' \
and adapter.extra.get('mode') in ('boxed', 'oracle', 'fc') \
and sample.target not in ('', None):
# oracle channel for mock verification so full pipelines run offline
messages = messages + [ChatMessage(role='user',
content=f'MOCKTARGET::{sample.target}')]
async with sem:
out = await adapter.generate(messages, tools=tools, **gen_kwargs)
total_usage = total_usage + out.usage
text = out.text
if out.tool_calls: # fc tasks: serialize calls as the prediction
import json
text = (text + '\n' if text else '') + json.dumps(
[c.to_openai()['function'] for c in out.tool_calls], ensure_ascii=False)
done_count += 1
_progress(progress, done_count, len(samples), t0, total_usage)
return {'raw': text, 'usage': out.usage.model_dump()}
work = _apply_limits(samples, limit, limit_per_task)
# checkpointing: restore completed samples, generate only the rest
ckpt_store = None
if checkpoint:
from ..eval.checkpoint import CheckpointStore, checkpoint_path
if isinstance(checkpoint, str):
ckpt = checkpoint
else:
ckpt = checkpoint_path(os.path.expanduser('~/.cache/evalharness'),
dataset_name, adapter.model or str(adapter))
ckpt_store = CheckpointStore(ckpt, model=adapter.model or str(adapter))
restored = ckpt_store.load()
else:
restored = {}
keys = []
pending = []
preds_by_key: Dict[str, Dict[str, Any]] = {}
for i, s in enumerate(work):
k = CheckpointStore.key_for(s, i) if ckpt_store else str(i)
keys.append(k)
if k in restored:
preds_by_key[k] = restored[k]
else:
pending.append((i, s))
if ckpt_store is not None and restored:
print(f'checkpoint: restored {len(restored)} predictions '
f'({len(pending)} to generate) -> {ckpt_store.path}', flush=True)
async def run_one(i_s):
i, s = i_s
pred = await one(s)
if ckpt_store is not None:
ckpt_store.append(keys[i], pred)
return i, pred
fresh = await asyncio.gather(*(run_one((i, s)) for i, s in pending))
for i, pred in fresh:
preds_by_key[keys[i]] = pred
preds = [preds_by_key[k] for k in keys]
usages = [p.get('usage', {}) for p in preds]
return preds, usages, total_usage
def _apply_limits(samples: List[Sample], total: Optional[int],
per_task: Optional[int], dataset=None,
shuffle: bool = True, seed: int = 42) -> List[Sample]:
"""total: cap the WHOLE run (ours semantics). per_task: cap each subset/
category (evalscope's --limit semantics) -- first N per group_key.
shuffle+seed mirror evalscope run.py: dataset_args.shuffle=True with
--seed 42 -> random.Random(seed).shuffle BEFORE limiting, so both
frameworks evaluate the IDENTICAL sample subset."""
if shuffle:
import random
random.Random(seed).shuffle(samples)
if per_task:
seen: Dict[str, int] = {}
out = []
for s in samples:
key = str((s.metadata or {}).get('category')
or (s.metadata or {}).get('subject')
or (s.metadata or {}).get('test_category')
or getattr(getattr(dataset, 'spec', None), 'subset', 'default'))
if seen.get(key, 0) < per_task:
seen[key] = seen.get(key, 0) + 1
out.append(s)
samples = out
if total:
samples = samples[:total]
return samples
def _progress(progress: bool, done: int, total: int, t0: float, usage: Usage) -> None:
if progress and (done % 20 == 0 or done == total):
rate = done / max(time.time() - t0, 1e-6)
print(f' [{done}/{total}] {rate:.1f} samples/s tokens={usage.total_tokens}', flush=True)
async def run_eval(
dataset: Union[Dataset, List[Sample]],
model_spec: str,
recipe: Optional[EvalRecipe] = None,
*,
concurrency: int = 32,
limit: Optional[int] = None,
gen_kwargs: Optional[Dict[str, Any]] = None,
judge_spec: Optional[str] = None,
judge: Optional[Any] = None,
progress: bool = True,
env: str = '',
system: str = '',
max_turns: int = 8,
max_input_chars: int = 0,
max_input_tokens: int = 0,
limit_per_task: Optional[int] = None,
checkpoint: Union[bool, str] = False,
dataset_name: str = 'adhoc',
few_shot_num: int = -1,
prompt_style: str = 'strict_letter',
) -> EvalReport:
"""Generate + score in one call. Model spec examples:
'mock', 'mock:boxed', 'openai/http://gpu03:8000/v1?qwen3-8b', 'deploy:vllm/qwen3-8b'.
env: environment plugin name ('bfcl_mock') -> agent message pump per
sample; omit for single-turn generation.
'auto' = logprob when the adapter supports it.
few_shot_num: -1 = the dataset's declared paper default (mmlu 5, bbh 3,
gsm8k 4, ...); 0 = zero-shot; N = explicit override.
prompt_style: 'strict_letter' (default, evalscope-style 'ANSWER: X')
| 'cot' (reasoning-friendly).
"""
spec = getattr(dataset, 'spec', None)
if few_shot_num < 0:
few_shot_num = (spec.few_shot_num if spec is not None else 0)
adapter = _make_adapter(model_spec)
name = spec.name if spec is not None else 'adhoc'
if recipe is None:
from ..eval.recipe import EvalRecipe, get_eval
try:
recipe = get_eval(name)
except KeyError:
if name != 'adhoc':
raise
recipe = EvalRecipe(name='adhoc', extract='identity',
scorers={'acc': {'name': 'exact', 'mode': 'raw'}})
samples = list(dataset)[:limit] if limit else list(dataset)
samples = _apply_limits(samples, limit, limit_per_task)
if progress:
mode = f'agent env={env}' if env else 'single-turn'
print(f'generating: {adapter} on {len(samples)} samples '
f'({mode}, concurrency={concurrency})', flush=True)
env_factory = None
if env:
from ..agent import ENV_REGISTRY, get_env
if env not in ENV_REGISTRY:
raise KeyError(f'unknown env {env!r}; available: {", ".join(ENV_REGISTRY.names())}')
probe = get_env(env)
if getattr(probe, 'needs_adapter', False):
env_factory = lambda: get_env(env, adapter=adapter) # noqa: E731
else:
env_factory = lambda: get_env(env) # noqa: E731
# paper-faithful few-shot exemplars: official hand-written hooks first
# (bbh CoT), else the dataset's own dev/train split
few_shot_samples = None
few_shot_text = None
if few_shot_num:
from ..data.registry import get_dataset_provider
prov = get_dataset_provider(name)
hook = getattr(prov, 'few_shot_hook', None)
if hook is not None:
few_shot_text = hook('', spec.subset if spec else 'default', few_shot_num)
if few_shot_text is None:
fs_split = (spec.few_shot_split if spec is not None else None) or 'dev'
try:
import dataclasses
fs_spec = dataclasses.replace(spec, split=fs_split) if spec is not None else None
if fs_spec is not None:
from ..data.loader import load_raw_records
fn = prov.resolve_record_fn()
fs_raw = load_raw_records(fs_spec)
few_shot_samples = [fn(r) for r in fs_raw[:few_shot_num]]
except Exception as e:
print(f'few-shot: could not load {fs_split} split ({type(e).__name__}: '
f'{str(e)[:80]}); continuing 0-shot', flush=True)
try:
preds, _usages, usage = await generate_predictions(
adapter, samples, concurrency, progress=progress,
gen_kwargs={**(getattr(spec, 'gen_config', {}) or {}), **(gen_kwargs or {})},
env_factory=env_factory,
system=system, max_turns=max_turns, max_input_chars=max_input_chars,
max_input_tokens=max_input_tokens,
dataset_spec=spec,
limit_per_task=limit_per_task,
checkpoint=checkpoint,
dataset_name=name,
few_shot_num=few_shot_num,
few_shot_samples=few_shot_samples, few_shot_text=few_shot_text,
prompt_style=prompt_style)
finally:
await adapter.close()
if judge is None and judge_spec:
judge_adapter = _make_adapter(judge_spec)
judge = _judge_callable(judge_adapter)
report = evaluate(
samples, preds, recipe,
model=model_spec,
judge=judge,
extra_metadata={'gen_input_tokens': usage.input_tokens,
'gen_output_tokens': usage.output_tokens,
'gen_total_tokens': usage.total_tokens},
)
report.model = model_spec
report.dataset = name
# performance profile: pool success rate + latency/ttft percentiles
try:
from .aggregator import get_aggregator
perf = get_aggregator('perf_stats')(report.samples, 'acc')
if hasattr(adapter, 'stats'):
perf.update({f'pool_{k}': round(v, 3) if isinstance(v, float) else v
for k, v in adapter.request_stats().items()})
report.metric_groups['perf'] = perf
except Exception:
pass
return report
def _make_adapter(spec: str) -> ModelAdapter:
"""Model spec forms:
- 'mock[:mode]' offline adapter
- 'openai-pool/<base-url-template>?model' with {port} placeholder:
e.g. 'openai-pool/http://127.0.0.1:{8123..8130}/v1?Qwen3-8B' -> N ports
- else resolve_adapter(spec) single endpoint
Pooled specs are CACHED per spec: all benches share one pool so the
round-robin counter stays global (independent pools would each restart
at the first backend and starve the rest).
"""
from .adapter import _ADAPTER_CACHE as _CACHE
cache_key = spec
if cache_key in _CACHE:
return _CACHE[cache_key]
opts = {}
while True:
for f in ('!nothink', '!textools', '!perf'):
if spec.endswith(f):
spec = spec[:-len(f)]
opts[f] = True
break
else:
break
if spec.startswith('openai-pool/'):
from .pool import pooled
rest = spec[len('openai-pool/'):]
m = __import__('re').search(r'\{(\d+)\.\.(\d+)\}', rest)
if not m:
raise ValueError("openai-pool needs a {start..end} port range")
lo, hi = int(m.group(1)), int(m.group(2))
base_url, _, model = rest.partition('?')
specs = []
for port in range(lo, hi + 1):
specs.append(f'openai/{base_url.replace(m.group(0), str(port))}?{model}')
adapter = pooled(specs)
elif spec.partition(':')[0] == 'mock' and ':' in spec and '/' not in spec.partition(':')[0]:
adapter = resolve_adapter('mock')
adapter.extra['mode'] = spec.partition(':')[2] or 'echo'
return adapter
else:
adapter = resolve_adapter(spec)
members = adapter.adapters if hasattr(adapter, 'adapters') else [adapter]
for a in members:
if opts.get('!nothink'):
a.extra['no_think'] = True
if opts.get('!textools'):
a.extra['tools_mode'] = 'text'
if opts.get('!perf'):
a.extra['collect_perf'] = True
_CACHE[cache_key] = adapter
return adapter
def _judge_callable(judge_adapter: ModelAdapter):
"""Sync judge bridge. Works inside a running event loop (evaluate() may be
called from async run_eval): the coroutine runs on a private loop in a
worker thread."""
def ask(messages) -> str:
import asyncio
if isinstance(messages, list) and messages and isinstance(messages[0], dict):
messages = [ChatMessage(role=m.get('role', 'user'), content=m.get('content', ''))
for m in messages]
async def go():
out = await judge_adapter.generate(messages)
return out.text
try:
asyncio.get_running_loop()
except RuntimeError:
return asyncio.run(go())
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
return pool.submit(asyncio.run, go()).result()
return ask