File size: 2,337 Bytes
0cf8154 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """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))
# Documents that contain special-token text must encode as ordinary bytes.
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)]
|