SlayerLab/tokenizers

1#!/usr/bin/env python32"""Load and evaluate the non-Hugging-Face tokenizer JSONs in this dataset.3 4The adapters deliberately do not pretend that missing configuration is known.5``load_custom_tokenizer`` returns a usable byte-level BPE core, plus a fidelity6classification describing whether its intended pre-tokenization is reproducible7from the artifact alone.8"""9 10from __future__ import annotations11 12import json13from dataclasses import dataclass14from pathlib import Path15from typing import Callable16 17try:18    import regex19except ImportError:  # pragma: no cover - surfaced only for regex tokenizers20    regex = None21 22 23GPT2_PATTERN = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"24 25 26def bytes_to_unicode() -> dict[int, str]:27    """The reversible byte alphabet used by GPT-2/minBPE-style artifacts."""28    visible = list(range(ord("!"), ord("~") + 1))29    visible += list(range(ord("¡"), ord("¬") + 1))30    visible += list(range(ord("®"), ord("ÿ") + 1))31    chars = visible[:]32    extra = 033    for byte in range(256):34        if byte not in visible:35            visible.append(byte)36            chars.append(256 + extra)37            extra += 138    return dict(zip(visible, map(chr, chars)))39 40 41@dataclass42class AdaptedTokenizer:43    source_format: str44    fidelity: str45    fidelity_note: str46    vocab_size: int47    merge_count: int48    merge_ranks: dict[tuple[int, int], int]49    token_bytes: dict[int, bytes]50    pretokenizer: Callable[[str], list[str]] | None = None51 52    def _encode_bytes(self, data: bytes) -> list[int]:53        ids = list(data)54        while len(ids) >= 2:55            candidate = min(56                ((self.merge_ranks[pair], pair) for pair in zip(ids, ids[1:]) if pair in self.merge_ranks),57                default=None,58            )59            if candidate is None:60                break61            new_id, pair = candidate62            out: list[int] = []63            i = 064            while i < len(ids):65                if i + 1 < len(ids) and (ids[i], ids[i + 1]) == pair:66                    out.append(new_id)67                    i += 268                else:69                    out.append(ids[i])70                    i += 171            ids = out72        return ids73 74    def encode(self, text: str) -> list[int]:75        chunks = self.pretokenizer(text) if self.pretokenizer else [text]76        return [token for chunk in chunks for token in self._encode_bytes(chunk.encode("utf-8"))]77 78    def decode(self, ids: list[int]) -> str:79        return b"".join(self.token_bytes[token] for token in ids).decode("utf-8")80 81 82def _regex_split(pattern: str) -> Callable[[str], list[str]]:83    if regex is None:84        raise RuntimeError("The 'regex' package is required by this tokenizer")85    compiled = regex.compile(pattern)86    return lambda text: compiled.findall(text)87 88 89def _from_symbol_bpe(document: dict, path: Path) -> AdaptedTokenizer:90    model = document["model"]91    vocab: dict[str, int] = model["vocab"]92    byte_for_symbol = {symbol: byte for byte, symbol in bytes_to_unicode().items()}93    token_bytes: dict[int, bytes] = {}94    for symbol, token_id in vocab.items():95        token_bytes[token_id] = bytes(byte_for_symbol[c] for c in symbol)96 97    ranks = {}98    for line in model["merges"]:99        left, right = line.split(" ")100        ranks[(vocab[left], vocab[right])] = vocab[left + right]101 102    meta = document.get("meta", {})103    variant = str(meta.get("variant", ""))104    if "naive" in path.name:105        pretok, fidelity, note = None, "exact", "Artifact specifies no pre-tokenization."106    elif "slayer-v2" in path.name:107        # README specifies cl100k + full digit runs, but the exact cl100k108        # expression is not serialized. This is enough to inspect the BPE core,109        # not enough to claim benchmark parity.110        pretok, fidelity, note = None, "core_only", "Exact cl100k pre-tokenizer is absent from JSON; BPE core is lossless but intended boundaries are not reproducible from the artifact alone."111    elif isinstance(meta.get("regex_pretok"), str):112        pretok, fidelity, note = _regex_split(meta["regex_pretok"]), "exact", "Exact pre-tokenizer regex is serialized in artifact metadata."113    elif variant == "fast":114        pretok, fidelity, note = _regex_split(GPT2_PATTERN), "documented", "README identifies GPT-2 pre-tokenization, but the exact expression is not serialized."115    else:116        pretok, fidelity, note = None, "core_only", "Pre-tokenization is not fully specified."117    return AdaptedTokenizer("symbol_bpe", fidelity, note, len(vocab), len(ranks), ranks, token_bytes, pretok)118 119 120def _from_integer_bpe(document: dict) -> AdaptedTokenizer:121    merges = document.get("merges") or document.get("reguly_merge")122    vocab = document["vocab"]123    ranks = {(int(left), int(right)): int(new) for left, right, new in merges}124    # Merge triples are the authoritative lossless representation. Some early125    # Kasia artifacts rendered invalid standalone UTF-8 bytes as U+FFFD in126    # ``vocab``; reconstructing recursively avoids inheriting that display loss.127    token_bytes = {token_id: bytes([token_id]) for token_id in range(256)}128    for left, right, new in merges:129        token_bytes[int(new)] = token_bytes[int(left)] + token_bytes[int(right)]130 131    pattern = document.get("pretokenizer_regex")132    if pattern:133        pretok, fidelity, note = _regex_split(pattern), "exact", "Pre-tokenizer regex is serialized in the artifact."134    else:135        pretok, fidelity, note = None, "exact", "Artifact defines raw-stream byte BPE without pre-tokenization."136    return AdaptedTokenizer("integer_bpe", fidelity, note, len(vocab), len(ranks), ranks, token_bytes, pretok)137 138 139def _from_vocab_export(document: dict) -> AdaptedTokenizer:140    inverse = {symbol: byte for byte, symbol in bytes_to_unicode().items()}141    vocab: dict[str, int] = document["token_to_id"]142    token_bytes = {token_id: bytes(inverse[c] for c in symbol) for symbol, token_id in vocab.items()}143    ranks = {(int(left), int(right)): int(new) for left, right, new in document["merges"]}144    return AdaptedTokenizer(145        "vocab_export", "core_only",146        "Vocabulary and merge ranks are complete, but the intended Polish regex pre-tokenizer is documented only in the write-up, not serialized in JSON.",147        len(vocab), len(ranks), ranks, token_bytes, None,148    )149 150 151def load_custom_tokenizer(path: str | Path) -> AdaptedTokenizer:152    path = Path(path)153    document = json.loads(path.read_text(encoding="utf-8"))154    return load_custom_tokenizer_document(document, path)155 156 157def load_custom_tokenizer_document(document: dict, source_path: str | Path) -> AdaptedTokenizer:158    """Load an artifact already parsed from JSON, retaining its source filename hints."""159    path = Path(source_path)160    if isinstance(document.get("model"), dict) and isinstance(document["model"].get("merges"), list):161        return _from_symbol_bpe(document, path)162    if "token_to_id" in document and "merges" in document:163        return _from_vocab_export(document)164    if ("merges" in document or "reguly_merge" in document) and "vocab" in document:165        return _from_integer_bpe(document)166    raise ValueError(f"Unsupported custom tokenizer schema: {path}")167