| """Packed token corpora for STRATA language-model pretraining. |
| |
| Documents are tokenised, separated by an end-of-sequence token, and packed into |
| a single flat token array on disk (``tokens.bin``) with a JSON sidecar |
| (``meta.json``). Training reads fixed-length, non-overlapping windows via |
| :class:`PackedLMDataset` and a memory map, so multi-billion-token corpora never |
| need to fit in RAM. |
| |
| Tokenisation is passed in as a callable returning token ids, which keeps this |
| module independent of any specific tokenizer and lets the byte tokenizer use a |
| fast vectorised path instead of building span lattices per byte. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from collections import Counter |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Callable, Iterable |
|
|
| import numpy as np |
|
|
| TokenizeIds = Callable[[str], np.ndarray] |
|
|
| TOKENS_FILENAME = "tokens.bin" |
| META_FILENAME = "meta.json" |
|
|
|
|
| def choose_token_dtype(vocab_size: int) -> np.dtype: |
| """Smallest unsigned integer dtype that can hold every id in the vocab.""" |
|
|
| if vocab_size <= 0: |
| raise ValueError("vocab_size must be positive") |
| if vocab_size <= np.iinfo(np.uint16).max + 1: |
| return np.dtype(np.uint16) |
| if vocab_size <= np.iinfo(np.uint32).max + 1: |
| return np.dtype(np.uint32) |
| raise ValueError(f"vocab_size {vocab_size} too large for uint32 token ids") |
|
|
|
|
| @dataclass(slots=True) |
| class PackedCorpusMeta: |
| """Sidecar metadata describing a packed token corpus.""" |
|
|
| num_tokens: int |
| dtype: str |
| vocab_size: int |
| eos_id: int |
| tokenizer_name: str |
| documents: int |
| per_language_tokens: dict[str, int] |
| sha256: str |
| seed: int | None = None |
| source: dict[str, object] = field(default_factory=dict) |
|
|
| def to_dict(self) -> dict[str, object]: |
| return { |
| "num_tokens": self.num_tokens, |
| "dtype": self.dtype, |
| "vocab_size": self.vocab_size, |
| "eos_id": self.eos_id, |
| "tokenizer_name": self.tokenizer_name, |
| "documents": self.documents, |
| "per_language_tokens": self.per_language_tokens, |
| "sha256": self.sha256, |
| "seed": self.seed, |
| "source": self.source, |
| } |
|
|
| @classmethod |
| def from_dict(cls, data: dict[str, object]) -> "PackedCorpusMeta": |
| return cls( |
| num_tokens=int(data["num_tokens"]), |
| dtype=str(data["dtype"]), |
| vocab_size=int(data["vocab_size"]), |
| eos_id=int(data["eos_id"]), |
| tokenizer_name=str(data.get("tokenizer_name", "")), |
| documents=int(data.get("documents", 0)), |
| per_language_tokens=dict(data.get("per_language_tokens", {}) or {}), |
| sha256=str(data.get("sha256", "")), |
| seed=data.get("seed"), |
| source=dict(data.get("source", {}) or {}), |
| ) |
|
|
| def write(self, out_dir: Path) -> Path: |
| path = out_dir / META_FILENAME |
| path.write_text(json.dumps(self.to_dict(), indent=2, ensure_ascii=False) + "\n", encoding="utf-8") |
| return path |
|
|
| @classmethod |
| def read(cls, out_dir: Path) -> "PackedCorpusMeta": |
| return cls.from_dict(json.loads((out_dir / META_FILENAME).read_text(encoding="utf-8"))) |
|
|
|
|
| def build_packed_corpus( |
| *, |
| documents: Iterable[tuple[str, str]], |
| tokenize_ids: TokenizeIds, |
| eos_id: int, |
| vocab_size: int, |
| tokenizer_name: str, |
| out_dir: Path, |
| max_tokens: int | None = None, |
| per_language_tokens: dict[str, int] | None = None, |
| seed: int | None = None, |
| source: dict[str, object] | None = None, |
| flush_every_tokens: int = 8_000_000, |
| logger=None, |
| ) -> PackedCorpusMeta: |
| """Tokenise ``(language, text)`` documents and pack them into ``tokens.bin``. |
| |
| ``per_language_tokens`` caps how many tokens each language contributes (for |
| balanced multilingual mixtures); ``max_tokens`` caps the total. Both are |
| honoured mid-document at token granularity. |
| """ |
|
|
| out_dir.mkdir(parents=True, exist_ok=True) |
| dtype = choose_token_dtype(vocab_size) |
| if not (0 <= eos_id < vocab_size): |
| raise ValueError(f"eos_id {eos_id} out of range for vocab_size {vocab_size}") |
|
|
| tokens_path = out_dir / TOKENS_FILENAME |
| caps = dict(per_language_tokens or {}) |
| lang_counts: Counter[str] = Counter() |
| hasher = hashlib.sha256() |
| total = 0 |
| docs = 0 |
| buffer: list[np.ndarray] = [] |
| buffered = 0 |
|
|
| def flush() -> None: |
| nonlocal buffered |
| if not buffer: |
| return |
| chunk = np.concatenate(buffer) |
| chunk.tofile(handle) |
| hasher.update(chunk.tobytes()) |
| buffer.clear() |
| buffered = 0 |
|
|
| with tokens_path.open("wb") as handle: |
| for language, text in documents: |
| if max_tokens is not None and total >= max_tokens: |
| break |
| if caps and all(lang_counts[lang] >= cap for lang, cap in caps.items()): |
| break |
| if language in caps and lang_counts[language] >= caps[language]: |
| continue |
| if not text: |
| continue |
| ids = np.asarray(tokenize_ids(text), dtype=dtype) |
| if ids.size == 0: |
| continue |
| ids = np.append(ids, np.asarray([eos_id], dtype=dtype)) |
|
|
| |
| if language in caps: |
| remaining = caps[language] - lang_counts[language] |
| if remaining <= 0: |
| continue |
| if ids.size > remaining: |
| ids = ids[:remaining] |
| if max_tokens is not None and total + ids.size > max_tokens: |
| ids = ids[: max_tokens - total] |
| if ids.size == 0: |
| continue |
| if ids.max() >= vocab_size: |
| raise ValueError( |
| f"token id {int(ids.max())} >= vocab_size {vocab_size}; " |
| "tokenizer/model vocab mismatch" |
| ) |
|
|
| buffer.append(ids) |
| buffered += ids.size |
| total += ids.size |
| lang_counts[language] += int(ids.size) |
| docs += 1 |
| if buffered >= flush_every_tokens: |
| flush() |
| if logger is not None: |
| logger.info("packed %d tokens (%d docs)", total, docs) |
| flush() |
|
|
| meta = PackedCorpusMeta( |
| num_tokens=total, |
| dtype=dtype.name, |
| vocab_size=vocab_size, |
| eos_id=eos_id, |
| tokenizer_name=tokenizer_name, |
| documents=docs, |
| per_language_tokens=dict(lang_counts), |
| sha256=hasher.hexdigest(), |
| seed=seed, |
| source=dict(source or {}), |
| ) |
| meta.write(out_dir) |
| if logger is not None: |
| logger.info("packed corpus: %d tokens, %d docs -> %s", total, docs, tokens_path) |
| return meta |
|
|
|
|
| class PackedLMDataset: |
| """Fixed-length, non-overlapping windows over a packed token corpus. |
| |
| Implements the ``torch.utils.data.Dataset`` protocol (``__len__`` / |
| ``__getitem__``) without importing torch at module load; each item is a 1-D |
| ``int64`` tensor of length ``seq_len`` suitable for causal-LM training where |
| the model shifts inputs and labels internally. |
| """ |
|
|
| def __init__(self, corpus_dir: str | Path, *, seq_len: int) -> None: |
| if seq_len <= 1: |
| raise ValueError("seq_len must be > 1") |
| self.corpus_dir = Path(corpus_dir) |
| self.meta = PackedCorpusMeta.read(self.corpus_dir) |
| self.seq_len = seq_len |
| self.dtype = np.dtype(self.meta.dtype) |
| tokens_path = self.corpus_dir / TOKENS_FILENAME |
| if not tokens_path.exists(): |
| raise FileNotFoundError(f"packed tokens not found: {tokens_path}") |
| self._tokens = np.memmap(tokens_path, dtype=self.dtype, mode="r") |
| self.num_windows = (len(self._tokens)) // seq_len |
| if self.num_windows == 0: |
| raise ValueError( |
| f"corpus has {len(self._tokens)} tokens, too few for seq_len {seq_len}" |
| ) |
|
|
| def __len__(self) -> int: |
| return self.num_windows |
|
|
| def __getitem__(self, index: int): |
| import torch |
|
|
| if index < 0: |
| index += self.num_windows |
| if not 0 <= index < self.num_windows: |
| raise IndexError(index) |
| start = index * self.seq_len |
| window = np.asarray(self._tokens[start : start + self.seq_len], dtype=np.int64) |
| return torch.from_numpy(window) |
|
|
| def total_tokens_used(self) -> int: |
| return self.num_windows * self.seq_len |
|
|
