| from __future__ import annotations |
|
|
| from collections.abc import Sequence |
|
|
| import torch |
| from torch.utils.data import Dataset |
|
|
| from .io import assigned_split |
| from .schemas import RewardRecord, route_text |
|
|
|
|
| class RouterDataset(Dataset): |
| def __init__(self, records: Sequence[RewardRecord], split: str | None = None): |
| self.records = [ |
| record |
| for record in records |
| if split is None or assigned_split(record.task_id, record.split) == split |
| ] |
|
|
| def __len__(self) -> int: |
| return len(self.records) |
|
|
| def __getitem__(self, index: int) -> dict: |
| record = self.records[index] |
| return { |
| "task_id": record.task_id, |
| "text": route_text(record.prompt, record.domain, record.tags), |
| "rewards": torch.tensor(record.rewards, dtype=torch.float32), |
| } |
|
|
|
|
| class RouterCollator: |
| def __init__(self, tokenizer, max_length: int): |
| self.tokenizer = tokenizer |
| self.max_length = max_length |
|
|
| def __call__(self, examples: list[dict]) -> dict: |
| encoded = self.tokenizer( |
| [example["text"] for example in examples], |
| padding=True, |
| truncation=True, |
| max_length=self.max_length, |
| return_tensors="pt", |
| ) |
| encoded["rewards"] = torch.stack([example["rewards"] for example in examples]) |
| encoded["task_ids"] = [example["task_id"] for example in examples] |
| return encoded |
|
|
|
|
