vmal/3-digit-basic-calc

1# Copyright 2026. Released under the MIT license.2"""Character/control-token tokenizer for the 3-digit-basic-calc calculator model.3 4Digits, operators and symbols are single characters; the scratchpad grammar adds5multi-character control tokens (``<add>``, ``<step>``, ``<qmul>`` ...). A minus6sign is emitted as the unary token ``~`` at the start of an expression, after7``=``, or immediately after a binary operator, so ``12*-34`` and ``5--3`` stay8unambiguous while ``-`` remains the binary subtraction operator.9"""10 11from __future__ import annotations12 13import json14import os15 16from transformers import PreTrainedTokenizer17 18 19SPECIAL_TOKENS = ["<pad>", "<bos>", "<eos>", "<unk>"]20DIGITS = list("0123456789")21OPERATORS = ["+", "-", "/", "*"]22SYMBOLS = ["=", ".", "~"]23CONTROL_TOKENS = [24    "<add>", "<sub>", "<mul>", "<div>", "<pos>", "<neg>",25    "<state>", "<step>", "<ans>", "<nan>", "<col>", "<qmul>", "<rem>",26]27TOKENS = SPECIAL_TOKENS + DIGITS + OPERATORS + SYMBOLS + CONTROL_TOKENS28_UNARY_PREDECESSORS = {"=", "+", "-", "/", "*"}29 30 31class ThreeDigitBasicCalcTokenizer(PreTrainedTokenizer):32    vocab_files_names = {"vocab_file": "vocab.json"}33    model_input_names = ["input_ids", "attention_mask"]34 35    def __init__(self, vocab_file=None, **kwargs):36        if vocab_file and os.path.isfile(vocab_file):37            with open(vocab_file, encoding="utf-8") as handle:38                self._vocab = json.load(handle)39        else:40            self._vocab = {tok: i for i, tok in enumerate(TOKENS)}41        self._ids_to_tokens = {i: t for t, i in self._vocab.items()}42        self._ordered_controls = sorted(CONTROL_TOKENS, key=len, reverse=True)43        kwargs.setdefault("pad_token", "<pad>")44        kwargs.setdefault("bos_token", "<bos>")45        kwargs.setdefault("eos_token", "<eos>")46        kwargs.setdefault("unk_token", "<unk>")47        # Generation takes the logits from the final sequence position.  Left48        # padding therefore keeps the final position on a real prompt token for49        # every member of a mixed-length batch.50        kwargs.setdefault("padding_side", "left")51        super().__init__(**kwargs)52 53    @property54    def vocab_size(self):55        return len(self._vocab)56 57    def get_vocab(self):58        return dict(self._vocab)59 60    def _tokenize(self, text):61        tokens = []62        index = 063        while index < len(text):64            control = next(65                (t for t in self._ordered_controls if text.startswith(t, index)),66                None,67            )68            if control is not None:69                tokens.append(control)70                index += len(control)71                continue72            ch = text[index]73            if ch == "-" and (74                    index == 0 or text[index - 1] in _UNARY_PREDECESSORS):75                tokens.append("~")76            else:77                tokens.append(ch)78            index += 179        return tokens80 81    def _convert_token_to_id(self, token):82        return self._vocab.get(token, self._vocab["<unk>"])83 84    def _convert_id_to_token(self, index):85        return self._ids_to_tokens.get(index, "<unk>")86 87    def convert_tokens_to_string(self, tokens):88        out = []89        for tok in tokens:90            if tok in SPECIAL_TOKENS:91                continue92            out.append("-" if tok == "~" else tok)93        return "".join(out)94 95    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):96        bos = [self._vocab["<bos>"]]97        if token_ids_1 is None:98            return bos + token_ids_099        return bos + token_ids_0 + token_ids_1100 101    def save_vocabulary(self, save_directory, filename_prefix=None):102        path = os.path.join(103            save_directory,104            (filename_prefix + "-" if filename_prefix else "") + "vocab.json",105        )106        with open(path, "w", encoding="utf-8") as handle:107            json.dump(self._vocab, handle, ensure_ascii=False, indent=2)108        return (path,)109