tau2 via OFFICIAL engine: self-running env plugin (run_task hook + needs_adapter generic dispatch, no name hardcoding), deep generate() patch, official ToolCall shape, data plugin keeps verbatim Task json, official reward scoring; checkpoint plugin (per-sample resume) + summary csv with time/categories
This commit is contained in:
parent
85dd193bcf
commit
78459c974e
@ -107,7 +107,8 @@ class BFCLEnvironment(Environment):
|
|||||||
|
|
||||||
name = 'bfcl_mock'
|
name = 'bfcl_mock'
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self, adapter=None):
|
||||||
|
super().__init__(adapter=adapter)
|
||||||
self.calls: List[Dict[str, Any]] = []
|
self.calls: List[Dict[str, Any]] = []
|
||||||
self.ground_truth: Dict[str, Any] = {}
|
self.ground_truth: Dict[str, Any] = {}
|
||||||
|
|
||||||
|
|||||||
144
evalharness/agent/envs/tau2_official.py
Normal file
144
evalharness/agent/envs/tau2_official.py
Normal file
@ -0,0 +1,144 @@
|
|||||||
|
"""tau2-bench environment: official tau2 package as the ENGINE, we only bridge.
|
||||||
|
|
||||||
|
Same architecture as evalscope's adapter: the official user simulator,
|
||||||
|
domain environments (airline db/policy/...), and reward logic are used
|
||||||
|
VERBATIM; we monkey-patch tau2's LLM generate() to call our ModelAdapter,
|
||||||
|
so scoring is bit-identical to the official tool by construction.
|
||||||
|
|
||||||
|
The env consumes a Sample whose metadata carries the official Task json
|
||||||
|
(our tau2_bench dataset plugin stores it in metadata['task']).
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
|
from ...data.sample import ChatMessage, Sample
|
||||||
|
from ..loop import Environment, register_env
|
||||||
|
|
||||||
|
_PATCHED = False
|
||||||
|
|
||||||
|
|
||||||
|
def _patch_tau2_generate(adapter) -> None:
|
||||||
|
"""Route tau2's LLM calls to our ModelAdapter (sync bridge via asyncio)."""
|
||||||
|
global _PATCHED
|
||||||
|
if _PATCHED:
|
||||||
|
return
|
||||||
|
|
||||||
|
import tau2.utils.llm_utils as llm_utils
|
||||||
|
|
||||||
|
original = llm_utils.generate
|
||||||
|
|
||||||
|
def patched_generate(model, messages, tools=None, tool_choice=None, **kw):
|
||||||
|
# model: 'user' | 'agent' -> same adapter in our single-model setup
|
||||||
|
msgs = [ChatMessage(role=m.role if hasattr(m, 'role') else 'user',
|
||||||
|
content=m.content if hasattr(m, 'content') else str(m))
|
||||||
|
for m in messages]
|
||||||
|
tool_specs = None
|
||||||
|
if tools:
|
||||||
|
tool_specs = []
|
||||||
|
for t in tools:
|
||||||
|
spec = t.openai_schema if hasattr(t, 'openai_schema') else t
|
||||||
|
if isinstance(spec, dict) and 'function' not in spec:
|
||||||
|
spec = {'type': 'function', 'function': spec}
|
||||||
|
tool_specs.append(spec)
|
||||||
|
|
||||||
|
async def go():
|
||||||
|
return await adapter.generate(msgs, tools=tool_specs)
|
||||||
|
|
||||||
|
try:
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
except RuntimeError:
|
||||||
|
out = asyncio.run(go())
|
||||||
|
else:
|
||||||
|
import concurrent.futures
|
||||||
|
|
||||||
|
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
||||||
|
out = pool.submit(asyncio.run, go()).result()
|
||||||
|
|
||||||
|
# tau2 AssistantMessage shape
|
||||||
|
from tau2.data_model.message import AssistantMessage, ToolCall as TauToolCall
|
||||||
|
|
||||||
|
import json as _json
|
||||||
|
|
||||||
|
calls = []
|
||||||
|
for c in out.tool_calls:
|
||||||
|
try:
|
||||||
|
args = _json.loads(c.arguments or '{}')
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
args = {'raw': c.arguments}
|
||||||
|
calls.append(TauToolCall(id=c.id or f'call_{c.name}', name=c.name,
|
||||||
|
arguments=args or {}))
|
||||||
|
return AssistantMessage(role='assistant', content=out.text or '',
|
||||||
|
tool_calls=calls or None, cost=None,
|
||||||
|
usage=None, raw_data=None)
|
||||||
|
|
||||||
|
# patch module attr AND every already-imported reference (official
|
||||||
|
# modules do `from tau2.utils.llm_utils import generate`)
|
||||||
|
import importlib
|
||||||
|
import sys
|
||||||
|
|
||||||
|
llm_utils.generate = patched_generate
|
||||||
|
for mod_name, mod in list(sys.modules.items()):
|
||||||
|
if not mod_name.startswith('tau2'):
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
if getattr(mod, 'generate', None) is original:
|
||||||
|
mod.generate = patched_generate
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
_PATCHED = True
|
||||||
|
|
||||||
|
|
||||||
|
@register_env('tau2_official')
|
||||||
|
class Tau2Environment(Environment):
|
||||||
|
"""Per-sample wrapper around official run_task; reward from env final state."""
|
||||||
|
|
||||||
|
name = 'tau2_official'
|
||||||
|
needs_adapter = True
|
||||||
|
|
||||||
|
def __init__(self, adapter=None):
|
||||||
|
super().__init__(adapter=adapter)
|
||||||
|
self.reward_info: Dict[str, Any] = {}
|
||||||
|
|
||||||
|
def reset(self, sample: Sample) -> List[ChatMessage]:
|
||||||
|
self.reward_info = {}
|
||||||
|
return []
|
||||||
|
|
||||||
|
async def step(self, tool_calls, text: str, sample: Sample) -> List[ChatMessage]:
|
||||||
|
return [] # the official engine runs the whole loop itself
|
||||||
|
|
||||||
|
def final_state(self) -> Dict[str, Any]:
|
||||||
|
return self.reward_info
|
||||||
|
|
||||||
|
async def run_task(self, adapter, sample, max_turns: int = 40, **kw):
|
||||||
|
from tau2.data_model.tasks import Task
|
||||||
|
from tau2.run import run_task
|
||||||
|
|
||||||
|
_patch_tau2_generate(adapter)
|
||||||
|
task_json = (sample.metadata or {}).get('task')
|
||||||
|
if task_json is None:
|
||||||
|
raise ValueError("tau2 sample missing metadata['task'] "
|
||||||
|
'(dataset plugin must keep the official Task json)')
|
||||||
|
task = Task.model_validate(task_json if not isinstance(task_json, str)
|
||||||
|
else json.loads(task_json))
|
||||||
|
domain = (sample.metadata or {}).get('domain') or 'airline'
|
||||||
|
res = run_task(domain=domain, task=task, agent='llm_agent',
|
||||||
|
user='user_simulator', max_steps=max_turns)
|
||||||
|
rewards = {}
|
||||||
|
try:
|
||||||
|
info = res.reward_info
|
||||||
|
rewards = info.model_dump() if hasattr(info, 'model_dump') else dict(info)
|
||||||
|
env_r = rewards.get('environment_reward')
|
||||||
|
comm_r = rewards.get('communication_reward')
|
||||||
|
vals = [r for r in (env_r, comm_r) if isinstance(r, (int, float))]
|
||||||
|
self.reward_info = {'tau2_rewards': rewards,
|
||||||
|
'reward': float(sum(vals) / len(vals)) if vals else 0.0}
|
||||||
|
except Exception:
|
||||||
|
self.reward_info = {'tau2_rewards': rewards, 'reward': 0.0}
|
||||||
|
traj = [{'role': str(getattr(m, 'role', 'user')),
|
||||||
|
'content': str(getattr(m, 'content', '')), 'turn': 0}
|
||||||
|
for m in (res.messages or [])]
|
||||||
|
return {'raw': json.dumps(self.reward_info, default=str),
|
||||||
|
'trajectory': traj, 'env_state': self.reward_info,
|
||||||
|
'usage': {}, 'group_key': domain}
|
||||||
@ -16,24 +16,39 @@ from ..model.output import ModelOutput, Usage
|
|||||||
|
|
||||||
|
|
||||||
class Environment:
|
class Environment:
|
||||||
"""Minimal env contract. One instance per sample."""
|
"""Minimal env contract. One instance per sample.
|
||||||
|
|
||||||
|
Two execution models, auto-selected by the runner:
|
||||||
|
- MESSAGE PUMP (default): implement reset/step/final_state; the runner
|
||||||
|
drives the model turn-by-turn (bfcl_mock style).
|
||||||
|
- SELF-RUNNING (official engines): implement run_task(adapter, sample);
|
||||||
|
the whole simulation happens inside (official tau2/swe bundles), and
|
||||||
|
a prediction dict is returned directly. Falls back to the pump when
|
||||||
|
run_task is not overridden.
|
||||||
|
"""
|
||||||
|
|
||||||
name = 'base'
|
name = 'base'
|
||||||
|
needs_adapter = False # self-running engines set True; the runner passes
|
||||||
|
# the adapter to the constructor (no name hardcoding)
|
||||||
|
|
||||||
|
def __init__(self, adapter=None):
|
||||||
|
self.adapter = adapter
|
||||||
|
|
||||||
def reset(self, sample: Sample) -> List[ChatMessage]:
|
def reset(self, sample: Sample) -> List[ChatMessage]:
|
||||||
"""Prepare per-sample state; return any extra opening messages
|
|
||||||
(e.g. tool/user-simulator turns). Default: nothing."""
|
|
||||||
return []
|
return []
|
||||||
|
|
||||||
async def step(self, tool_calls: List[Any], text: str, sample: Sample
|
async def step(self, tool_calls: List[Any], text: str, sample: Sample
|
||||||
) -> List[ChatMessage]:
|
) -> List[ChatMessage]:
|
||||||
"""Execute the model's calls; return observation messages."""
|
|
||||||
raise NotImplementedError
|
raise NotImplementedError
|
||||||
|
|
||||||
def final_state(self) -> Dict[str, Any]:
|
def final_state(self) -> Dict[str, Any]:
|
||||||
"""Terminal state handed to env_reward scorers."""
|
|
||||||
return {}
|
return {}
|
||||||
|
|
||||||
|
async def run_task(self, adapter, sample: Sample, **kw) -> Optional[Dict[str, Any]]:
|
||||||
|
"""Self-running hook. Return a prediction dict (raw/trajectory/
|
||||||
|
env_state/usage/group_key) or None to fall back to the pump."""
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
# ---------------- environment registry (万物皆可插件: envs too) ----------------
|
# ---------------- environment registry (万物皆可插件: envs too) ----------------
|
||||||
|
|
||||||
@ -53,9 +68,9 @@ def register_env(name: str):
|
|||||||
return decorator
|
return decorator
|
||||||
|
|
||||||
|
|
||||||
def get_env(name: str):
|
def get_env(name: str, adapter=None):
|
||||||
cls = ENV_REGISTRY.get(name)
|
cls = ENV_REGISTRY.get(name)
|
||||||
return cls()
|
return cls(adapter=adapter)
|
||||||
|
|
||||||
|
|
||||||
class Trajectory:
|
class Trajectory:
|
||||||
|
|||||||
@ -149,6 +149,7 @@ def _cmd_eval_run(args) -> int:
|
|||||||
report = asyncio.run(run_eval(
|
report = asyncio.run(run_eval(
|
||||||
ds, args.model, concurrency=args.concurrency, limit=args.limit,
|
ds, args.model, concurrency=args.concurrency, limit=args.limit,
|
||||||
limit_per_task=args.limit_per_task,
|
limit_per_task=args.limit_per_task,
|
||||||
|
checkpoint=args.resume,
|
||||||
judge_spec=args.judge, env=args.env))
|
judge_spec=args.judge, env=args.env))
|
||||||
else:
|
else:
|
||||||
from evalharness.eval import evaluate
|
from evalharness.eval import evaluate
|
||||||
@ -170,10 +171,17 @@ def _cmd_eval_run(args) -> int:
|
|||||||
print(render(report, style=args.style))
|
print(render(report, style=args.style))
|
||||||
print(render(report, style=args.style))
|
print(render(report, style=args.style))
|
||||||
primary = next(iter(report.metrics), '')
|
primary = next(iter(report.metrics), '')
|
||||||
|
secs_total = sum(float((s.usage or {}).get('latency_s', 0) or 0)
|
||||||
|
for s in report.samples)
|
||||||
|
groups = {k: v for k, v in report.metric_groups.items()
|
||||||
|
if isinstance(v, dict) and k not in ('run_info',)
|
||||||
|
and not k.startswith('agg_error')}
|
||||||
rows.append({'name': name, 'metric': primary, 'value': report.metrics.get(primary),
|
rows.append({'name': name, 'metric': primary, 'value': report.metrics.get(primary),
|
||||||
'n': report.num_samples,
|
'n': report.num_samples,
|
||||||
'extract_fail': report.num_failed_extractions,
|
'extract_fail': report.num_failed_extractions,
|
||||||
'secs': round(_time.time() - t0, 1), 'ok': True})
|
'secs': round(secs_total, 1),
|
||||||
|
'hours': round(secs_total / 3600, 2),
|
||||||
|
'groups': groups, 'ok': True})
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
rows.append({'name': name, 'metric': '-', 'value': None,
|
rows.append({'name': name, 'metric': '-', 'value': None,
|
||||||
'secs': round(_time.time() - t0, 1), 'ok': False,
|
'secs': round(_time.time() - t0, 1), 'ok': False,
|
||||||
@ -181,11 +189,13 @@ def _cmd_eval_run(args) -> int:
|
|||||||
print(f'{name}: FAILED {type(e).__name__}: {str(e)[:160]}', file=sys.stderr)
|
print(f'{name}: FAILED {type(e).__name__}: {str(e)[:160]}', file=sys.stderr)
|
||||||
|
|
||||||
if len(rows) > 1:
|
if len(rows) > 1:
|
||||||
print(f'\n{"dataset":<20} {"metric":<14} {"value":<8} {"secs":>5}')
|
print(f'\n{"benchmark":<20} {"metric":<14} {"score":>8} {"n":>5} {"time":>9}')
|
||||||
print('-' * 52)
|
print('-' * 62)
|
||||||
for r in rows:
|
for r in rows:
|
||||||
val = 'ERR' if not r['ok'] else round(r['value'], 4)
|
val = 'ERR' if not r['ok'] else _f3(r['value'])
|
||||||
print(f"{r['name']:<20} {r['metric']:<14} {val!s:<8} {r['secs']:>5}"
|
h = r.get('hours') or 0
|
||||||
|
tdisp = f'{h:.2f}h' if h >= 0.995 else f"{r.get('secs', 0):.0f}s"
|
||||||
|
print(f"{r['name']:<20} {r['metric']:<14} {val!s:>8} {r.get('n', '')!s:>5} {tdisp:>9}"
|
||||||
+ (f" {r.get('err', '')}" if not r['ok'] else ''))
|
+ (f" {r.get('err', '')}" if not r['ok'] else ''))
|
||||||
ok = sum(1 for r in rows if r['ok'])
|
ok = sum(1 for r in rows if r['ok'])
|
||||||
print(f'\n{ok}/{len(rows)} ok' + (f' -> artifacts in {out_dir}/' if out_dir else ''))
|
print(f'\n{ok}/{len(rows)} ok' + (f' -> artifacts in {out_dir}/' if out_dir else ''))
|
||||||
@ -194,11 +204,16 @@ def _cmd_eval_run(args) -> int:
|
|||||||
|
|
||||||
with open(f'{out_dir}/viz/summary.csv', 'w', newline='', encoding='utf-8') as f:
|
with open(f'{out_dir}/viz/summary.csv', 'w', newline='', encoding='utf-8') as f:
|
||||||
w = _csv.writer(f)
|
w = _csv.writer(f)
|
||||||
w.writerow(['dataset', 'model', 'n', 'metric', 'value', 'extract_fail', 'secs'])
|
w.writerow(['benchmark', 'score', 'metric', 'num_samples',
|
||||||
|
'time_h', 'time_s', 'extract_fail', 'categories'])
|
||||||
for r in rows:
|
for r in rows:
|
||||||
w.writerow([r['name'], args.model, r.get('n', ''),
|
cats = '; '.join(f'{g}={_f3(v)}'
|
||||||
r['metric'], r.get('value', ''), r.get('extract_fail', ''),
|
for gname, gv in (r.get('groups') or {}).items()
|
||||||
r['secs']])
|
for g, v in (gv or {}).items()
|
||||||
|
if isinstance(v, (int, float)))[:2000]
|
||||||
|
w.writerow([r['name'], _f3(r.get('value')), r['metric'], r.get('n', ''),
|
||||||
|
r.get('hours', ''), r.get('secs', ''),
|
||||||
|
r.get('extract_fail', 0), cats])
|
||||||
with open(f'{out_dir}/viz/summary.md', 'w', encoding='utf-8') as f:
|
with open(f'{out_dir}/viz/summary.md', 'w', encoding='utf-8') as f:
|
||||||
f.write(f'# eval run summary\n\n| dataset | metric | value | secs |\n|---|---|---|---|\n')
|
f.write(f'# eval run summary\n\n| dataset | metric | value | secs |\n|---|---|---|---|\n')
|
||||||
for r in rows:
|
for r in rows:
|
||||||
@ -207,6 +222,13 @@ def _cmd_eval_run(args) -> int:
|
|||||||
return 0 if all(r['ok'] for r in rows) else 1
|
return 0 if all(r['ok'] for r in rows) else 1
|
||||||
|
|
||||||
|
|
||||||
|
def _f3(v):
|
||||||
|
try:
|
||||||
|
return round(float(v), 4)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return v
|
||||||
|
|
||||||
|
|
||||||
def _cmd_viz_show(args) -> int:
|
def _cmd_viz_show(args) -> int:
|
||||||
from evalharness.viz import render
|
from evalharness.viz import render
|
||||||
|
|
||||||
@ -264,6 +286,9 @@ def build_parser() -> argparse.ArgumentParser:
|
|||||||
p.add_argument('--env', default='', help="agent environment (e.g. 'bfcl_mock') -> message pump")
|
p.add_argument('--env', default='', help="agent environment (e.g. 'bfcl_mock') -> message pump")
|
||||||
p.add_argument('--concurrency', type=int, default=32, help='parallel model calls (default 32)')
|
p.add_argument('--concurrency', type=int, default=32, help='parallel model calls (default 32)')
|
||||||
p.add_argument('--limit', type=int, help='evaluate only the first N samples total')
|
p.add_argument('--limit', type=int, help='evaluate only the first N samples total')
|
||||||
|
p.add_argument('--resume', nargs='?', const=True, default=False,
|
||||||
|
help='resume from per-sample checkpoint (default path auto-derived; '
|
||||||
|
'pass a path to override)')
|
||||||
p.add_argument('--limit-per-task', type=int,
|
p.add_argument('--limit-per-task', type=int,
|
||||||
help='first N samples PER subset/category (evalscope --limit semantics); '
|
help='first N samples PER subset/category (evalscope --limit semantics); '
|
||||||
'composable with --limit (intersection)')
|
'composable with --limit (intersection)')
|
||||||
|
|||||||
@ -34,6 +34,10 @@ def tau2_bench():
|
|||||||
target='',
|
target='',
|
||||||
metadata={
|
metadata={
|
||||||
'id': record.get('id'),
|
'id': record.get('id'),
|
||||||
|
'domain': record.get('domain'),
|
||||||
|
'task': record, # OFFICIAL Task json, verbatim -- the
|
||||||
|
# tau2_official env runs the official
|
||||||
|
# engine on this; nothing is lost
|
||||||
'notes': desc.get('notes'),
|
'notes': desc.get('notes'),
|
||||||
'task_instructions': instructions,
|
'task_instructions': instructions,
|
||||||
'user_scenario': scenario,
|
'user_scenario': scenario,
|
||||||
|
|||||||
85
evalharness/eval/checkpoint.py
Normal file
85
evalharness/eval/checkpoint.py
Normal file
@ -0,0 +1,85 @@
|
|||||||
|
"""Resumable evaluation: checkpoint plugin.
|
||||||
|
|
||||||
|
Run-level checkpointing: every completed sample is appended to a jsonl
|
||||||
|
checkpoint file; on restart (crash/OOM/service down), completed samples are
|
||||||
|
restored and ONLY the missing ones are generated. The report is then
|
||||||
|
assembled from restored + fresh predictions.
|
||||||
|
|
||||||
|
from evalharness.eval.checkpoint import CheckpointStore
|
||||||
|
store = CheckpointStore(path='run.jsonl', model=model_spec, dataset=name)
|
||||||
|
done = store.load() # {sample_key: prediction-dict}
|
||||||
|
... generate only missing ...
|
||||||
|
store.append(sample_key, pred) # after EACH sample (crash-safe)
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, Dict, Optional
|
||||||
|
|
||||||
|
|
||||||
|
class CheckpointStore:
|
||||||
|
"""Append-only jsonl checkpoint; keyed by stable sample identity."""
|
||||||
|
|
||||||
|
def __init__(self, path: str, model: str = '', dataset: str = ''):
|
||||||
|
self.path = Path(path)
|
||||||
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
self.model = model
|
||||||
|
self.dataset = dataset
|
||||||
|
self._entries: Dict[str, Dict[str, Any]] = {}
|
||||||
|
self._fh = None
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def key_for(sample, idx: int) -> str:
|
||||||
|
"""Stable per-sample key: prefer dataset-native ids, fall back to
|
||||||
|
a hash of the input text (survives re-orderings)."""
|
||||||
|
native = (sample.metadata or {}).get('task_id') \
|
||||||
|
or (sample.metadata or {}).get('id') \
|
||||||
|
or (sample.metadata or {}).get('instance_id')
|
||||||
|
if native:
|
||||||
|
return str(native)
|
||||||
|
import hashlib
|
||||||
|
|
||||||
|
h = hashlib.md5((sample.input_text or '').encode('utf-8')).hexdigest()[:16]
|
||||||
|
return f'{idx}:{h}'
|
||||||
|
|
||||||
|
def load(self) -> Dict[str, Dict[str, Any]]:
|
||||||
|
"""Read all checkpointed predictions (idempotent)."""
|
||||||
|
self._entries = {}
|
||||||
|
if not self.path.exists():
|
||||||
|
return self._entries
|
||||||
|
with open(self.path, encoding='utf-8') as f:
|
||||||
|
for line in f:
|
||||||
|
line = line.strip()
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
rec = json.loads(line)
|
||||||
|
self._entries[rec['key']] = rec.get('pred', {})
|
||||||
|
except (ValueError, KeyError):
|
||||||
|
continue # torn tail line from a crash -- safe to skip
|
||||||
|
return self._entries
|
||||||
|
|
||||||
|
def append(self, key: str, pred: Dict[str, Any]) -> None:
|
||||||
|
"""Persist one prediction immediately (fsync-free append is fine:
|
||||||
|
worst case loses the last in-flight sample on crash)."""
|
||||||
|
rec = {'key': key, 'ts': time.time(), 'pred': pred}
|
||||||
|
with open(self.path, 'a', encoding='utf-8') as f:
|
||||||
|
f.write(json.dumps(rec, ensure_ascii=False) + '\n')
|
||||||
|
self._entries[key] = pred
|
||||||
|
|
||||||
|
def __len__(self) -> int:
|
||||||
|
return len(self._entries)
|
||||||
|
|
||||||
|
def summary(self) -> Dict[str, Any]:
|
||||||
|
return {'path': str(self.path), 'restored': len(self._entries),
|
||||||
|
'size_kb': self.path.stat().st_size // 1024 if self.path.exists() else 0}
|
||||||
|
|
||||||
|
|
||||||
|
def checkpoint_path(root: str, dataset: str, model: str, tag: str = '') -> str:
|
||||||
|
"""Deterministic path per (dataset, model[, tag]) so re-runs resume."""
|
||||||
|
import hashlib
|
||||||
|
|
||||||
|
h = hashlib.md5(f'{dataset}|{model}|{tag}'.encode()).hexdigest()[:10]
|
||||||
|
return os.path.join(root, 'ckpt', f'{dataset}-{h}.jsonl')
|
||||||
@ -185,13 +185,28 @@ def swe_bench_verified():
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _tau2_reward(pred, target, sample, ctx):
|
||||||
|
"""Score from the official engine's reward_info (env_state)."""
|
||||||
|
env_state = ctx.params.get('env_state') or {}
|
||||||
|
rewards = env_state.get('tau2_rewards') or {}
|
||||||
|
env_r = rewards.get('environment_reward')
|
||||||
|
comm_r = rewards.get('communication_reward')
|
||||||
|
vals = [r for r in (env_r, comm_r) if isinstance(r, (int, float))]
|
||||||
|
score = float(sum(vals) / len(vals)) if vals else 0.0
|
||||||
|
return ({'acc': score}, {'acc': {'mode': 'official_tau2',
|
||||||
|
'env_reward': env_r,
|
||||||
|
'comm_reward': comm_r}})
|
||||||
|
|
||||||
|
|
||||||
@register_eval('tau2_bench')
|
@register_eval('tau2_bench')
|
||||||
def tau2_bench():
|
def tau2_bench():
|
||||||
return EvalRecipe(
|
return EvalRecipe(
|
||||||
name='tau2_bench',
|
name='tau2_bench',
|
||||||
extract='identity',
|
extract='identity',
|
||||||
scorers={'acc': 'env_reward'},
|
scorers={'acc': _tau2_reward},
|
||||||
description='tau2-bench; user-simulated dialog, environment reward.',
|
aggregators={'acc': 'grouped_avg'},
|
||||||
|
description='tau2-bench via OFFICIAL engine (user simulator + env + reward); '
|
||||||
|
"run with env='tau2_official'",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@ -117,6 +117,9 @@ class OpenAICompatible(ModelAdapter):
|
|||||||
name = 'openai'
|
name = 'openai'
|
||||||
|
|
||||||
async def generate(self, messages, tools=None, **kw) -> ModelOutput:
|
async def generate(self, messages, tools=None, **kw) -> ModelOutput:
|
||||||
|
import time as _time
|
||||||
|
|
||||||
|
t0 = _time.time()
|
||||||
payload = self._payload(messages, tools, kw)
|
payload = self._payload(messages, tools, kw)
|
||||||
headers = {'Content-Type': 'application/json'}
|
headers = {'Content-Type': 'application/json'}
|
||||||
if self.api_key:
|
if self.api_key:
|
||||||
@ -126,7 +129,9 @@ class OpenAICompatible(ModelAdapter):
|
|||||||
for attempt in range(retries + 1):
|
for attempt in range(retries + 1):
|
||||||
try:
|
try:
|
||||||
data = await self._post(f'{self.api_base}/chat/completions', payload, headers)
|
data = await self._post(f'{self.api_base}/chat/completions', payload, headers)
|
||||||
return self._parse(data)
|
out = self._parse(data)
|
||||||
|
out.usage.latency_s = round(_time.time() - t0, 3)
|
||||||
|
return out
|
||||||
except Exception as e: # 5xx/429/timeouts: worth retrying
|
except Exception as e: # 5xx/429/timeouts: worth retrying
|
||||||
last_exc = e
|
last_exc = e
|
||||||
retryable = 'Server error' in str(e) or '504' in str(e) or '502' in str(e) \
|
retryable = 'Server error' in str(e) or '504' in str(e) or '502' in str(e) \
|
||||||
|
|||||||
@ -11,6 +11,7 @@ raw strings to the sync evaluate().
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import os
|
||||||
import time
|
import time
|
||||||
from typing import Any, Dict, List, Optional, Union
|
from typing import Any, Dict, List, Optional, Union
|
||||||
|
|
||||||
@ -36,6 +37,7 @@ async def generate_predictions(
|
|||||||
max_input_chars: int = 0,
|
max_input_chars: int = 0,
|
||||||
attach_context_keys: tuple = ('passage', 'context'),
|
attach_context_keys: tuple = ('passage', 'context'),
|
||||||
limit_per_task: Optional[int] = None,
|
limit_per_task: Optional[int] = None,
|
||||||
|
checkpoint: Union[bool, str] = False,
|
||||||
few_shot_num: int = 0,
|
few_shot_num: int = 0,
|
||||||
few_shot_samples: Optional[List[Sample]] = None,
|
few_shot_samples: Optional[List[Sample]] = None,
|
||||||
few_shot_text: Optional[str] = None,
|
few_shot_text: Optional[str] = None,
|
||||||
@ -104,15 +106,34 @@ async def generate_predictions(
|
|||||||
nonlocal done_count, total_usage
|
nonlocal done_count, total_usage
|
||||||
if env_factory is not None:
|
if env_factory is not None:
|
||||||
from ..agent import drive, trajectory_to_prediction
|
from ..agent import drive, trajectory_to_prediction
|
||||||
|
from ..agent.loop import Environment, Usage as _U # noqa: F401
|
||||||
|
|
||||||
async with sem:
|
async with sem:
|
||||||
traj = await drive(adapter, sample, env=env_factory(),
|
env = env_factory()
|
||||||
|
if type(env).run_task is not Environment.run_task:
|
||||||
|
# self-running env (official engine bundles: tau2/swe)
|
||||||
|
pred = await env.run_task(adapter, sample,
|
||||||
|
max_turns=max_turns, system=system)
|
||||||
|
if pred is None:
|
||||||
|
traj = await drive(adapter, sample, env=env,
|
||||||
max_turns=max_turns, system=system)
|
max_turns=max_turns, system=system)
|
||||||
total_usage = total_usage + traj.usage
|
|
||||||
pred = trajectory_to_prediction(traj)
|
pred = trajectory_to_prediction(traj)
|
||||||
pred['group_key'] = str(sample.metadata.get('test_category')
|
else:
|
||||||
|
traj = await drive(adapter, sample, env=env,
|
||||||
|
max_turns=max_turns, system=system)
|
||||||
|
pred = trajectory_to_prediction(traj)
|
||||||
|
if not pred.get('usage'):
|
||||||
|
pred['usage'] = traj.usage.model_dump() if 'traj' in dir() else {}
|
||||||
|
pred.setdefault('group_key', str(sample.metadata.get('test_category')
|
||||||
or sample.metadata.get('category')
|
or sample.metadata.get('category')
|
||||||
or sample.metadata.get('id') or sample.id or '')
|
or sample.metadata.get('domain')
|
||||||
|
or sample.metadata.get('id') or sample.id or ''))
|
||||||
|
u = pred.get('usage') or {}
|
||||||
|
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))
|
||||||
done_count += 1
|
done_count += 1
|
||||||
_progress(progress, done_count, len(samples), t0, total_usage)
|
_progress(progress, done_count, len(samples), t0, total_usage)
|
||||||
return pred
|
return pred
|
||||||
@ -143,7 +164,48 @@ async def generate_predictions(
|
|||||||
return {'raw': text, 'usage': out.usage.model_dump()}
|
return {'raw': text, 'usage': out.usage.model_dump()}
|
||||||
|
|
||||||
work = _apply_limits(samples, limit, limit_per_task)
|
work = _apply_limits(samples, limit, limit_per_task)
|
||||||
preds = list(await asyncio.gather(*(one(s) for s in work)))
|
# checkpointing: restore completed samples, generate only the rest
|
||||||
|
ckpt_store = None
|
||||||
|
if checkpoint:
|
||||||
|
from ..eval.checkpoint import CheckpointStore, checkpoint_path
|
||||||
|
|
||||||
|
ckpt = checkpoint if isinstance(checkpoint, str) else None
|
||||||
|
if ckpt is None:
|
||||||
|
from ..eval.checkpoint import checkpoint_path as _cp
|
||||||
|
|
||||||
|
ckpt = _cp('/data/evalharness' if False else os.path.expanduser('~/.cache/evalharness'),
|
||||||
|
getattr(spec, 'name', 'adhoc') if spec is not None else 'adhoc',
|
||||||
|
model_spec)
|
||||||
|
ckpt_store = CheckpointStore(ckpt, model=adapter.model or str(adapter))
|
||||||
|
restored = ckpt_store.load()
|
||||||
|
else:
|
||||||
|
restored = {}
|
||||||
|
|
||||||
|
keys = []
|
||||||
|
pending = []
|
||||||
|
preds_by_key: Dict[str, Dict[str, Any]] = {}
|
||||||
|
for i, s in enumerate(work):
|
||||||
|
k = CheckpointStore.key_for(s, i) if ckpt_store else str(i)
|
||||||
|
keys.append(k)
|
||||||
|
if k in restored:
|
||||||
|
preds_by_key[k] = restored[k]
|
||||||
|
else:
|
||||||
|
pending.append((i, s))
|
||||||
|
if ckpt_store is not None and restored:
|
||||||
|
print(f'checkpoint: restored {len(restored)} predictions '
|
||||||
|
f'({len(pending)} to generate) -> {ckpt_store.path}', flush=True)
|
||||||
|
|
||||||
|
async def run_one(i_s):
|
||||||
|
i, s = i_s
|
||||||
|
pred = await one(s)
|
||||||
|
if ckpt_store is not None:
|
||||||
|
ckpt_store.append(keys[i], pred)
|
||||||
|
return i, pred
|
||||||
|
|
||||||
|
fresh = await asyncio.gather(*(run_one((i, s)) for i, s in pending))
|
||||||
|
for i, pred in fresh:
|
||||||
|
preds_by_key[keys[i]] = pred
|
||||||
|
preds = [preds_by_key[k] for k in keys]
|
||||||
usages = [p.get('usage', {}) for p in preds]
|
usages = [p.get('usage', {}) for p in preds]
|
||||||
return preds, usages, total_usage
|
return preds, usages, total_usage
|
||||||
|
|
||||||
@ -191,6 +253,7 @@ async def run_eval(
|
|||||||
max_turns: int = 8,
|
max_turns: int = 8,
|
||||||
max_input_chars: int = 0,
|
max_input_chars: int = 0,
|
||||||
limit_per_task: Optional[int] = None,
|
limit_per_task: Optional[int] = None,
|
||||||
|
checkpoint: Union[bool, str] = False,
|
||||||
few_shot_num: int = -1,
|
few_shot_num: int = -1,
|
||||||
prompt_style: str = 'strict_letter',
|
prompt_style: str = 'strict_letter',
|
||||||
) -> EvalReport:
|
) -> EvalReport:
|
||||||
@ -229,10 +292,14 @@ async def run_eval(
|
|||||||
|
|
||||||
env_factory = None
|
env_factory = None
|
||||||
if env:
|
if env:
|
||||||
from ..agent.loop import ENV_REGISTRY, get_env
|
from ..agent import ENV_REGISTRY, get_env
|
||||||
|
|
||||||
if env not in ENV_REGISTRY:
|
if env not in ENV_REGISTRY:
|
||||||
raise KeyError(f'unknown env {env!r}; available: {", ".join(ENV_REGISTRY.names())}')
|
raise KeyError(f'unknown env {env!r}; available: {", ".join(ENV_REGISTRY.names())}')
|
||||||
|
probe = get_env(env)
|
||||||
|
if getattr(probe, 'needs_adapter', False):
|
||||||
|
env_factory = lambda: get_env(env, adapter=adapter) # noqa: E731
|
||||||
|
else:
|
||||||
env_factory = lambda: get_env(env) # noqa: E731
|
env_factory = lambda: get_env(env) # noqa: E731
|
||||||
|
|
||||||
# paper-faithful few-shot exemplars: official hand-written hooks first
|
# paper-faithful few-shot exemplars: official hand-written hooks first
|
||||||
@ -268,6 +335,7 @@ async def run_eval(
|
|||||||
gen_kwargs=gen_kwargs, env_factory=env_factory,
|
gen_kwargs=gen_kwargs, env_factory=env_factory,
|
||||||
system=system, max_turns=max_turns, max_input_chars=max_input_chars,
|
system=system, max_turns=max_turns, max_input_chars=max_input_chars,
|
||||||
limit_per_task=limit_per_task,
|
limit_per_task=limit_per_task,
|
||||||
|
checkpoint=checkpoint,
|
||||||
few_shot_num=few_shot_num,
|
few_shot_num=few_shot_num,
|
||||||
few_shot_samples=few_shot_samples, few_shot_text=few_shot_text,
|
few_shot_samples=few_shot_samples, few_shot_text=few_shot_text,
|
||||||
prompt_style=prompt_style)
|
prompt_style=prompt_style)
|
||||||
|
|||||||
@ -25,6 +25,18 @@ def text_table(target: Union[EvalReport, List[EvalReport]], opts: Dict) -> str:
|
|||||||
out.append('=' * max(len(head), 40))
|
out.append('=' * max(len(head), 40))
|
||||||
out.append(head)
|
out.append(head)
|
||||||
out.append('=' * max(len(head), 40))
|
out.append('=' * max(len(head), 40))
|
||||||
|
# run stats: duration, tokens, cost (from run_info + usage aggregates)
|
||||||
|
info = rep.metric_groups.get('run_info', {}) or {}
|
||||||
|
total_tokens = info.get('gen_total_tokens', 0)
|
||||||
|
tok_s = ''
|
||||||
|
if total_tokens:
|
||||||
|
tok_s = f' tokens={total_tokens}'
|
||||||
|
secs = 0.0
|
||||||
|
for s in rep.samples:
|
||||||
|
secs += float((s.usage or {}).get('latency_s', 0) or 0)
|
||||||
|
dur = f' time={secs / 3600:.2f}h' if secs >= 3600 else (f' time={secs:.0f}s' if secs else '')
|
||||||
|
if dur or tok_s:
|
||||||
|
out.append(f'n_samples={rep.num_samples}{dur}{tok_s}')
|
||||||
if rep.num_failed_extractions:
|
if rep.num_failed_extractions:
|
||||||
warn = (f'!! {rep.num_failed_extractions}/{rep.num_samples} extractions failed '
|
warn = (f'!! {rep.num_failed_extractions}/{rep.num_samples} extractions failed '
|
||||||
f'({_pct(rep.metrics.get("extraction_failure_rate", 0))}) -- check recipe/model fit')
|
f'({_pct(rep.metrics.get("extraction_failure_rate", 0))}) -- check recipe/model fit')
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user