ScalableMath/Lean-STaR-plus

1# coding=utf-82# Copyright (c) The InternLM team and The HuggingFace Inc. team. All rights reserved.3#4# This code is based on transformers/src/transformers/models/llama/tokenization_llama.py5#6# Licensed under the Apache License, Version 2.0 (the "License");7# you may not use this file except in compliance with the License.8# You may obtain a copy of the License at9#10#     http://www.apache.org/licenses/LICENSE-2.011#12# Unless required by applicable law or agreed to in writing, software13# distributed under the License is distributed on an "AS IS" BASIS,14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.15# See the License for the specific language governing permissions and16# limitations under the License.17 18"""Tokenization classes for InternLM."""19import os20from shutil import copyfile21from typing import Any, Dict, List, Optional, Tuple22 23import sentencepiece as spm24from transformers.tokenization_utils import PreTrainedTokenizer25from transformers.utils import logging26 27logger = logging.get_logger(__name__)28 29VOCAB_FILES_NAMES = {"vocab_file": "./tokenizer.model"}30 31PRETRAINED_VOCAB_FILES_MAP = {}32 33 34# Modified from transformers.model.llama.tokenization_llama.LlamaTokenizer35class InternLM2Tokenizer(PreTrainedTokenizer):36    """37    Construct a InternLM2 tokenizer. Based on byte-level Byte-Pair-Encoding.38 39    Args:40        vocab_file (`str`):41            Path to the vocabulary file.42    """43 44    vocab_files_names = VOCAB_FILES_NAMES45    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP46    model_input_names = ["input_ids", "attention_mask"]47    _auto_class = "AutoTokenizer"48 49    def __init__(50        self,51        vocab_file,52        unk_token="<unk>",53        bos_token="<s>",54        eos_token="</s>",55        pad_token="</s>",56        sp_model_kwargs: Optional[Dict[str, Any]] = None,57        add_bos_token=True,58        add_eos_token=False,59        decode_with_prefix_space=False,60        clean_up_tokenization_spaces=False,61        **kwargs,62    ):63        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs64        self.vocab_file = vocab_file65        self.add_bos_token = add_bos_token66        self.add_eos_token = add_eos_token67        self.decode_with_prefix_space = decode_with_prefix_space68        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)69        self.sp_model.Load(vocab_file)70        self._no_prefix_space_tokens = None71        super().__init__(72            bos_token=bos_token,73            eos_token=eos_token,74            unk_token=unk_token,75            pad_token=pad_token,76            clean_up_tokenization_spaces=clean_up_tokenization_spaces,77            **kwargs,78        )79 80    @property81    def no_prefix_space_tokens(self):82        if self._no_prefix_space_tokens is None:83            vocab = self.convert_ids_to_tokens(list(range(self.vocab_size)))84            self._no_prefix_space_tokens = {i for i, tok in enumerate(vocab) if not tok.startswith("▁")}85        return self._no_prefix_space_tokens86 87    @property88    def vocab_size(self):89        """Returns vocab size"""90        return self.sp_model.get_piece_size()91 92    @property93    def bos_token_id(self) -> Optional[int]:94        return self.sp_model.bos_id()95 96    @property97    def eos_token_id(self) -> Optional[int]:98        return self.sp_model.eos_id()99 100    def get_vocab(self):101        """Returns vocab as a dict"""102        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}103        vocab.update(self.added_tokens_encoder)104        return vocab105 106    def _tokenize(self, text):107        """Returns a tokenized string."""108        return self.sp_model.encode(text, out_type=str)109 110    def _convert_token_to_id(self, token):111        """Converts a token (str) in an id using the vocab."""112        return self.sp_model.piece_to_id(token)113 114    def _convert_id_to_token(self, index):115        """Converts an index (integer) in a token (str) using the vocab."""116        token = self.sp_model.IdToPiece(index)117        return token118 119    def _maybe_add_prefix_space(self, tokens, decoded):120        if tokens and tokens[0] not in self.no_prefix_space_tokens:121            return " " + decoded122        else:123            return decoded124 125    def convert_tokens_to_string(self, tokens):126        """Converts a sequence of tokens (string) in a single string."""127        current_sub_tokens = []128        out_string = ""129        prev_is_special = False130        for token in tokens:131            # make sure that special tokens are not decoded using sentencepiece model132            if token in self.all_special_tokens:133                if not prev_is_special:134                    out_string += " "135                out_string += self.sp_model.decode(current_sub_tokens) + token136                prev_is_special = True137                current_sub_tokens = []138            else:139                current_sub_tokens.append(token)140                prev_is_special = False141        out_string += self.sp_model.decode(current_sub_tokens)142        out_string = self.clean_up_tokenization(out_string)143        out_string = self._maybe_add_prefix_space(tokens=tokens, decoded=out_string)144        return out_string[1:]145 146    def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:147        """148        Save the vocabulary and special tokens file to a directory.149 150        Args:151            save_directory (`str`):152                The directory in which to save the vocabulary.153 154        Returns:155            `Tuple(str)`: Paths to the files saved.156        """157        if not os.path.isdir(save_directory):158            logger.error(f"Vocabulary path ({save_directory}) should be a directory")159            return160        out_vocab_file = os.path.join(161            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]162        )163 164        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):165            copyfile(self.vocab_file, out_vocab_file)166        elif not os.path.isfile(self.vocab_file):167            with open(out_vocab_file, "wb") as fi:168                content_spiece_model = self.sp_model.serialized_model_proto()169                fi.write(content_spiece_model)170 171        return (out_vocab_file,)172 173    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):174        if self.add_bos_token:175            bos_token_ids = [self.bos_token_id]176        else:177            bos_token_ids = []178 179        output = bos_token_ids + token_ids_0180 181        if token_ids_1 is not None:182            output = output + token_ids_1183 184        if self.add_eos_token:185            output = output + [self.eos_token_id]186 187        return output188 189    def get_special_tokens_mask(190        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False191    ) -> List[int]:192        """193        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding194        special tokens using the tokenizer `prepare_for_model` method.195 196        Args:197            token_ids_0 (`List[int]`):198                List of IDs.199            token_ids_1 (`List[int]`, *optional*):200                Optional second list of IDs for sequence pairs.201            already_has_special_tokens (`bool`, *optional*, defaults to `False`):202                Whether or not the token list is already formatted with special tokens for the model.203 204        Returns:205            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.206        """207        if already_has_special_tokens:208            return super().get_special_tokens_mask(209                token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True210            )211 212        if token_ids_1 is None:213            return [1] + ([0] * len(token_ids_0)) + [1]214        return [1] + ([0] * len(token_ids_0)) + [1, 1] + ([0] * len(token_ids_1)) + [1]215 216    def create_token_type_ids_from_sequences(217        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None218    ) -> List[int]:219        """220        Create a mask from the two sequences passed to be used in a sequence-pair classification task. T5 does not make221        use of token type ids, therefore a list of zeros is returned.222 223        Args:224            token_ids_0 (`List[int]`):225                List of IDs.226            token_ids_1 (`List[int]`, *optional*):227                Optional second list of IDs for sequence pairs.228 229        Returns:230            `List[int]`: List of zeros.231        """232        eos = [self.eos_token_id]233 234        if token_ids_1 is None:235            return len(token_ids_0 + eos) * [0]236        return len(token_ids_0 + eos + token_ids_1 + eos) * [0]237