| 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") |
|
|
|
|
