| """Export a ZENITH BPE tokenizer to the Hugging Face `tokenizers` format. |
| |
| The fast encoder only executes ZENITH's own vocabulary and merges; parity with |
| `BPETokenizer.encode` is enforced by tests. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| from tokenizers import AddedToken, Regex, Tokenizer, decoders, models, pre_tokenizers |
|
|
| from zenith.tokenizer.bpe import PRETOKENIZE_PATTERN, SPECIAL_TOKENS, BPETokenizer |
|
|
|
|
| def _bytes_to_unicode() -> dict[int, str]: |
| """GPT-2's reversible byte -> printable character map, used by HF ByteLevel.""" |
| bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1)) |
| cs = bs[:] |
| n = 0 |
| for b in range(256): |
| if b not in bs: |
| bs.append(b) |
| cs.append(256 + n) |
| n += 1 |
| return {b: chr(c) for b, c in zip(bs, cs)} |
|
|
|
|
| _BYTE_CHARS = _bytes_to_unicode() |
|
|
|
|
| def _as_unicode(b: bytes) -> str: |
| return "".join(_BYTE_CHARS[x] for x in b) |
|
|
|
|
| def export_hf(tok: BPETokenizer, path: Path) -> None: |
| vocab: dict[str, int] = {} |
| for i in range(256 + len(tok.merges)): |
| s = _as_unicode(tok.token_bytes(i)) |
| if s in vocab: |
| raise ValueError(f"token {i} duplicates the bytes of token {vocab[s]}; HF BPE cannot represent this") |
| vocab[s] = i |
| for s, i in tok.special_ids.items(): |
| vocab[s] = i |
| merges = [(_as_unicode(tok.token_bytes(a)), _as_unicode(tok.token_bytes(b))) for a, b in tok.merges] |
|
|
| fast = Tokenizer(models.BPE(vocab=vocab, merges=merges)) |
| fast.pre_tokenizer = pre_tokenizers.Sequence( |
| [ |
| pre_tokenizers.Split(Regex(PRETOKENIZE_PATTERN), behavior="isolated"), |
| pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False), |
| ] |
| ) |
| fast.decoder = decoders.ByteLevel() |
| fast.add_special_tokens([AddedToken(s, special=True, normalized=False) for s in SPECIAL_TOKENS]) |
| fast.save(str(path)) |
|
|
|
|
| def load_fast(path: Path) -> Tokenizer: |
| fast = Tokenizer.from_file(str(path)) |
| |
| fast.encode_special_tokens = True |
| return fast |
|
|
|
|
| def encode_fast(fast: Tokenizer, texts: list[str]) -> list[list[int]]: |
| return [e.ids for e in fast.encode_batch(texts, add_special_tokens=False)] |
|
|
