| """Token shards (uint16 .bin files), train/val split and the corpus manifest.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from collections.abc import Iterable |
| from datetime import datetime, timezone |
| from pathlib import Path |
| from typing import Literal |
|
|
| import numpy as np |
|
|
|
|
| def split_of(record_id: str) -> Literal["train", "val"]: |
| return "val" if int(hashlib.sha256(record_id.encode()).hexdigest()[:8], 16) % 200 == 0 else "train" |
|
|
|
|
| def _sha256(path: Path) -> str: |
| h = hashlib.sha256() |
| with open(path, "rb") as f: |
| for chunk in iter(lambda: f.read(1 << 20), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| class ShardWriter: |
| """Appends `[bos] + ids + [eos]` documents; starts a new shard after `shard_tokens`.""" |
|
|
| def __init__(self, out_dir: Path, prefix: str, bos_id: int, eos_id: int, shard_tokens: int = 100_000_000): |
| self.out_dir = Path(out_dir) |
| self.out_dir.mkdir(parents=True, exist_ok=True) |
| self.prefix, self.bos_id, self.eos_id, self.shard_tokens = prefix, bos_id, eos_id, shard_tokens |
| self.tokens = 0 |
| self.bytes = 0 |
| self._buf: list[np.ndarray] = [] |
| self._buf_tokens = 0 |
| self._shards: list[dict] = [] |
|
|
| def add(self, ids: list[int], n_bytes: int = 0) -> None: |
| arr = np.asarray([self.bos_id, *ids, self.eos_id], dtype="<u2") |
| self._buf.append(arr) |
| self._buf_tokens += len(arr) |
| self.tokens += len(arr) |
| self.bytes += n_bytes |
| if self._buf_tokens >= self.shard_tokens: |
| self._flush() |
|
|
| def _flush(self) -> None: |
| if not self._buf: |
| return |
| path = self.out_dir / f"{self.prefix}_{len(self._shards):03d}.bin" |
| np.concatenate(self._buf).tofile(path) |
| self._shards.append({"path": str(path), "tokens": self._buf_tokens, "sha256": _sha256(path)}) |
| self._buf, self._buf_tokens = [], 0 |
|
|
| def close(self) -> list[dict]: |
| self._flush() |
| return self._shards |
|
|
|
|
| def write_shards(token_docs: Iterable[list[int]], out_dir: Path, prefix: str, bos_id: int, eos_id: int, |
| shard_tokens: int = 100_000_000) -> list[dict]: |
| w = ShardWriter(out_dir, prefix, bos_id, eos_id, shard_tokens) |
| for ids in token_docs: |
| w.add(ids) |
| return w.close() |
|
|
|
|
| def write_manifest(out_dir: Path, stats: dict) -> None: |
| manifest = {"created": datetime.now(timezone.utc).isoformat(timespec="seconds"), **stats} |
| (Path(out_dir) / "manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True)) |
|
|
