205 lines
7.7 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,
env_factory=None,
system: str = '',
max_turns: int = 8,
) -> tuple:
"""Fan out model calls; returns (pred-dicts, total_usage).
Without env_factory: single-turn generation (text or tool-call JSON).
With env_factory(sample)->Environment: the agent message pump runs per
sample and predictions carry trajectory/env_state/usage.
"""
gen_kwargs = gen_kwargs or {}
sem = asyncio.Semaphore(concurrency)
total_usage = Usage()
done_count = 0
t0 = time.time()
usages: List[Dict[str, Any]] = []
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
async with sem:
traj = await drive(adapter, sample, env=env_factory(),
max_turns=max_turns, system=system)
total_usage = total_usage + traj.usage
pred = trajectory_to_prediction(traj)
pred['group_key'] = str(sample.metadata.get('test_category')
or sample.metadata.get('category')
or sample.metadata.get('id') or sample.id or '')
done_count += 1
_progress(progress, done_count, len(samples), t0, total_usage)
return pred
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') 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 = samples[:limit] if limit else samples
preds = list(await asyncio.gather(*(one(s) for s in work)))
usages = [p.get('usage', {}) for p in preds]
return preds, usages, total_usage
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,
) -> 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.
"""
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:
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 BFCLEnvironment
envs = {'bfcl_mock': BFCLEnvironment}
if env not in envs:
raise KeyError(f"unknown env {env!r}; available: {', '.join(envs)}")
env_factory = envs[env]
try:
preds, _usages, usage = await generate_predictions(
adapter, samples, concurrency, progress=progress,
gen_kwargs=gen_kwargs, env_factory=env_factory,
system=system, max_turns=max_turns)
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
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