169 lines
6.4 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 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,
) -> tuple:
"""Fan out model calls; returns (raws, total_usage).
Each sample becomes one user message (or its ChatMessage list is used
verbatim for multi-turn samples). Tool declarations from sample.tools
are passed through so fc/agent recipes degrade gracefully today and
agent loops can reuse this adapter untouched.
"""
gen_kwargs = gen_kwargs or {}
sem = asyncio.Semaphore(concurrency)
total_usage = Usage()
done_count = 0
t0 = time.time()
raws: List[str] = []
usages: List[Dict[str, Any]] = []
async def one(sample: Sample) -> tuple:
nonlocal done_count, total_usage
messages = ([ChatMessage(role='user', content=sample.input)] 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') == 'boxed' \
and sample.target not in ('', None):
# oracle channel for mock:boxed so full pipelines verify 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)
usage = out.usage.model_dump()
done_count += 1
if progress and (done_count % 20 == 0 or done_count == len(samples)):
rate = done_count / max(time.time() - t0, 1e-6)
print(f' [{done_count}/{len(samples)}] {rate:.1f} samples/s '
f'tokens={total_usage.total_tokens}', flush=True)
return text, usage
work = samples[:limit] if limit else samples
pairs = await asyncio.gather(*(one(s) for s in work))
raws = [p[0] for p in pairs]
usages = [p[1] for p in pairs]
return raws, usages, total_usage
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,
) -> EvalReport:
"""Generate + score in one call. Model spec examples:
'mock', 'mock:boxed', 'openai/http://gpu03:8000/v1?qwen3-8b', 'deploy:vllm/qwen3-8b'.
"""
adapter = _make_adapter(model_spec)
spec = getattr(dataset, 'spec', None)
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)
if progress:
print(f'generating: {adapter} on {len(samples)} samples '
f'(concurrency={concurrency})', flush=True)
try:
raws, usages, usage = await generate_predictions(adapter, samples, concurrency,
progress=progress, gen_kwargs=gen_kwargs)
finally:
await adapter.close()
if judge is None and judge_spec:
judge_adapter = _make_adapter(judge_spec)
judge = _judge_callable(judge_adapter)
preds = [{'raw': r, 'usage': u} for r, u in zip(raws, usages)]
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
return report
def _make_adapter(spec: str) -> ModelAdapter:
"""'mock:boxed' -> MockAdapter(mode='boxed'); else resolve_adapter().
The colon-mode syntax exists ONLY for 'mock': adapter names contain no
scheme/colon, so 'mock:xxx' is safe while URLs ('openai/http://...')
must never be split on ':'.
"""
base, sep, mode = spec.partition(':')
if sep and '/' not in base and base == 'mock':
adapter = resolve_adapter('mock')
adapter.extra['mode'] = mode or 'echo'
return adapter
return resolve_adapter(spec)
def _judge_callable(judge_adapter: ModelAdapter):
async def ask(messages) -> str:
out = await judge_adapter.generate([ChatMessage(role='user', content=str(m)) for m in messages]
if isinstance(messages, list) and messages and isinstance(messages[0], dict)
else messages)
return out.text
import asyncio
def sync_ask(messages):
return asyncio.run(ask(messages))
return sync_ask