| """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: |
| 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"]) |
|
|
