| """Tokenizer wrapper with strict identity tracking. |
| |
| Engram hashes token ids, so a tokenizer swap silently rehashes every table row |
| (docs/ENGRAM_NOTES.md section 1.2). The checkpoint therefore stores this |
| tokenizer's `identity()`, and `checkpoint.load()` refuses to resume when it |
| differs. This is the "tokenizer frozen once a run starts" rule from the brief, |
| enforced rather than documented. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| __all__ = ["Tokenizer", "train_tokenizer", "SPECIAL_TOKENS"] |
|
|
| |
| SPECIAL_TOKENS = ["<pad>", "<bos>", "<eos>", "<unk>"] |
| PAD_ID, BOS_ID, EOS_ID, UNK_ID = 0, 1, 2, 3 |
|
|
|
|
| class Tokenizer: |
| """Thin wrapper over a HuggingFace `tokenizers.Tokenizer`.""" |
|
|
| def __init__(self, path: str | Path): |
| from tokenizers import Tokenizer as HFTokenizer |
|
|
| self.path = Path(path) |
| tok_file = self.path / "tokenizer.json" if self.path.is_dir() else self.path |
| if not tok_file.exists(): |
| raise FileNotFoundError( |
| f"no tokenizer.json at {tok_file}. Train one with:\n" |
| f" python scripts/train_tokenizer.py --input <corpus> --out {self.path}" |
| ) |
| self._tok = HFTokenizer.from_file(str(tok_file)) |
| self._blob = tok_file.read_bytes() |
| self.pad_id, self.bos_id, self.eos_id, self.unk_id = PAD_ID, BOS_ID, EOS_ID, UNK_ID |
| |
| |
| |
| |
| for token, expected in zip(SPECIAL_TOKENS, (PAD_ID, BOS_ID, EOS_ID, UNK_ID)): |
| actual = self._tok.token_to_id(token) |
| if actual is not None and actual != expected: |
| raise ValueError( |
| f"{tok_file} maps {token!r} to id {actual}, but this " |
| f"codebase hardcodes id {expected} (packed shards and " |
| f"Engram hashes depend on it)" |
| ) |
|
|
| |
|
|
| def __len__(self) -> int: |
| return self._tok.get_vocab_size() |
|
|
| @property |
| def vocab_size(self) -> int: |
| return self._tok.get_vocab_size() |
|
|
| @property |
| def backend_tokenizer(self): |
| return self._tok |
|
|
| def encode(self, text: str, add_special: bool = False) -> list[int]: |
| ids = self._tok.encode(text, add_special_tokens=False).ids |
| if add_special: |
| ids = [self.bos_id, *ids, self.eos_id] |
| return ids |
|
|
| def decode(self, ids, skip_special_tokens: bool = True) -> str: |
| return self._tok.decode(list(ids), skip_special_tokens=skip_special_tokens) |
|
|
| def id_to_token(self, i: int): |
| return self._tok.id_to_token(i) |
|
|
| |
|
|
| def identity(self) -> dict: |
| """Everything needed to prove two runs used the same tokenizer.""" |
| return { |
| "sha256": hashlib.sha256(self._blob).hexdigest(), |
| "vocab_size": self.vocab_size, |
| "pad_id": self.pad_id, |
| "bos_id": self.bos_id, |
| "eos_id": self.eos_id, |
| } |
|
|
| def save_into(self, directory: str | Path) -> None: |
| """Copy the tokenizer into a checkpoint directory, so the checkpoint is |
| self-contained and a resume cannot pick up a drifted file.""" |
| directory = Path(directory) |
| directory.mkdir(parents=True, exist_ok=True) |
| (directory / "tokenizer.json").write_bytes(self._blob) |
| (directory / "tokenizer_identity.json").write_text( |
| json.dumps(self.identity(), indent=2), encoding="utf-8" |
| ) |
|
|
|
|
| def train_tokenizer( |
| files: list[str | Path], |
| out_dir: str | Path, |
| vocab_size: int = 32768, |
| min_frequency: int = 2, |
| ) -> Tokenizer: |
| """Train a byte-level BPE and write `tokenizer.json`. |
| |
| Byte-level so the vocabulary can never produce <unk> on unseen text, which |
| keeps Engram's hashing total. |
| """ |
| from tokenizers import Tokenizer as HFTokenizer |
| from tokenizers import decoders, models, pre_tokenizers, trainers |
|
|
| out_dir = Path(out_dir) |
| out_dir.mkdir(parents=True, exist_ok=True) |
|
|
| tok = HFTokenizer(models.BPE(unk_token=None)) |
| tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) |
| tok.decoder = decoders.ByteLevel() |
| trainer = trainers.BpeTrainer( |
| vocab_size=vocab_size, |
| min_frequency=min_frequency, |
| special_tokens=SPECIAL_TOKENS, |
| initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), |
| show_progress=True, |
| ) |
| tok.train([str(f) for f in files], trainer) |
| tok.save(str(out_dir / "tokenizer.json")) |
|
|
| wrapper = Tokenizer(out_dir) |
| wrapper.save_into(out_dir) |
| return wrapper |
|
|
