optimum-rbln 0.1.1__py3-none-any.whl → 0.1.4__py3-none-any.whl

optimum/rbln/transformers/models/llama/modeling_llama.py

@@ -39,12 +39,20 @@ from ....modeling_base import RBLNBaseModel
 from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
 from ....utils.runtime_utils import RBLNPytorchRuntime
 from ....utils.save_utils import maybe_save_preprocessors
+
+
+# FIXME:: Merge Two architecture Codes
 from .llama_architecture import (
     LlamaWrapper,
     wrap_llama,
     unwrap_llama,
 )
 
+from .llama_architecture_cb import (
+    LlamaDynamicBatchWrapper as LlamaWrapper_cb,
+    wrap_llama as wrap_llama_cb,
+)
+
 
 logger = logging.getLogger(__name__)
 
@@ -57,24 +65,12 @@ if TYPE_CHECKING:
     )
 
 
+SUPPORTED_BATCHING_MODES = ["static", "vllm"]
+
+
 class RBLNRuntimeModel(RBLNPytorchRuntime):
     mandatory_members = ["main_input_name"]
 
-    # RBLN_Runtimemodule
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: torch.LongTensor = None,
-        cache_position: torch.Tensor = None,
-        **kwargs: Dict[str, Any],
-    ):
-        logits = super().forward(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            cache_position=cache_position,
-        )
-        return logits
-
 
 class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
     """
@@ -95,13 +91,16 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
         self.batch_size = self.rbln_config.meta["rbln_batch_size"]
         self.max_seq_len = self.rbln_config.meta["rbln_max_seq_len"]
         self.prefill_chunk_size = self.rbln_config.meta["rbln_prefill_chunk_size"]
+        self.use_continuous_batch = self.rbln_config.meta["rbln_batching"] == "vllm"
 
+        prefill_batch_size = self.batch_size if not self.use_continuous_batch else 1
         self.prefill_attention_mask = torch.zeros(
-            self.batch_size, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.int64
+            prefill_batch_size, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.int64
         )
         self.causal_mask = 1 - torch.triu(
-            torch.ones(self.batch_size, 1, self.prefill_chunk_size, self.prefill_chunk_size), diagonal=1
+            torch.ones(prefill_batch_size, 1, self.prefill_chunk_size, self.prefill_chunk_size), diagonal=1
         )
+        self.decoder_attention_mask = torch.zeros(self.batch_size, 1, 1, self.max_seq_len, dtype=torch.int64)
 
         self.prefill_decoder = RBLNRuntimeModel(runtime=self.runtimes[0], main_input_name="input_ids")
         self.decoder = RBLNRuntimeModel(runtime=self.runtimes[1], main_input_name="input_ids")
@@ -164,7 +163,10 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
 
         rbln_config_kwargs, rbln_constructor_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
 
-        origin_mehtods = wrap_llama()
+        # FIXME :: This should be moved when wrapping removed.
+        use_continuous_batch = rbln_config_kwargs.get("rbln_batching", "static") == "vllm"
+        origin_mehtods = wrap_llama_cb() if use_continuous_batch else wrap_llama()
+
         model: LlamaForCausalLM = TasksManager.get_model_from_task(
             task=task,
             model_name_or_path=model_id,
@@ -191,14 +193,18 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
                 preprocessors=preprocessors, model_config=model.config, **rbln_config_kwargs
             )
 
-        def compile_llama():
-            wrapped_model = LlamaWrapper(model).eval()
+        def compile_llama(use_continuous_batch, wrapper_cls):
+            wrapped_model = wrapper_cls(model).eval()
 
             prefill_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][0]
             dec_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][1]
 
             prefill_example_inputs = prefill_rbln_runtime_config.get_dummy_inputs(fill=0)
-            dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=0)
+            dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=4)
+
+            if use_continuous_batch:
+                batch_index_index = 3
+                dec_example_inputs[batch_index_index].fill_(-1)  # fill batch_position -1 to indicate it is decoder.
 
             prefill_scripted_model = torch.jit.trace(wrapped_model, prefill_example_inputs)
             dec_scripted_model = torch.jit.trace(wrapped_model, dec_example_inputs)
@@ -213,8 +219,9 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
             )
 
             # Caching prefill_decoder/decoder I/O
+            cache_index_offset = 4 if use_continuous_batch else 3
             connections = [
-                (prefill_ir.outputs[1 + i], prefill_ir.inputs[3 + i])
+                (prefill_ir.outputs[1 + i], prefill_ir.inputs[cache_index_offset + i])
                 for i in range(model.config.num_hidden_layers * 2)
             ]
 
@@ -229,7 +236,8 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
             )
             compiled_model.save(save_dir_path / f"{DEFAULT_COMPILED_MODEL_NAME}.rbln")
 
-        compile_llama()
+        wrapper_cls = LlamaWrapper_cb if use_continuous_batch else LlamaWrapper
+        compile_llama(use_continuous_batch=use_continuous_batch, wrapper_cls=wrapper_cls)
         unwrap_llama(origin_mehtods)
 
         rbln_config.save(save_dir_path)
@@ -249,6 +257,7 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
         model_config: "PretrainedConfig",
         rbln_max_seq_len: Optional[int] = None,
         rbln_batch_size: Optional[int] = None,
+        rbln_batching: Optional[str] = None,
     ) -> RBLNConfig:
         meta = {}
 
@@ -256,21 +265,38 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
         if rbln_max_seq_len is None:
             rbln_max_seq_len = getattr(model_config, "max_position_embeddings", None)
         rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
+        rbln_batching = "static" if rbln_batching is None else rbln_batching
 
         meta["rbln_max_seq_len"] = rbln_max_seq_len
         meta["rbln_batch_size"] = rbln_batch_size
         meta["rbln_prefill_chunk_size"] = prefill_chunk_size
+        meta["rbln_batching"] = rbln_batching
+        use_continuous_batching = meta["rbln_batching"] == "vllm"
+
+        if rbln_batching not in SUPPORTED_BATCHING_MODES:
+            raise ValueError(
+                f'rbln_batching="{rbln_batching}" is not a supported batch mode, '
+                f"Possible: {SUPPORTED_BATCHING_MODES}"
+            )
 
-        def get_input_info(query_length):
+        def get_input_info(
+            batch_size,  # should be 1 if continous batch prefill
+            query_length,
+            continuous_batch=False,  # determines the shape of `cache position`
+        ):
             input_info = [
-                ("input_ids", [rbln_batch_size, query_length], "int64"),
-                ("attention_mask", [rbln_batch_size, 1, query_length, rbln_max_seq_len], "int64"),
+                ("input_ids", [batch_size, query_length], "int64"),
+                ("attention_mask", [batch_size, 1, query_length, rbln_max_seq_len], "int64"),
                 (
                     "cache_position",
-                    [],
+                    [batch_size, query_length] if continuous_batch else [],
                     "int32",
                 ),
             ]
+
+            if continuous_batch:
+                input_info.append(("batch_position", [], "int16"))
+
             input_info.extend(
                 [
                     (
@@ -286,10 +312,19 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
                     for i in range(model_config.num_hidden_layers * 2)
                 ]
             )
+
             return input_info
 
-        prefill_input_info = get_input_info(query_length=prefill_chunk_size)
-        dec_input_info = get_input_info(query_length=1)
+        prefill_input_info = get_input_info(
+            batch_size=1 if use_continuous_batching else rbln_batch_size,
+            query_length=prefill_chunk_size,
+            continuous_batch=use_continuous_batching,
+        )
+        dec_input_info = get_input_info(
+            batch_size=rbln_batch_size,
+            query_length=1,
+            continuous_batch=use_continuous_batching,
+        )
 
         prefill_rbln_runtime_config = RBLNRuntimeConfig(input_info=prefill_input_info)
         dec_rbln_runtime_config = RBLNRuntimeConfig(input_info=dec_input_info)
@@ -337,7 +372,6 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
 
         # In greedy decoding
         if past_cached_length == 0:
-
             # padding with prefill_chunk_size
             # TODO left padding + left padding has issue on stoppingcriteria(max_len)
             if cur_len % self.prefill_chunk_size != 0:
@@ -384,7 +418,13 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
 
         return model_inputs
 
-    def forward(
+    def forward(self, *args, **kwargs):
+        if self.use_continuous_batch:
+            return self.forward_cb(*args, **kwargs)
+        else:
+            return self.forward_static(*args, **kwargs)
+
+    def forward_static(
         self,
         input_ids: torch.LongTensor = None,
         attention_mask: Optional[torch.Tensor] = None,
@@ -393,7 +433,6 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
         query_length: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> Tuple[torch.FloatTensor]:
-
         if past_key_values is not None:
             past_key_values += query_length
 
@@ -425,3 +464,58 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
             logits=outputs,
             past_key_values=past_key_values,
         )
+
+    def forward_cb(
+        self,
+        input_ids: torch.LongTensor = None,
+        cache_position: Optional[torch.Tensor] = None,  # torch.tensor(,dtype=int32) (1,64) // (4,1)
+        batch_idx: int = None,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor]:
+        # prefill_decoder
+        if cache_position.shape[1] > 1:
+            query_length = input_ids.shape[1]
+            attention_mask = self.prefill_attention_mask.clone()
+            for step in range(0, query_length, self.prefill_chunk_size):
+                if step + self.prefill_chunk_size > query_length:
+                    input_ids = torch.nn.functional.pad(input_ids, (0, step + self.prefill_chunk_size - query_length))
+                    cache_position = torch.cat(
+                        [
+                            cache_position,
+                            torch.arange(
+                                query_length,
+                                step + self.prefill_chunk_size,
+                                dtype=torch.int32,
+                            ).unsqueeze(0),
+                        ],
+                        dim=-1,
+                    )
+
+                sliced_input_ids = input_ids[:, step : step + self.prefill_chunk_size]
+                sliced_cache_positions = cache_position[:, step : step + self.prefill_chunk_size]
+                attention_mask[:, :, :, :step] = 1
+                attention_mask[:, :, :, step : step + self.prefill_chunk_size] = self.causal_mask
+
+                outputs, _ = self.prefill_decoder(
+                    sliced_input_ids.contiguous(),
+                    attention_mask.contiguous(),
+                    sliced_cache_positions.contiguous(),
+                    torch.tensor(batch_idx, dtype=torch.int16),
+                )
+            outputs = outputs[:, query_length % self.prefill_chunk_size - 1].unsqueeze(1)
+        # decoder
+        else:
+            attention_mask = self.decoder_attention_mask.clone()
+            for b_idx in range(self.batch_size):
+                attention_mask[b_idx, :, :, : cache_position[b_idx].item() + 1] = 1
+
+            outputs = self.decoder(
+                input_ids.contiguous(),
+                attention_mask.contiguous(),
+                cache_position.contiguous(),
+                torch.tensor(0, dtype=torch.int16),
+            )[0]
+
+        return CausalLMOutputWithPast(
+            logits=outputs,
+        )

optimum/rbln/transformers/models/midm/__init__.py

@@ -0,0 +1,32 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+import os
+from os import environ
+
+
+this_path = os.path.abspath(__file__)
+local_dir = "/" + os.path.join(*this_path.split("/")[:-1]) + "/hf_hub_cached"
+environ["LOCAL_CACHE_ROOT_CUSTOM_CODE_MIDM"] = local_dir
+
+from .modeling_midm import RBLNMidmLMHeadModel

optimum/rbln/transformers/models/midm/hf_hub_cached/configuration_midm.py

@@ -0,0 +1,22 @@
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
+
+
+class MidmBitextConfig(GPT2Config):
+    model_type = "midm-bitext-S"
+
+    def __init__(
+        self,
+        use_absolute_position_embedding: bool = True,
+        use_rotary_position_embedding: bool = False,
+        rotary_percentage: float = 1.0,
+        normalization_type: str = "layernorm",
+        scale_qk_by_inverse_layer_idx: bool = False,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.use_absolute_position_embedding = use_absolute_position_embedding
+        self.use_rotary_position_embedding = use_rotary_position_embedding
+        self.rotary_percentage = rotary_percentage
+        self.normalization_type = normalization_type
+        self.scale_qk_by_inverse_layer_idx = scale_qk_by_inverse_layer_idx

optimum/rbln/transformers/models/midm/hf_hub_cached/midm_bitext_tokenization.py

@@ -0,0 +1,303 @@
+# coding=utf-8
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Tokenization class for model Midm_bitext_tonkenizer."""
+import os
+import re
+import warnings
+from shutil import copyfile
+from typing import Any, Dict, List, Optional, Tuple
+
+import sentencepiece as spm
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "midm_bitext_tokenizer.model"}
+
+PRETRAINED_VOCAB_FILES_MAP = {}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {}
+
+
+class Midm_bitext_Tokenizer(PreTrainedTokenizer):
+    """
+    Construct a Midm bitext tonkenizer. Based on [SentencePiece](https://github.com/google/sentencepiece).
+
+    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+    this superclass for more information regarding those methods.
+
+    Args:
+        vocab_file (`str`):
+            [SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that
+            contains the vocabulary necessary to instantiate a tokenizer.
+        eos_token (`str`, *optional*, defaults to `"</s>"`):
+            The end of sequence token.
+
+            <Tip>
+
+            When building a sequence using special tokens, this is not the token that is used for the end of sequence.
+            The token used is the `sep_token`.
+
+            </Tip>
+
+        unk_token (`str`, *optional*, defaults to `"<unk>"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+        pad_token (`str`, *optional*, defaults to `"<pad>"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        extra_ids (`int`, *optional*, defaults to 100):
+            Add a number of extra ids added to the end of the vocabulary for use as sentinels. These tokens are
+            accessible as "<extra_id_{%d}>" where "{%d}" is a number between 0 and extra_ids-1. Extra tokens are
+            indexed from the end of the vocabulary up to beginning.
+        additional_special_tokens (`List[str]`, *optional*):
+            Additional special tokens used by the tokenizer.
+        sp_model_kwargs (`dict`, *optional*):
+            Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
+            SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
+            to set:
+
+            - `enable_sampling`: Enable subword regularization.
+            - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
+
+              - `nbest_size = {0,1}`: No sampling is performed.
+              - `nbest_size > 1`: samples from the nbest_size results.
+              - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
+                using forward-filtering-and-backward-sampling algorithm.
+
+            - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
+              BPE-dropout.
+
+    Attributes:
+        sp_model (`SentencePieceProcessor`):
+            The *SentencePiece* processor that is used for every conversion (string, tokens and IDs).
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        eos_token="</s>",
+        unk_token="<unk>",
+        pad_token="<pad>",
+        extra_ids=100,
+        additional_special_tokens=None,
+        sp_model_kwargs: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> None:
+        # Add extra_ids to the special token list
+        if extra_ids > 0 and additional_special_tokens is None:
+            additional_special_tokens = [f"<extra_id_{i}>" for i in range(extra_ids)]
+        elif extra_ids > 0 and additional_special_tokens is not None:
+            # Check that we have the right number of extra_id special tokens
+            extra_tokens = len(set(filter(lambda x: bool("extra_id" in str(x)), additional_special_tokens)))
+            if extra_tokens != extra_ids:
+                raise ValueError(
+                    f"Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are provided to Midm_bitext_Tonkenizer. "
+                    "In this case the additional_special_tokens must include the extra_ids tokens"
+                )
+
+        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
+
+        # custom special tokens
+        # convert \n, \t in input text -> <[!newline]>, <[!tab]>
+        self.newline_token = "<[!newline]>"
+        self.tab_token = "<[!tab]>"
+
+        self.vocab_file = vocab_file
+        self._extra_ids = extra_ids
+
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(vocab_file)
+        super().__init__(
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            extra_ids=extra_ids,
+            additional_special_tokens=additional_special_tokens,
+            sp_model_kwargs=self.sp_model_kwargs,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self):
+        return self.sp_model.get_piece_size() + self._extra_ids
+
+    def get_vocab(self):
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def get_special_tokens_mask(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
+    ) -> List[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer `prepare_for_model` method.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+
+        Returns:
+            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+            )
+
+        # normal case: some special tokens
+        if token_ids_1 is None:
+            return ([0] * len(token_ids_0)) + [1]
+        return ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
+
+    def _add_eos_if_not_present(self, token_ids: List[int]) -> List[int]:
+        """Do not add eos again if user already added it."""
+        if len(token_ids) > 0 and token_ids[-1] == self.eos_token_id:
+            warnings.warn(
+                f"This sequence already has {self.eos_token}. In future versions this behavior may lead to duplicated eos tokens being added."
+            )
+            return token_ids
+        else:
+            return token_ids + [self.eos_token_id]
+
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Create a mask from the two sequences passed to be used in a sequence-pair classification task. Midm does not make
+        use of token type ids, therefore a list of zeros is returned.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of zeros.
+        """
+        eos = [self.eos_token_id]
+
+        if token_ids_1 is None:
+            return len(token_ids_0 + eos) * [0]
+        return len(token_ids_0 + eos + token_ids_1 + eos) * [0]
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. A sequence has the following format:
+
+        - single sequence: `X </s>`
+        - pair of sequences: `A </s> B </s>`
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+        """
+        token_ids_0 = self._add_eos_if_not_present(token_ids_0)
+        if token_ids_1 is None:
+            return token_ids_0
+        else:
+            token_ids_1 = self._add_eos_if_not_present(token_ids_1)
+            return token_ids_0 + token_ids_1
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["sp_model"] = None
+        return state
+
+    def __setstate__(self, d):
+        self.__dict__ = d
+
+        # for backward compatibility
+        if not hasattr(self, "sp_model_kwargs"):
+            self.sp_model_kwargs = {}
+
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(self.vocab_file)
+
+    def _tokenize(self, text: str) -> List[str]:
+        """Take as input a string and return a list of strings (tokens) for words/sub-words"""
+        text = text.replace("\n", self.newline_token)
+        text = text.replace("\t", self.tab_token)
+
+        return self.sp_model.encode(text, out_type=str)
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        if token.startswith("<extra_id_"):
+            match = re.match(r"<extra_id_(\d+)>", token)
+            num = int(match.group(1))
+            return self.vocab_size - num - 1
+        return self.sp_model.piece_to_id(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        if index < self.sp_model.get_piece_size():
+            token = self.sp_model.IdToPiece(index)
+        else:
+            token = f"<extra_id_{self.vocab_size - 1 - index}>"
+        return token
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        current_sub_tokens = []
+        out_string = ""
+        for token in tokens:
+            # make sure that special tokens are not decoded using sentencepiece model
+            if token in self.all_special_tokens:
+                out_string += self.sp_model.decode_pieces(current_sub_tokens) + token + " "
+                current_sub_tokens = []
+            else:
+                current_sub_tokens.append(token)
+        out_string += self.sp_model.decode_pieces(current_sub_tokens)
+
+        out_string.replace(self.newline_token, "\n")
+        out_string.replace(self.tab_token, "\t")
+
+        return out_string.strip()
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        out_vocab_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+        )
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+        elif not os.path.isfile(self.vocab_file):
+            with open(out_vocab_file, "wb") as fi:
+                content_spiece_model = self.sp_model.serialized_model_proto()
+                fi.write(content_spiece_model)
+
+        return (out_vocab_file,)