76 lines
3.0 KiB
Python
76 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from ..models import RolloutTrace
|
|
from ..storage import atomic_write_jsonl
|
|
from .agentrm import AgentRM
|
|
from .pre_score import PreScorer, adjust_agent_score
|
|
from .identity import composite_id
|
|
|
|
|
|
class TraceScorer:
|
|
"""Resolve pre-scores and AgentRM scores into one effective score per trace."""
|
|
|
|
def __init__(
|
|
self,
|
|
pre_scorer: PreScorer,
|
|
agentrm: AgentRM | None = None,
|
|
):
|
|
self.pre_scorer = pre_scorer
|
|
self.agentrm = agentrm if agentrm is not None else AgentRM()
|
|
|
|
def score_all(
|
|
self,
|
|
traces: list[RolloutTrace],
|
|
final_score_dir: Path,
|
|
agentrm_dir: Path | None = None,
|
|
) -> dict[str, float]:
|
|
final_score_dir.mkdir(parents=True, exist_ok=True)
|
|
agentrm_dir = agentrm_dir or final_score_dir
|
|
agentrm_dir.mkdir(parents=True, exist_ok=True)
|
|
results = self.pre_scorer.score_all(traces)
|
|
pre_scores = {result.trace_id: result for result in results}
|
|
if set(pre_scores) != {trace.trace_id for trace in traces}:
|
|
raise ValueError("pre-scorer returned incomplete or duplicate results")
|
|
atomic_write_jsonl(
|
|
final_score_dir / "pre_scores.jsonl",
|
|
[pre_scores[trace.trace_id].to_dict() for trace in traces],
|
|
)
|
|
agentrm_traces = [
|
|
trace for trace in traces if pre_scores[trace.trace_id].route == "agentrm"
|
|
]
|
|
requests = [trace.agentrm_request() for trace in agentrm_traces]
|
|
local_scores = agentrm_dir / "agentrm_scores.jsonl"
|
|
found = self.agentrm.score_requests(requests)
|
|
atomic_write_jsonl(local_scores, found)
|
|
|
|
by_key = {composite_id(row): float(row["score"]) for row in found}
|
|
resolved = []
|
|
scores: dict[str, float] = {}
|
|
for trace in traces:
|
|
pre_score = pre_scores[trace.trace_id]
|
|
if pre_score.route == "fixed_score":
|
|
if pre_score.score is None:
|
|
raise ValueError(f"fixed pre-score is missing for {trace.trace_id}")
|
|
score = pre_score.score
|
|
source = "pre_score"
|
|
else:
|
|
score = adjust_agent_score(by_key[trace.key.as_tuple()], pre_score)
|
|
source = "agentrm"
|
|
scores[trace.trace_id] = score
|
|
resolved.append({
|
|
"task_name": trace.task_name,
|
|
"compile_type": trace.compile_type,
|
|
"test_name": trace.test_name,
|
|
"trace_id": trace.trace_id,
|
|
"effective_score": score,
|
|
"score_source": source,
|
|
"pre_score_reason": pre_score.reason,
|
|
"agentrm_score": by_key.get(trace.key.as_tuple()),
|
|
"efficiency_cost": pre_score.efficiency_cost,
|
|
"efficiency_penalty": pre_score.efficiency_penalty,
|
|
})
|
|
atomic_write_jsonl(final_score_dir / "effective_scores.jsonl", resolved)
|
|
return scores
|