File size: 876 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 | from __future__ import annotations
from pathlib import Path
from .io import write_json
from .model import RouterModel
from .schemas import RewardRecord
from .training_common import evaluate_loader, make_loader, resolve_device
def evaluate_checkpoint(
records: list[RewardRecord],
checkpoint: str | Path,
split: str = "test",
output_path: str | Path | None = None,
) -> dict:
model = RouterModel.from_checkpoint(checkpoint, device=resolve_device())
if model.worker_ids != records[0].worker_ids:
raise ValueError("Checkpoint worker order does not match reward dataset")
loader = make_loader(records, model, split, batch_size=32, shuffle=False)
report = evaluate_loader(model, loader)
report["split"] = split
report["checkpoint"] = str(checkpoint)
if output_path:
write_json(output_path, report)
return report
|