SlitherCode/tiny-edu-166m

1import os2import json3import tiktoken4from transformers import PreTrainedTokenizer5 6VOCAB_FILES_NAMES = {"vocab_file": "tiktoken_encoding.json"}7 8 9class TiktokenTokenizer(PreTrainedTokenizer):10    """11    HuggingFace-compatible tokenizer wrapping tiktoken's cl100k_base.12    Produces byte-identical token IDs to tiktoken.get_encoding("cl100k_base").13    Tokens are represented internally as their raw bytes decoded via latin-114    (a lossless bijection for arbitrary byte sequences).15    """16 17    vocab_files_names = VOCAB_FILES_NAMES18    model_input_names = ["input_ids", "attention_mask"]19 20    def __init__(self, vocab_file=None, encoding_name="cl100k_base", **kwargs):21        self.encoding_name = encoding_name22        self._build_enc()23 24        eot_str = self._id_to_tok[self._enc.eot_token]  # "<|endoftext|>"25 26        # When loading from saved config, special tokens are already in kwargs —27        # pop them so we don't pass duplicates to super().__init__().28        kwargs.pop("bos_token", None)29        kwargs.pop("eos_token", None)30        kwargs.pop("pad_token", None)31        kwargs.pop("unk_token", None)32 33        super().__init__(34            encoding_name=encoding_name,35            bos_token=eot_str,36            eos_token=eot_str,37            pad_token=eot_str,38            unk_token=eot_str,39            **kwargs,40        )41 42    def _build_enc(self):43        self._enc = tiktoken.get_encoding(self.encoding_name)44        self._id_to_tok = {}45        self._tok_to_id = {}46        for i in range(self._enc.n_vocab):47            try:48                s = self._enc.decode_single_token_bytes(i).decode("latin-1")49            except Exception:50                s = f"<|special_{i}|>"51            self._id_to_tok[i] = s52            self._tok_to_id[s] = i53 54    # ── Required interface ─────────────────────────────────────────────────────55 56    @property57    def vocab_size(self):58        return self._enc.n_vocab  # 10027759 60    def get_vocab(self):61        return dict(self._tok_to_id)62 63    def _tokenize(self, text):64        ids = self._enc.encode(text, allowed_special="all")65        return [self._id_to_tok[i] for i in ids]66 67    def _convert_token_to_id(self, token):68        return self._tok_to_id.get(token, self._enc.eot_token)69 70    def _convert_id_to_token(self, index):71        return self._id_to_tok.get(index, "<|unk|>")72 73    def convert_tokens_to_string(self, tokens):74        raw = b"".join(t.encode("latin-1") for t in tokens)75        return raw.decode("utf-8", errors="replace")76 77    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):78        """No BOS/EOS added — matches bare tiktoken encode() behaviour."""79        if token_ids_1 is None:80            return token_ids_081        return token_ids_0 + token_ids_182 83    def save_vocabulary(self, save_directory, filename_prefix=None):84        os.makedirs(save_directory, exist_ok=True)85        fname = (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]86        vocab_file = os.path.join(save_directory, fname)87        with open(vocab_file, "w", encoding="utf-8") as f:88            json.dump({"encoding_name": self.encoding_name}, f)89        return (vocab_file,)90 91    # ── Pickle support (tiktoken objects aren't picklable) ─────────────────────92 93    def __getstate__(self):94        state = self.__dict__.copy()95        state.pop("_enc", None)96        state.pop("_id_to_tok", None)97        state.pop("_tok_to_id", None)98        return state99 100    def __setstate__(self, state):101        self.__dict__.update(state)102        self._build_enc()103