File size: 5,033 Bytes
88e15cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | from __future__ import annotations
import asyncio
from collections import Counter
from pathlib import Path
from rich.console import Console
from .config import WorkerPoolConfig
from .io import assigned_split, load_rewards
from .providers import WorkerRunner
from .rewards import compute_utility, deterministic_quality
from .schemas import RewardRecord, TaskRecord, WorkerResult
console = Console()
def _append_record(destination: Path, record: RewardRecord) -> None:
with destination.open("a", encoding="utf-8") as handle:
handle.write(record.model_dump_json() + "\n")
def _winner_label(record: RewardRecord) -> str:
best_reward = max(record.rewards)
winners = [
worker_id
for worker_id, reward in zip(record.worker_ids, record.rewards, strict=True)
if abs(reward - best_reward) <= 1e-12
]
if len(winners) == 1:
return winners[0]
return f"tie({','.join(winners)})"
async def _score_result(
runner: WorkerRunner,
task: TaskRecord,
result: WorkerResult,
config: WorkerPoolConfig,
) -> WorkerResult:
if task.grader.type == "llm_judge":
quality = await runner.judge(task, result.response)
else:
quality = deterministic_quality(task, result.response)
scored = result.model_copy(update={"quality": quality})
return scored.model_copy(
update={"utility": compute_utility(scored, config.reward)}
)
async def generate_reward_dataset(
tasks: list[TaskRecord],
config: WorkerPoolConfig,
output_path: str | Path,
limit: int | None = None,
resume: bool = False,
repetitions: int = 1,
) -> list[RewardRecord]:
if repetitions < 1:
raise ValueError("repetitions must be >= 1")
selected_tasks = tasks[:limit] if limit else tasks
destination = Path(output_path)
destination.parent.mkdir(parents=True, exist_ok=True)
records: list[RewardRecord] = []
if resume and destination.exists() and destination.stat().st_size:
records = load_rewards(destination)
if records[0].worker_ids != config.worker_ids:
raise ValueError(
"Cannot resume: existing reward file has a different worker order"
)
if any(record.repetitions != repetitions for record in records):
raise ValueError(
"Cannot resume: existing reward file uses a different repetition count"
)
elif not resume:
destination.write_text("", encoding="utf-8")
completed_ids = {record.task_id for record in records}
remaining_tasks = [
task for task in selected_tasks
if task.task_id not in completed_ids
]
runner = WorkerRunner(config)
try:
for index, task in enumerate(remaining_tasks, start=1):
# Worker-major ordering:
# qwen run1, run2, run3,
# deepseek run1, run2, run3,
# gemini run1, run2, run3
raw_results = await asyncio.gather(
*(
runner.complete(worker, task)
for worker in config.workers
for _ in range(repetitions)
)
)
scored_results = await asyncio.gather(
*(
_score_result(runner, task, result, config)
for result in raw_results
)
)
result_groups = [
scored_results[
worker_index * repetitions:
(worker_index + 1) * repetitions
]
for worker_index in range(len(config.workers))
]
averaged_rewards = [
sum(result.utility for result in group) / len(group)
for group in result_groups
]
record = RewardRecord(
task_id=task.task_id,
prompt=task.prompt,
domain=task.domain,
split=assigned_split(task.task_id, task.split),
tags=task.tags,
metadata=task.metadata,
repetitions=repetitions,
worker_ids=config.worker_ids,
rewards=averaged_rewards,
results=list(scored_results),
)
records.append(record)
await asyncio.to_thread(
_append_record,
destination,
record,
)
console.print(
f"[{index}/{len(remaining_tasks)}] "
f"{task.task_id}: "
f"best={_winner_label(record)} "
f"avg_rewards="
f"{[round(value, 4) for value in record.rewards]}"
)
finally:
await runner.close()
winners = Counter(
_winner_label(record)
for record in records
)
console.print(
f"Dataset now has {len(records)} records at {output_path}; "
f"oracle winners={dict(winners)}"
)
return records
|