File size: 2,312 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 | from __future__ import annotations
import hashlib
import json
from collections.abc import Iterable
from pathlib import Path
from typing import TypeVar
from pydantic import BaseModel
from .schemas import RewardRecord, TaskRecord
T = TypeVar("T", bound=BaseModel)
def read_jsonl(path: str | Path, model: type[T]) -> list[T]:
records: list[T] = []
with Path(path).open("r", encoding="utf-8") as handle:
for line_number, line in enumerate(handle, start=1):
if not line.strip():
continue
try:
records.append(model.model_validate_json(line))
except Exception as exc:
raise ValueError(f"Invalid JSONL at {path}:{line_number}: {exc}") from exc
return records
def write_jsonl(path: str | Path, records: Iterable[BaseModel]) -> None:
destination = Path(path)
destination.parent.mkdir(parents=True, exist_ok=True)
temporary = destination.with_suffix(destination.suffix + ".tmp")
with temporary.open("w", encoding="utf-8") as handle:
for record in records:
handle.write(record.model_dump_json() + "\n")
temporary.replace(destination)
def load_tasks(path: str | Path) -> list[TaskRecord]:
return read_jsonl(path, TaskRecord)
def load_rewards(path: str | Path) -> list[RewardRecord]:
records = read_jsonl(path, RewardRecord)
if not records:
raise ValueError(f"No reward records found in {path}")
expected = records[0].worker_ids
for record in records[1:]:
if record.worker_ids != expected:
raise ValueError(
f"Worker order changed at task {record.task_id}; expected {expected}, "
f"got {record.worker_ids}"
)
return records
def assigned_split(task_id: str, explicit: str | None) -> str:
if explicit:
return explicit
bucket = int(hashlib.sha256(task_id.encode("utf-8")).hexdigest()[:8], 16) % 10
if bucket == 0:
return "test"
if bucket == 1:
return "validation"
return "train"
def write_json(path: str | Path, payload: object) -> None:
destination = Path(path)
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
|