| """Subword tokenizer wrapper for YAML-BERT v9. |
| |
| Wraps the HF `tokenizers` library so the rest of yaml_bert depends on |
| this stable interface, not on HF internals. |
| |
| Vocabulary semantics: |
| - Special tokens reserved at training time: [PAD], [UNK], [MASK], [LONG_VALUE] |
| - Otherwise byte-level BPE; any string can be encoded. |
| |
| Long-value rule (values only; keys never get this treatment): |
| - value length >= LONG_VALUE_THRESHOLD chars → single [LONG_VALUE] token |
| - MAX_CHARS_VALUE < value length < LONG_VALUE_THRESHOLD → truncate to MAX_CHARS_VALUE chars, then BPE |
| - shorter → BPE in full |
| """ |
| from __future__ import annotations |
|
|
| from tokenizers import Tokenizer |
|
|
| PAD_TOKEN = "[PAD]" |
| UNK_TOKEN = "[UNK]" |
| MASK_TOKEN = "[MASK]" |
| LONG_VALUE_TOKEN = "[LONG_VALUE]" |
|
|
| MAX_CHARS_VALUE = 64 |
| LONG_VALUE_THRESHOLD = 256 |
|
|
|
|
| class SubwordTokenizer: |
| """Wraps an HF Tokenizer with YAML-BERT-specific value-length rules.""" |
|
|
| def __init__(self, hf_tokenizer: Tokenizer) -> None: |
| self._tok = hf_tokenizer |
| self.pad_id = hf_tokenizer.token_to_id(PAD_TOKEN) |
| self.unk_id = hf_tokenizer.token_to_id(UNK_TOKEN) |
| self.mask_id = hf_tokenizer.token_to_id(MASK_TOKEN) |
| self.long_value_id = hf_tokenizer.token_to_id(LONG_VALUE_TOKEN) |
| for name, val in ( |
| (PAD_TOKEN, self.pad_id), (UNK_TOKEN, self.unk_id), |
| (MASK_TOKEN, self.mask_id), (LONG_VALUE_TOKEN, self.long_value_id), |
| ): |
| if val is None: |
| raise ValueError( |
| f"SubwordTokenizer: required special token {name!r} not " |
| f"in tokenizer vocab — was the tokenizer trained with " |
| f"this special token reserved?" |
| ) |
|
|
| @classmethod |
| def load(cls, path: str) -> "SubwordTokenizer": |
| return cls(Tokenizer.from_file(path)) |
|
|
| @property |
| def vocab_size(self) -> int: |
| return self._tok.get_vocab_size() |
|
|
| def encode_token(self, token: str, *, is_value: bool) -> list[int]: |
| """Encode one linearizer-node token to a list of subword ids. |
| |
| See module docstring for the value-length rule. Keys are always |
| BPE-encoded in full regardless of length. |
| """ |
| if is_value: |
| if len(token) >= LONG_VALUE_THRESHOLD: |
| return [self.long_value_id] |
| if len(token) > MAX_CHARS_VALUE: |
| token = token[:MAX_CHARS_VALUE] |
| return self._tok.encode(token).ids |
|
|
| def decode(self, ids: list[int]) -> str: |
| return self._tok.decode(ids) |
|
|
