File size: 3,108 Bytes
026f63e | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """Text tokenizers behind one interface (numpy only, no torch):
tok.encode_batch(texts, max_len, dynamic=False) -> (n, L) int64, padded with tok.pad_id
tok.encode_choice_sets(choice_sets, max_len) -> (B, N, L); absent choices are all-pad rows
`ByteTokenizer` is the original MiniJEV byte-level tokenizer (vocab 260, pad 0).
`HFTokenizer` wraps a pretrained encoder's `tokenizer.json` through the Rust
`tokenizers` package (no transformers/torch needed at inference).
"""
from __future__ import annotations
from pathlib import Path
from typing import Sequence
import numpy as np
from minijev import tokenizer as byte_tok
class ByteTokenizer:
kind = "byte"
pad_id = byte_tok.PAD
vocab_size = byte_tok.VOCAB_SIZE
def encode_batch(self, texts: Sequence[str], max_len: int, dynamic: bool = False) -> np.ndarray:
out = byte_tok.encode_batch(texts, max_len)
if dynamic:
n = int((out != self.pad_id).sum(1).max()) if len(out) else 1
out = out[:, : max(1, n)]
return out
def encode_choice_sets(self, choice_sets, max_len):
return byte_tok.encode_choice_sets(choice_sets, max_len)
class HFTokenizer:
kind = "hf"
def __init__(self, name_or_path: str):
from tokenizers import Tokenizer
p = Path(name_or_path)
if p.is_dir():
p = p / "tokenizer.json"
if not p.exists():
from huggingface_hub import hf_hub_download
p = Path(hf_hub_download(name_or_path, "tokenizer.json"))
self.path = str(p)
self.tok = Tokenizer.from_file(str(p))
self.tok.no_padding()
self.tok.no_truncation()
pad = self.tok.token_to_id("<pad>")
self.pad_id = int(pad if pad is not None else (self.tok.token_to_id("[PAD]") or 0))
self.vocab_size = self.tok.get_vocab_size()
def encode_batch(self, texts: Sequence[str], max_len: int, dynamic: bool = False) -> np.ndarray:
encs = self.tok.encode_batch(list(texts))
L = max_len
if dynamic:
L = max(1, min(max_len, max((len(e.ids) for e in encs), default=1)))
out = np.full((len(texts), L), self.pad_id, np.int64)
for i, e in enumerate(encs):
ids = e.ids
if len(ids) > L: # keep the final special token (</s> / [SEP])
ids = ids[: L - 1] + ids[-1:]
out[i, : len(ids)] = ids
return out
def encode_choice_sets(self, choice_sets, max_len):
n = max(len(c) for c in choice_sets)
flat = [c for cs in choice_sets for c in cs]
enc = self.encode_batch(flat, max_len, dynamic=True)
out = np.full((len(choice_sets), n, enc.shape[1]), self.pad_id, np.int64)
k = 0
for i, cs in enumerate(choice_sets):
out[i, : len(cs)] = enc[k:k + len(cs)]
k += len(cs)
return out
def make_tokenizer(text_cfg: dict | None):
text_cfg = text_cfg or {}
if text_cfg.get("type", "byte") == "byte":
return ByteTokenizer()
return HFTokenizer(text_cfg.get("tokenizer") or text_cfg["name"])
|