Buckets:
| """Manifest-indexed packed trajectories and deterministic query scheduling.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| from collections import OrderedDict | |
| from pathlib import Path | |
| import torch | |
| class PackedTrajectoryCache: | |
| """Read a prepared trajectory cache through its stable manifest indices.""" | |
| def __init__(self, root: str | Path, *, max_open_shards: int = 2): | |
| self.root = Path(root) | |
| if max_open_shards <= 0: | |
| raise ValueError("max_open_shards must be positive") | |
| if not (self.root / "READY").is_file(): | |
| raise ValueError(f"cache is incomplete: {self.root}") | |
| self.rows = tuple( | |
| json.loads(line) | |
| for line in (self.root / "manifest.jsonl").read_text().splitlines() | |
| if line | |
| ) | |
| sample_ids = [row["sample_id"] for row in self.rows] | |
| if len(sample_ids) != len(set(sample_ids)): | |
| raise ValueError("cache manifest contains duplicate sample IDs") | |
| self.max_open_shards = max_open_shards | |
| self._shards: OrderedDict[str, list[dict]] = OrderedDict() | |
| def indices(self, split: str) -> tuple[int, ...]: | |
| result = tuple( | |
| index for index, row in enumerate(self.rows) if row["split"] == split | |
| ) | |
| if not result: | |
| raise ValueError(f"cache split is empty: {split}") | |
| return result | |
| def sample_id(self, index: int) -> str: | |
| return str(self.rows[index]["sample_id"]) | |
| def _items(self, relative: str) -> list[dict]: | |
| if relative in self._shards: | |
| self._shards.move_to_end(relative) | |
| return self._shards[relative] | |
| payload = torch.load( | |
| self.root / relative, | |
| map_location="cpu", | |
| weights_only=False, | |
| mmap=True, | |
| ) | |
| if payload.get("schema") != "zebra_uniar_cache_v1": | |
| raise ValueError(f"unsupported cache shard schema: {relative}") | |
| self._shards[relative] = payload["items"] | |
| while len(self._shards) > self.max_open_shards: | |
| self._shards.popitem(last=False) | |
| return self._shards[relative] | |
| def load(self, index: int) -> dict: | |
| row = self.rows[index] | |
| item = self._items(row["cache_shard"])[row["cache_index"]] | |
| if item["sample_id"] != row["sample_id"]: | |
| raise ValueError("cache item and manifest sample IDs differ") | |
| return item | |
| def answer_reference(item: dict) -> dict[str, str]: | |
| """Extract the single canonical answer from a packed trajectory item.""" | |
| answers = [ | |
| event.get("text") | |
| for event in item["trajectory"]["assistant"] | |
| if event.get("role") == "answer" | |
| ] | |
| if len(answers) != 1 or not isinstance(answers[0], str) or not answers[0].strip(): | |
| raise ValueError("trajectory must contain exactly one non-empty answer") | |
| return {"answer": answers[0], "schema": "answer_v1"} | |
| def _epoch_order(indices: tuple[int, ...], *, seed: int, epoch: int) -> tuple[int, ...]: | |
| return tuple( | |
| sorted( | |
| indices, | |
| key=lambda index: ( | |
| hashlib.sha256(f"{seed}\0{epoch}\0{index}".encode()).digest(), | |
| index, | |
| ), | |
| ) | |
| ) | |
| def query_indices_for_update( | |
| train_indices: tuple[int, ...], | |
| *, | |
| batch_size: int, | |
| seed: int, | |
| update: int, | |
| shuffle: bool = True, | |
| ) -> tuple[int, ...]: | |
| """Return one exact logical query budget from deterministic epoch orders.""" | |
| if not train_indices or len(set(train_indices)) != len(train_indices): | |
| raise ValueError("train_indices must be non-empty and unique") | |
| if batch_size <= 0 or update < 0: | |
| raise ValueError("batch_size must be positive and update non-negative") | |
| result = [] | |
| absolute = update * batch_size | |
| while len(result) < batch_size: | |
| epoch, offset = divmod(absolute, len(train_indices)) | |
| order = ( | |
| _epoch_order(train_indices, seed=seed, epoch=epoch) | |
| if shuffle | |
| else train_indices | |
| ) | |
| take = min(batch_size - len(result), len(order) - offset) | |
| result.extend(order[offset : offset + take]) | |
| absolute += take | |
| return tuple(result) | |
| # Retain the established OPD name while all modules share one implementation. | |
| rollout_indices_for_update = query_indices_for_update | |
| __all__ = [ | |
| "answer_reference", | |
| "PackedTrajectoryCache", | |
| "query_indices_for_update", | |
| "rollout_indices_for_update", | |
| ] | |
Xet Storage Details
- Size:
- 4.5 kB
- Xet hash:
- eb0d6a89ed736821aa1432a44777f3c5d72b5cc134a43ffc179762f1d23b4d89
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.