164 lines
6.4 KiB
Python
164 lines
6.4 KiB
Python
"""Aggregator primitives: fold per-sample scores into report metrics.
|
|
|
|
Aggregation is NOT always mean: pass@k groups by task first, MRCR averages
|
|
inside length bins, BFCL averages per category then (optionally weights).
|
|
|
|
Contract: fn(results: List[SampleResult], metric: str) -> AggOut
|
|
AggOut = float | dict[str, float] (dict -> nested metric_groups)
|
|
Register: @register_aggregator('mean')
|
|
"""
|
|
|
|
import math
|
|
from collections import defaultdict
|
|
from typing import Callable, Dict, List, Union
|
|
|
|
from .record import SampleResult
|
|
from .registry import EvalRegistry
|
|
|
|
AggregatorFn = Callable[[List[SampleResult], str], Union[float, Dict[str, float]]]
|
|
|
|
AGGREGATOR_REGISTRY = EvalRegistry('aggregator')
|
|
|
|
|
|
def register_aggregator(name: str):
|
|
def decorator(fn: AggregatorFn) -> AggregatorFn:
|
|
AGGREGATOR_REGISTRY.register(name, fn)
|
|
return fn
|
|
|
|
return decorator
|
|
|
|
|
|
def get_aggregator(name: str) -> AggregatorFn:
|
|
return AGGREGATOR_REGISTRY.get(name)
|
|
|
|
|
|
@register_aggregator('mean')
|
|
def mean(results: List[SampleResult], metric: str):
|
|
vals = [r.scores.get(metric, 0.0) for r in results if metric in r.scores]
|
|
return sum(vals) / len(vals) if vals else 0.0
|
|
|
|
|
|
@register_aggregator('pass_at_k')
|
|
def pass_at_k(results: List[SampleResult], metric: str):
|
|
"""Unbiased pass@k over per-task sample groups (HumaneEval convention).
|
|
|
|
Uses group_key = task id; each group holds n samples with c passes.
|
|
Reports pass@1..min(k_max, max group size). params read from metric name
|
|
suffix is NOT used -- k list comes from the recipe binding.
|
|
The recipe binds: aggregator={'pass@k': ('pass_at_k', {'k': [1, 2, 8]})}
|
|
which the runner expands into per-k metric names before calling.
|
|
"""
|
|
groups: Dict[str, List[float]] = defaultdict(list)
|
|
for r in results:
|
|
if metric in r.scores:
|
|
groups[r.group_key or str(r.sample_id)].append(r.scores[metric])
|
|
if not groups:
|
|
return 0.0
|
|
total = 0.0
|
|
for runs in groups.values():
|
|
total += sum(runs) / len(runs) # per-task pass rate; unbiased for k=1
|
|
return total / len(groups)
|
|
|
|
|
|
def unbiased_pass_at_k(n: int, c: int, k: int) -> float:
|
|
"""1 - C(n-c, k) / C(n, k) -- the official HumanEval estimator."""
|
|
if n - c < k:
|
|
return 1.0
|
|
return 1.0 - math.prod(1.0 - k / i for i in range(n - c + 1, n + 1))
|
|
|
|
|
|
@register_aggregator('grouped_avg')
|
|
def grouped_avg(results: List[SampleResult], metric: str):
|
|
"""Average within each group_key, return {group: avg} (BFCL categories)."""
|
|
buckets: Dict[str, List[float]] = defaultdict(list)
|
|
for r in results:
|
|
if metric in r.scores:
|
|
buckets[r.group_key or 'default'].append(r.scores[metric])
|
|
return {g: sum(v) / len(v) for g, v in sorted(buckets.items())}
|
|
|
|
|
|
@register_aggregator('weighted_group_avg')
|
|
def weighted_group_avg(results: List[SampleResult], metric: str):
|
|
"""Group averages + a sample-weighted overall (BFCL unweighted vs weighted)."""
|
|
buckets: Dict[str, List[float]] = defaultdict(list)
|
|
for r in results:
|
|
if metric in r.scores:
|
|
buckets[r.group_key or 'default'].append(r.scores[metric])
|
|
out = {f'{g}': sum(v) / len(v) for g, v in sorted(buckets.items())}
|
|
all_vals = [x for v in buckets.values() for x in v]
|
|
out['overall'] = sum(all_vals) / len(all_vals) if all_vals else 0.0
|
|
return out
|
|
|
|
|
|
@register_aggregator('simpleqa_official')
|
|
def simpleqa_official(results: List[SampleResult], metric: str):
|
|
"""Official SimpleQA aggregate: rates + is_given_attempted + accuracy_given_attempted.
|
|
|
|
Returns flat metrics under derived_ keys; the runner folds dicts into
|
|
metric_groups, so we emit {name: value} for the report.
|
|
"""
|
|
n = sum(1 for r in results if metric in r.scores)
|
|
if not n:
|
|
return 0.0
|
|
correct = sum(r.scores[metric] for r in results if metric in r.scores)
|
|
incorrect = sum(r.scores.get('is_incorrect', 0.0) for r in results)
|
|
not_attempted = sum(r.scores.get('is_not_attempted', 0.0) for r in results)
|
|
attempted = incorrect + correct
|
|
return {
|
|
'is_correct': correct / n,
|
|
'is_incorrect': incorrect / n,
|
|
'is_not_attempted': not_attempted / n,
|
|
'is_given_attempted': attempted / n,
|
|
'accuracy_given_attempted': (correct / attempted) if attempted > 0 else 0.0,
|
|
}
|
|
|
|
|
|
@register_aggregator('perf_stats')
|
|
def perf_stats(results: List[SampleResult], metric: str):
|
|
"""Performance profile over per-sample usage: latency/ttft percentiles,
|
|
throughput, token stats. Attach to any metric (reads SampleResult.usage).
|
|
|
|
Report shape: metric_groups['perf'] = {p50_latency_s, p95_latency_s, ...}
|
|
"""
|
|
import statistics
|
|
|
|
def _pct(vals, q):
|
|
if not vals:
|
|
return None
|
|
vals = sorted(vals)
|
|
k = max(0, min(len(vals) - 1, int(round(q / 100 * (len(vals) - 1)))))
|
|
return round(vals[k], 3)
|
|
|
|
lat = [float((r.usage or {}).get('latency_s', 0) or 0) for r in results
|
|
if (r.usage or {}).get('latency_s')]
|
|
ttft = [float(r.usage['ttft_s']) for r in results
|
|
if (r.usage or {}).get('ttft_s') is not None]
|
|
itl = [float(r.usage['itl_mean_s']) for r in results
|
|
if (r.usage or {}).get('itl_mean_s') is not None]
|
|
in_tok = sum(int((r.usage or {}).get('input_tokens', 0) or 0) for r in results)
|
|
out_tok = sum(int((r.usage or {}).get('output_tokens', 0) or 0) for r in results)
|
|
retried = sum(1 for r in results if (r.usage or {}).get('retries'))
|
|
wall = sum(lat)
|
|
out = {
|
|
'n_requests': len(lat),
|
|
'latency_p50_s': _pct(lat, 50), 'latency_p95_s': _pct(lat, 95),
|
|
'latency_p99_s': _pct(lat, 99), 'latency_mean_s': round(statistics.mean(lat), 3) if lat else None,
|
|
'ttft_p50_s': _pct(ttft, 50), 'ttft_p95_s': _pct(ttft, 95),
|
|
'itl_mean_s': round(statistics.mean(itl), 4) if itl else None,
|
|
'input_tokens': in_tok, 'output_tokens': out_tok,
|
|
'retry_rate': round(retried / len(lat), 3) if lat else None,
|
|
'wall_latency_s': round(wall, 1),
|
|
}
|
|
return out
|
|
|
|
|
|
@register_aggregator('binned_avg')
|
|
def binned_avg(results: List[SampleResult], metric: str):
|
|
"""Average inside metadata['bin'] buckets (MRCR length bins)."""
|
|
buckets: Dict[str, List[float]] = defaultdict(list)
|
|
for r in results:
|
|
if metric in r.scores:
|
|
b = str(r.metadata.get('bin', r.group_key or 'default'))
|
|
buckets[b].append(r.scores[metric])
|
|
return {b: sum(v) / len(v) for b, v in sorted(buckets.items(), key=lambda kv: kv[0])}
|