166 lines
6.9 KiB
Python

"""The evaluation runner: Dataset x predictions -> EvalReport.
Pure orchestration, no I/O hidden inside: predictions arrive as a list
(loaded from a jsonl of model outputs, a Session store, or built inline),
results aggregate into an EvalReport that visualizers consume.
Judge wiring: pass judge=<callable(messages)->str> once a ModelAdapter
exists; llm_judge recipes work immediately after that, no recipe change.
"""
import traceback
from typing import Callable, Dict, Iterable, List, Optional, Sequence, Union
from ..data.dataset import Dataset
from ..data.sample import Sample
from .aggregator import mean as _mean_agg
from .recipe import EvalRecipe
from .record import EvalReport, SampleResult
from .scorer import ScoreContext
def evaluate(
dataset: Union[Dataset, List[Sample]],
predictions: Sequence[Union[str, Dict]],
recipe: Optional[EvalRecipe] = None,
*,
model: str = '',
judge: Optional[Callable] = None,
extra_metadata: Optional[Dict] = None,
) -> EvalReport:
"""Score a dataset against raw predictions.
dataset: a Dataset or a plain list of Samples (views/slices).
predictions: str per sample (raw model output) or dicts with
{'raw': str, 'group_key': ..., 'metadata': {...}} overrides.
"""
samples: List[Sample] = list(dataset)
spec = getattr(dataset, 'spec', None)
ds_name = spec.name if spec is not None else samples[0].metadata.get('dataset', 'adhoc') if samples else 'adhoc'
ds_subset = spec.subset if spec is not None else ''
if len(predictions) != len(samples):
raise ValueError(f'{len(predictions)} predictions for {len(samples)} samples')
if recipe is None:
from .recipe import get_eval
recipe = get_eval(ds_name)
extractor = recipe.resolve_extract()
scorers = recipe.resolve_scorers()
aggregators = recipe.resolve_aggregators()
ctx = ScoreContext(judge=judge, params={})
# If any scorer executes in docker with per-sample images, overlap pulls
# with scoring (run sample N while N+1..N+lookahead images download).
bp = None
if _needs_bg_prefetch(recipe, samples):
from ..sandbox import BackgroundPrefetcher, images_for_samples
bp = BackgroundPrefetcher(images_for_samples(samples), workers=4, lookahead=8)
bp.__enter__()
results: List[SampleResult] = []
for sample, pred in zip(samples, predictions):
raw = pred if isinstance(pred, str) else str(pred.get('raw', ''))
override = {} if isinstance(pred, str) else pred
result = SampleResult(
sample_id=sample.id,
dataset=ds_name,
subset=ds_subset,
task_type=sample.task_type,
raw_prediction=raw,
target=sample.target,
group_key=str(override.get('group_key')
or sample.metadata.get('group_key')
or (sample.metadata.get('task_id') or sample.metadata.get('id') or '')),
metadata={k: v for k, v in (sample.metadata or {}).items()
if k in ('category', 'subject', 'test_category', 'bin', 'difficulty')},
)
if isinstance(pred, dict) and pred.get('metadata'):
result.metadata.update(pred['metadata'])
if isinstance(pred, dict) and pred.get('trajectory'):
result.trajectory = pred['trajectory']
if isinstance(pred, dict) and pred.get('env_state'):
result.env_state = pred['env_state']
if isinstance(pred, dict) and pred.get('usage'):
result.usage = pred['usage']
try:
if bp is not None and sample.sandbox and sample.sandbox.image:
bp.ensure(sample.sandbox.image) # wait only if this one still pulling
value, ok, note = extractor(raw, sample)
result.extracted_prediction = value
result.extraction_ok = ok
result.extraction_note = note
if not ok:
result.extraction_note = note or 'extractor returned not-ok'
for metric, scorer in scorers.items():
try:
sctx = ctx
if result.env_state and 'env_state' not in ctx.params:
sctx = ScoreContext(judge=ctx.judge, judge_model=ctx.judge_model,
params={**ctx.params,
'env_state': result.env_state})
scores, details = scorer(value if ok else '', sample.target, sample, sctx)
result.scores.update(scores)
result.score_details.update(details)
except Exception as e: # one metric failing must not kill the run
result.scores[metric] = 0.0
result.score_details[metric] = {'error': f'{type(e).__name__}: {e}'}
except Exception as e:
result.error = f'{type(e).__name__}: {e}\n{traceback.format_exc(limit=2)}'
results.append(result)
report = EvalReport(
dataset=ds_name,
recipe=recipe.name or dataset.spec.name,
model=model,
num_samples=len(results),
num_failed_extractions=sum(1 for r in results if not r.extraction_ok),
samples=results,
)
_aggregate_into(report, results, recipe, aggregators)
if bp is not None:
report.metric_groups['run_info'] = {
**report.metric_groups.get('run_info', {}),
**{f'img_{k}': v for k, v in bp.stats().items()},
}
bp.__exit__(None, None, None)
if extra_metadata:
report.metric_groups['run_info'] = {k: v for k, v in extra_metadata.items()
if isinstance(v, (int, float, str))}
return report
def _needs_bg_prefetch(recipe, samples) -> bool:
"""True when the recipe executes in docker AND samples declare images."""
try:
for spec in recipe.scorers.values():
params = spec if isinstance(spec, dict) else {}
if params.get('name') == 'execution' and params.get('sandbox') == 'docker':
return any(s.sandbox and s.sandbox.image for s in samples[:50])
except Exception:
return False
return False
def _aggregate_into(report: EvalReport, results, recipe: EvalRecipe, aggregators) -> None:
for metric in recipe.scorers:
agg = aggregators.get(metric)
if agg is None:
agg = _mean_agg
try:
out = agg(results, metric)
except Exception as e:
report.metric_groups[f'agg_error_{metric}'] = {'error': str(e)[:200]}
continue
if isinstance(out, dict):
report.metric_groups[metric] = out
vals = [v for v in out.values() if isinstance(v, (int, float))]
if vals:
report.metrics[metric] = sum(vals) / len(vals)
else:
report.metrics[metric] = float(out)
report.metrics['extraction_failure_rate'] = (
report.num_failed_extractions / report.num_samples if report.num_samples else 0.0
)