diff --git a/evalharness/model/runner.py b/evalharness/model/runner.py index b03a225..fef72b3 100644 --- a/evalharness/model/runner.py +++ b/evalharness/model/runner.py @@ -443,6 +443,20 @@ async def generate_predictions( preds_by_key[keys[i]] = pred preds = [preds_by_key[k] for k in keys] usages = [p.get('usage', {}) for p in preds] + # include RESTORED predictions' usage (they carry it in the ckpt); + # previously only fresh generations counted -> restored benches showed 0 + fresh_keys = {keys[i] for i, _ in pending} + for k, p in preds_by_key.items(): + if k in fresh_keys: + continue # already counted via one()'s total_usage updates + u = p.get('usage') or {} + if not any(u.get(kk) for kk in ('input_tokens', 'output_tokens')): + continue + 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)) if status_callback and pending: status_callback(f'Generation complete: {len(preds)} responses collected') return preds, usages, total_usage