File size: 3,889 Bytes
58223a8 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | from __future__ import annotations
import json
from pathlib import Path
from transformers import PreTrainedTokenizer
CANONICAL_VOCAB = {"<BOS>": 0, "+": 1, "=": 2, **{str(digit): digit + 3 for digit in range(10)}}
class AdditionTokenizer(PreTrainedTokenizer):
vocab_files_names = {"vocab_file": "vocab.json"}
model_input_names = ["input_ids", "attention_mask"]
def __init__(self, vocab_file: str | None = None, **kwargs) -> None:
if vocab_file is None:
vocab = dict(CANONICAL_VOCAB)
else:
with Path(vocab_file).open("r", encoding="utf-8") as handle:
vocab = json.load(handle)
if vocab != CANONICAL_VOCAB:
raise ValueError("AdditionTokenizer requires the canonical 13-token vocabulary.")
self._vocab = vocab
self._ids_to_tokens = {token_id: token for token, token_id in vocab.items()}
kwargs.pop("bos_token", None)
kwargs.pop("eos_token", None)
kwargs.pop("pad_token", None)
kwargs.pop("unk_token", None)
super().__init__(
bos_token="<BOS>",
eos_token=None,
pad_token=None,
unk_token=None,
**kwargs,
)
@property
def vocab_size(self) -> int:
return len(self._vocab)
def get_vocab(self) -> dict[str, int]:
return dict(self._vocab)
def _tokenize(self, text: str, **kwargs) -> list[str]:
compact = "".join(text.split())
invalid = sorted(set(compact) - set("0123456789+="))
if invalid:
raise ValueError(f"Unsupported characters for addition tokenizer: {''.join(invalid)}")
return list(compact)
def _convert_token_to_id(self, token: str) -> int:
try:
return self._vocab[token]
except KeyError as exc:
raise ValueError(f"Unknown addition token: {token!r}") from exc
def _convert_id_to_token(self, index: int) -> str:
try:
return self._ids_to_tokens[index]
except KeyError as exc:
raise ValueError(f"Unknown addition token ID: {index}") from exc
def convert_tokens_to_string(self, tokens: list[str]) -> str:
return "".join(tokens)
def build_inputs_with_special_tokens(
self,
token_ids_0: list[int],
token_ids_1: list[int] | None = None,
) -> list[int]:
if token_ids_1 is not None:
raise ValueError("AdditionTokenizer does not support sequence pairs.")
return [self.bos_token_id, *token_ids_0]
def get_special_tokens_mask(
self,
token_ids_0: list[int],
token_ids_1: list[int] | None = None,
already_has_special_tokens: bool = False,
) -> list[int]:
if already_has_special_tokens:
return [int(token_id == self.bos_token_id) for token_id in token_ids_0]
if token_ids_1 is not None:
raise ValueError("AdditionTokenizer does not support sequence pairs.")
return [1, *([0] * len(token_ids_0))]
def create_token_type_ids_from_sequences(
self,
token_ids_0: list[int],
token_ids_1: list[int] | None = None,
) -> list[int]:
if token_ids_1 is not None:
raise ValueError("AdditionTokenizer does not support sequence pairs.")
return [0] * (len(token_ids_0) + 1)
def save_vocabulary(
self,
save_directory: str,
filename_prefix: str | None = None,
) -> tuple[str]:
directory = Path(save_directory)
directory.mkdir(parents=True, exist_ok=True)
filename = f"{filename_prefix + '-' if filename_prefix else ''}vocab.json"
path = directory / filename
path.write_text(json.dumps(self._vocab, indent=2, sort_keys=True) + "\n", encoding="utf-8")
return (str(path),)
AdditionTokenizer.register_for_auto_class("AutoTokenizer")
|