




'''simple docstring'''






import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule

class  UpperCAmelCase__						(       A ):
    lowerCAmelCase_				    =		42
    lowerCAmelCase_				    =		jnp.floataa
    lowerCAmelCase_				    =		True



    def   lowerCamelCase_  (	self      :       Any     ):
          super().setup()
          _lowerCamelCase				:   str								=   nn.Dense(5,dtype=self.dtype     )



    def __call__(	self      :       Optional[int],*__A      :       Optional[int],**__A      :       Union[str, Any]     ):
          _lowerCamelCase				:   Optional[int]								=   super().__call__(*__A,**__A     )
          _lowerCamelCase				:   int								=   self.cls(outputs[2]     )
          return outputs[:2] + (cls_out,)


class  UpperCAmelCase__						(       A ):
    lowerCAmelCase_				    =		FlaxBigBirdForNaturalQuestionsModule







def 		A_       (	_lowerCAmelCase   :    Optional[Any]				,	_lowerCAmelCase   :    Dict				,	_lowerCAmelCase   :    int				,	_lowerCAmelCase   :    Optional[int]				,	_lowerCAmelCase   :    Tuple				,	_lowerCAmelCase   :    Dict							):



      """simple docstring"""




      def cross_entropy(_lowerCAmelCase   :    Dict				,	_lowerCAmelCase   :    str				,	_lowerCAmelCase   :    Tuple=None							):
            _lowerCamelCase				:   Optional[int]								=   logits.shape[-1]
            _lowerCamelCase				:   Optional[Any]								=   (labels[..., None] == jnp.arange(_lowerCAmelCase							)[None]).astype("f4"							)
            _lowerCamelCase				:   str								=   jax.nn.log_softmax(_lowerCAmelCase				,	axis=-1							)
            _lowerCamelCase				:   int								=   -jnp.sum(labels * logits				,	axis=-1							)
            if reduction is not None:
                  _lowerCamelCase				:   List[Any]								=   reduction(_lowerCAmelCase							)
            return loss

      _lowerCamelCase				:   int								=   partial(_lowerCAmelCase				,	reduction=jnp.mean							)
      _lowerCamelCase				:   List[Any]								=   cross_entropy(_lowerCAmelCase				,	_lowerCAmelCase							)
      _lowerCamelCase				:   Union[str, Any]								=   cross_entropy(_lowerCAmelCase				,	_lowerCAmelCase							)
      _lowerCamelCase				:   Optional[int]								=   cross_entropy(_lowerCAmelCase				,	_lowerCAmelCase							)
      return (start_loss + end_loss + pooled_loss) / 3

@dataclass
class  UpperCAmelCase__						:
    lowerCAmelCase_				    =		'google/bigbird-roberta-base'
    lowerCAmelCase_				    =		3000
    lowerCAmelCase_				    =		10500

    lowerCAmelCase_				    =		128
    lowerCAmelCase_				    =		3

    lowerCAmelCase_				    =		1
    lowerCAmelCase_				    =		5

    # tx_args
    lowerCAmelCase_				    =		3E-5
    lowerCAmelCase_				    =		0.0
    lowerCAmelCase_				    =		20000
    lowerCAmelCase_				    =		0.0_0_9_5

    lowerCAmelCase_				    =		'bigbird-roberta-natural-questions'
    lowerCAmelCase_				    =		'training-expt'
    lowerCAmelCase_				    =		'data/nq-training.jsonl'
    lowerCAmelCase_				    =		'data/nq-validation.jsonl'



    def   lowerCamelCase_  (	self      :       Optional[int]     ):
          os.makedirs(self.base_dir,exist_ok=__A     )
          _lowerCamelCase				:   Union[str, Any]								=   os.path.join(self.base_dir,self.save_dir     )
          _lowerCamelCase				:   Any								=   self.batch_size_per_device * jax.device_count()


@dataclass
class  UpperCAmelCase__						:
    lowerCAmelCase_				    =		42
    lowerCAmelCase_				    =		4096  # no dynamic padding on TPUs



    def __call__(	self      :       Union[str, Any],__A      :       List[str]     ):
          _lowerCamelCase				:   Dict								=   self.collate_fn(__A     )
          _lowerCamelCase				:   int								=   jax.tree_util.tree_map(__A,__A     )
          return batch



    def   lowerCamelCase_  (	self      :       List[Any],__A      :       Optional[Any]     ):
          _lowerCamelCase				:   Union[str, Any]								=   self.fetch_inputs(features["input_ids"]     )
          _lowerCamelCase				:   Dict								=   {
              "input_ids": jnp.array(__A,dtype=jnp.intaa     ),
              "attention_mask": jnp.array(__A,dtype=jnp.intaa     ),
              "start_labels": jnp.array(features["start_token"],dtype=jnp.intaa     ),
              "end_labels": jnp.array(features["end_token"],dtype=jnp.intaa     ),
              "pooled_labels": jnp.array(features["category"],dtype=jnp.intaa     ),
          }
          return batch



    def   lowerCamelCase_  (	self      :       Union[str, Any],__A      :       list     ):
          _lowerCamelCase				:   str								=   [self._fetch_inputs(__A     ) for ids in input_ids]
          return zip(*__A     )



    def   lowerCamelCase_  (	self      :       Tuple,__A      :       list     ):
          _lowerCamelCase				:   str								=   [1 for _ in range(len(__A     )     )]
          while len(__A     ) < self.max_length:
                input_ids.append(self.pad_id     )
                attention_mask.append(0     )
          return input_ids, attention_mask







def 		A_       (	_lowerCAmelCase   :    Any				,	_lowerCAmelCase   :    List[str]				,	_lowerCAmelCase   :    Dict=None							):



      """simple docstring"""




      if seed is not None:
            _lowerCamelCase				:   Union[str, Any]								=   dataset.shuffle(seed=_lowerCAmelCase							)
      for i in range(len(_lowerCAmelCase							) // batch_size							):
            _lowerCamelCase				:   Any								=   dataset[i * batch_size : (i + 1) * batch_size]
            yield dict(_lowerCAmelCase							)







@partial(jax.pmap				,	axis_name="batch"							)
def 		A_       (	_lowerCAmelCase   :    Union[str, Any]				,	_lowerCAmelCase   :    List[str]				,	**_lowerCAmelCase   :    Optional[int]							):



      """simple docstring"""




      def loss_fn(_lowerCAmelCase   :    str							):
            _lowerCamelCase				:   List[Any]								=   model_inputs.pop("start_labels"							)
            _lowerCamelCase				:   Union[str, Any]								=   model_inputs.pop("end_labels"							)
            _lowerCamelCase				:   Optional[int]								=   model_inputs.pop("pooled_labels"							)

            _lowerCamelCase				:   List[str]								=   state.apply_fn(**_lowerCAmelCase				,	params=_lowerCAmelCase				,	dropout_rng=_lowerCAmelCase				,	train=_lowerCAmelCase							)
            _lowerCamelCase				:   Optional[Any]								=   outputs

            return state.loss_fn(
                _lowerCAmelCase				,	_lowerCAmelCase				,	_lowerCAmelCase				,	_lowerCAmelCase				,	_lowerCAmelCase				,	_lowerCAmelCase				,	)

      _lowerCamelCase				:   List[Any]								=   jax.random.split(_lowerCAmelCase							)
      _lowerCamelCase				:   Tuple								=   jax.value_and_grad(_lowerCAmelCase							)
      _lowerCamelCase				:   int								=   grad_fn(state.params							)
      _lowerCamelCase				:   Optional[Any]								=   jax.lax.pmean({"loss": loss}				,	axis_name="batch"							)
      _lowerCamelCase				:   Union[str, Any]								=   jax.lax.pmean(_lowerCAmelCase				,	"batch"							)

      _lowerCamelCase				:   Union[str, Any]								=   state.apply_gradients(grads=_lowerCAmelCase							)
      return state, metrics, new_drp_rng







@partial(jax.pmap				,	axis_name="batch"							)
def 		A_       (	_lowerCAmelCase   :    int				,	**_lowerCAmelCase   :    Tuple							):



      """simple docstring"""




      _lowerCamelCase				:   Optional[Any]								=   model_inputs.pop("start_labels"							)
      _lowerCamelCase				:   Any								=   model_inputs.pop("end_labels"							)
      _lowerCamelCase				:   Any								=   model_inputs.pop("pooled_labels"							)

      _lowerCamelCase				:   Optional[int]								=   state.apply_fn(**_lowerCAmelCase				,	params=state.params				,	train=_lowerCAmelCase							)
      _lowerCamelCase				:   Dict								=   outputs

      _lowerCamelCase				:   List[str]								=   state.loss_fn(_lowerCAmelCase				,	_lowerCAmelCase				,	_lowerCAmelCase				,	_lowerCAmelCase				,	_lowerCAmelCase				,	_lowerCAmelCase							)
      _lowerCamelCase				:   Union[str, Any]								=   jax.lax.pmean({"loss": loss}				,	axis_name="batch"							)
      return metrics

class  UpperCAmelCase__						(       train_state.TrainState ):
    lowerCAmelCase_				    =		struct.field(pytree_node=A )


@dataclass
class  UpperCAmelCase__						:
    lowerCAmelCase_				    =		42
    lowerCAmelCase_				    =		42
    lowerCAmelCase_				    =		42
    lowerCAmelCase_				    =		42
    lowerCAmelCase_				    =		42
    lowerCAmelCase_				    =		42
    lowerCAmelCase_				    =		None



    def   lowerCamelCase_  (	self      :       Tuple,__A      :       List[Any],__A      :       int,__A      :       List[str],__A      :       Tuple=None     ):
          _lowerCamelCase				:   Optional[Any]								=   model.params
          _lowerCamelCase				:   List[Any]								=   TrainState.create(
              apply_fn=model.__call__,params=__A,tx=__A,loss_fn=__A,)
          if ckpt_dir is not None:
                _lowerCamelCase				:   List[Any]								=   restore_checkpoint(__A,__A     )
                _lowerCamelCase				:   List[str]								=   {
                    "lr": args.lr,
                    "init_lr": args.init_lr,
                    "warmup_steps": args.warmup_steps,
                    "num_train_steps": num_train_steps,
                    "weight_decay": args.weight_decay,
                }
                _lowerCamelCase				:   Any								=   build_tx(**__A     )
                _lowerCamelCase				:   int								=   train_state.TrainState(
                    step=__A,apply_fn=model.__call__,params=__A,tx=__A,opt_state=__A,)
                _lowerCamelCase				:   int								=   args
                _lowerCamelCase				:   str								=   data_collator
                _lowerCamelCase				:   Tuple								=   lr
                _lowerCamelCase				:   Tuple								=   params
          _lowerCamelCase				:   Tuple								=   jax_utils.replicate(__A     )
          return state



    def   lowerCamelCase_  (	self      :       List[Any],__A      :       int,__A      :       Dict,__A      :       Optional[Any]     ):
          _lowerCamelCase				:   str								=   self.args
          _lowerCamelCase				:   Union[str, Any]								=   len(__A     ) // args.batch_size

          _lowerCamelCase				:   Dict								=   jax.random.PRNGKey(0     )
          _lowerCamelCase				:   Optional[int]								=   jax.random.split(__A,jax.device_count()     )
          for epoch in range(args.max_epochs     ):
                _lowerCamelCase				:   List[str]								=   jnp.array(0,dtype=jnp.floataa     )
                _lowerCamelCase				:   Any								=   get_batched_dataset(__A,args.batch_size,seed=__A     )
                _lowerCamelCase				:   Union[str, Any]								=   0
                for batch in tqdm(__A,total=__A,desc=f'Running EPOCH-{epoch}'     ):
                      _lowerCamelCase				:   Optional[Any]								=   self.data_collator(__A     )
                      _lowerCamelCase				:   Optional[Any]								=   self.train_step_fn(__A,__A,**__A     )
                      running_loss += jax_utils.unreplicate(metrics["loss"]     )
                      i += 1
                      if i % args.logging_steps == 0:
                            _lowerCamelCase				:   Dict								=   jax_utils.unreplicate(state.step     )
                            _lowerCamelCase				:   Optional[Any]								=   running_loss.item() / i
                            _lowerCamelCase				:   Tuple								=   self.scheduler_fn(state_step - 1     )

                            _lowerCamelCase				:   Any								=   self.evaluate(__A,__A     )
                            _lowerCamelCase				:   str								=   {
                                "step": state_step.item(),
                                "eval_loss": eval_loss.item(),
                                "tr_loss": tr_loss,
                                "lr": lr.item(),
                            }
                            tqdm.write(str(__A     )     )
                            self.logger.log(__A,commit=__A     )

                      if i % args.save_steps == 0:
                            self.save_checkpoint(args.save_dir + f'-e{epoch}-s{i}',state=__A     )



    def   lowerCamelCase_  (	self      :       Union[str, Any],__A      :       int,__A      :       Tuple     ):
          _lowerCamelCase				:   List[Any]								=   get_batched_dataset(__A,self.args.batch_size     )
          _lowerCamelCase				:   str								=   len(__A     ) // self.args.batch_size
          _lowerCamelCase				:   Optional[Any]								=   jnp.array(0,dtype=jnp.floataa     )
          _lowerCamelCase				:   Dict								=   0
          for batch in tqdm(__A,total=__A,desc="Evaluating ... "     ):
                _lowerCamelCase				:   str								=   self.data_collator(__A     )
                _lowerCamelCase				:   Dict								=   self.val_step_fn(__A,**__A     )
                running_loss += jax_utils.unreplicate(metrics["loss"]     )
                i += 1
          return running_loss / i



    def   lowerCamelCase_  (	self      :       List[str],__A      :       Optional[Any],__A      :       Union[str, Any]     ):
          _lowerCamelCase				:   Any								=   jax_utils.unreplicate(__A     )
          print(f'SAVING CHECKPOINT IN {save_dir}',end=" ... "     )
          self.model_save_fn(__A,params=state.params     )
          with open(os.path.join(__A,"opt_state.msgpack"     ),"wb"     ) as f:
                f.write(to_bytes(state.opt_state     )     )
          joblib.dump(self.args,os.path.join(__A,"args.joblib"     )     )
          joblib.dump(self.data_collator,os.path.join(__A,"data_collator.joblib"     )     )
          with open(os.path.join(__A,"training_state.json"     ),"w"     ) as f:
                json.dump({"step": state.step.item()},__A     )
          print("DONE"     )







def 		A_       (	_lowerCAmelCase   :    Optional[Any]				,	_lowerCAmelCase   :    List[str]							):



      """simple docstring"""




      print(F'RESTORING CHECKPOINT FROM {save_dir}'				,	end=" ... "							)
      with open(os.path.join(_lowerCAmelCase				,	"flax_model.msgpack"							)				,	"rb"							) as f:
            _lowerCamelCase				:   str								=   from_bytes(state.params				,	f.read()							)

      with open(os.path.join(_lowerCAmelCase				,	"opt_state.msgpack"							)				,	"rb"							) as f:
            _lowerCamelCase				:   List[Any]								=   from_bytes(state.opt_state				,	f.read()							)

      _lowerCamelCase				:   Optional[int]								=   joblib.load(os.path.join(_lowerCAmelCase				,	"args.joblib"							)							)
      _lowerCamelCase				:   List[Any]								=   joblib.load(os.path.join(_lowerCAmelCase				,	"data_collator.joblib"							)							)

      with open(os.path.join(_lowerCAmelCase				,	"training_state.json"							)				,	"r"							) as f:
            _lowerCamelCase				:   str								=   json.load(_lowerCAmelCase							)
      _lowerCamelCase				:   Union[str, Any]								=   training_state["step"]

      print("DONE"							)
      return params, opt_state, step, args, data_collator







def 		A_       (	_lowerCAmelCase   :    Optional[int]				,	_lowerCAmelCase   :    int				,	_lowerCAmelCase   :    List[str]				,	_lowerCAmelCase   :    str							):



      """simple docstring"""




      _lowerCamelCase				:   List[str]								=   num_train_steps - warmup_steps
      _lowerCamelCase				:   Dict								=   optax.linear_schedule(init_value=_lowerCAmelCase				,	end_value=_lowerCAmelCase				,	transition_steps=_lowerCAmelCase							)
      _lowerCamelCase				:   Dict								=   optax.linear_schedule(init_value=_lowerCAmelCase				,	end_value=1E-7				,	transition_steps=_lowerCAmelCase							)
      _lowerCamelCase				:   List[Any]								=   optax.join_schedules(schedules=[warmup_fn, decay_fn]				,	boundaries=[warmup_steps]							)
      return lr







def 		A_       (	_lowerCAmelCase   :    Dict				,	_lowerCAmelCase   :    int				,	_lowerCAmelCase   :    Optional[int]				,	_lowerCAmelCase   :    Dict				,	_lowerCAmelCase   :    Optional[int]							):



      """simple docstring"""




      def weight_decay_mask(_lowerCAmelCase   :    List[Any]							):
            _lowerCamelCase				:   Tuple								=   traverse_util.flatten_dict(_lowerCAmelCase							)
            _lowerCamelCase				:   Tuple								=   {k: (v[-1] != "bias" and v[-2:] != ("LayerNorm", "scale")) for k, v in params.items()}
            return traverse_util.unflatten_dict(_lowerCAmelCase							)

      _lowerCamelCase				:   int								=   scheduler_fn(_lowerCAmelCase				,	_lowerCAmelCase				,	_lowerCAmelCase				,	_lowerCAmelCase							)

      _lowerCamelCase				:   Optional[Any]								=   optax.adamw(learning_rate=_lowerCAmelCase				,	weight_decay=_lowerCAmelCase				,	mask=_lowerCAmelCase							)
      return tx, lr

'''simple docstring'''






import os
import re
import unicodedata
from shutil import copyfile
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import is_torch_available, logging


if is_torch_available():
 import torch


if TYPE_CHECKING:
 from transformers.pipelines.conversational import Conversation


UpperCAmelCase_       :   str   =       logging.get_logger(__name__)
UpperCAmelCase_       :   Optional[int]   =       {'vocab_file': 'spiece.model'}

UpperCAmelCase_       :   Any   =       {
    'vocab_file': {
        'AI-Sweden/gpt-sw3-126m': 'https://huggingface.co/AI-Sweden/gpt-sw3-126m/resolve/main/spiece.model',
        'AI-Sweden/gpt-sw3-350m': 'https://huggingface.co/AI-Sweden/gpt-sw3-350m/resolve/main/spiece.model',
        'AI-Sweden/gpt-sw3-1.6b': 'https://huggingface.co/AI-Sweden/gpt-sw3-1.6b/resolve/main/spiece.model',
        'AI-Sweden/gpt-sw3-6.7b': 'https://huggingface.co/AI-Sweden/gpt-sw3-6.7b/resolve/main/spiece.model',
        'AI-Sweden/gpt-sw3-20b': 'https://huggingface.co/AI-Sweden/gpt-sw3-20b/resolve/main/spiece.model',
    }
}

UpperCAmelCase_       :   str   =       {
    'AI-Sweden/gpt-sw3-126m': 2048,
    'AI-Sweden/gpt-sw3-350m': 2048,
    'AI-Sweden/gpt-sw3-1.6b': 2048,
    'AI-Sweden/gpt-sw3-6.7b': 2048,
    'AI-Sweden/gpt-sw3-20b': 2048,
}

class  UpperCAmelCase__						(       A ):
    lowerCAmelCase_				    =		VOCAB_FILES_NAMES
    lowerCAmelCase_				    =		PRETRAINED_VOCAB_FILES_MAP
    lowerCAmelCase_				    =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    lowerCAmelCase_				    =		['input_ids', 'attention_mask']



    def __init__(	self      :       Dict,__A      :       List[str],__A      :       Any=False,__A      :       Tuple=False,__A      :       Dict=False,__A      :       str=None,__A      :       List[str]=None,__A      :       Any=None,__A      :       str=None,__A      :       Optional[Dict[str, Any]] = None,**__A      :       str,):
          _lowerCamelCase				:   List[Any]								=   {} if sp_model_kwargs is None else sp_model_kwargs

          _lowerCamelCase				:   int								=   kwargs.get("name_or_path"     )
          if name_or_path is None:
                logger.warning(
                    "name_or_path not provided, will work for all GPTSw3 models except gpt-sw3-7b,"
                    " you are testing the model, this can safely be ignored"     )
                _lowerCamelCase				:   Union[str, Any]								=   "None"

          # Default definitions for our 2 tokenizer versions, with None-checks to enable proper testing
          _lowerCamelCase				:   Tuple								=   "<|endoftext|>" if eos_token is None else eos_token
          _lowerCamelCase				:   List[str]								=   "<unk>" if unk_token is None else unk_token
          if "gpt-sw3-7b" in name_or_path:
                _lowerCamelCase				:   Union[str, Any]								=   unk_token if pad_token is None else pad_token
                _lowerCamelCase				:   str								=   eos_token if bos_token is None else bos_token
          else:
                _lowerCamelCase				:   List[str]								=   "<pad>" if pad_token is None else pad_token
                _lowerCamelCase				:   str								=   "<s>" if bos_token is None else bos_token

          super().__init__(
              do_lower_case=__A,remove_space=__A,keep_accents=__A,bos_token=__A,eos_token=__A,unk_token=__A,pad_token=__A,sp_model_kwargs=self.sp_model_kwargs,**__A,)

          _lowerCamelCase				:   Union[str, Any]								=   do_lower_case
          _lowerCamelCase				:   List[Any]								=   remove_space
          _lowerCamelCase				:   str								=   keep_accents
          _lowerCamelCase				:   List[Any]								=   vocab_file

          _lowerCamelCase				:   int								=   spm.SentencePieceProcessor(**self.sp_model_kwargs     )
          self.sp_model.Load(__A     )

          # Used for whitespace normalization in input texts
          # fmt : off
          _lowerCamelCase				:   Union[str, Any]								=   {" ", " ", " ", " ", " ", "　", " ", " ", " ", " ", "￼", ""}
          # fmt : on

          # Regular expression to remove non-printing characters (e.g. some unicode control chars) in preprocessing
          _lowerCamelCase				:   int								=   re.compile(
              f'[{"".join(map(__A,list(range(0,9     )     ) + list(range(1_1,3_2     )     ) + list(range(1_2_7,1_6_0     )     ) + [1_6_0, 1_7_3, 8_2_0_3]     )     )}]'     )



    def __getstate__(	self      :       Dict     ):
          _lowerCamelCase				:   int								=   self.__dict__.copy()
          _lowerCamelCase				:   Optional[Any]								=   None
          return state



    def __setstate__(	self      :       Tuple,__A      :       int     ):
          _lowerCamelCase				:   Optional[int]								=   d

          # for backward compatibility
          if not hasattr(self,"sp_model_kwargs"     ):
                _lowerCamelCase				:   List[str]								=   {}

          _lowerCamelCase				:   Tuple								=   spm.SentencePieceProcessor(**self.sp_model_kwargs     )
          self.sp_model.Load(self.vocab_file     )



    @property
    # Copied from transformers.models.albert.tokenization_albert.AlbertTokenizer.vocab_size
    def   lowerCamelCase_  (	self      :       int     ):
          return len(self.sp_model     )



    def   lowerCamelCase_  (	self      :       Dict,__A      :       str     ):
          _lowerCamelCase				:   Union[str, Any]								=   self.non_printing_characters_re.sub("",__A     )

          # Normalize whitespaces
          _lowerCamelCase				:   Optional[Any]								=   "".join([char if char not in self.whitespaces else " " for char in text]     )

          # NFC Unicode normalization
          _lowerCamelCase				:   Optional[Any]								=   unicodedata.normalize("NFC",__A     )
          return text



    def   lowerCamelCase_  (	self      :       Union[str, Any],__A      :       str,**__A      :       Optional[int]     ):
          _lowerCamelCase				:   str								=   self.preprocess_text(__A     )
          return self.sp_model.encode(__A,out_type=__A     )



    def   lowerCamelCase_  (	self      :       int,__A      :       str     ):
          return self.sp_model.PieceToId(__A     )



    def   lowerCamelCase_  (	self      :       Optional[int],__A      :       int     ):
          return self.sp_model.IdToPiece(__A     )



    @staticmethod
    def   lowerCamelCase_  (	__A      :       str     ):
          return out_string



    def   lowerCamelCase_  (	self      :       str,__A      :       List[str]     ):
          _lowerCamelCase				:   str								=   []
          _lowerCamelCase				:   List[Any]								=   ""
          _lowerCamelCase				:   Tuple								=   False
          for token in tokens:
                # make sure that special tokens are not decoded using sentencepiece model
                if token in self.all_special_tokens:
                      # TODO: Check if this is needed, as it ensures that decode(encode(doc)) != doc by adding extra whitespace in the decoded document
                      if not prev_is_special:
                            out_string += " "

                      out_string += self.sp_model.decode(__A     ) + token
                      _lowerCamelCase				:   Optional[int]								=   True
                      _lowerCamelCase				:   Optional[Any]								=   []
                else:
                      current_sub_tokens.append(__A     )
                      _lowerCamelCase				:   str								=   False
          out_string += self.sp_model.decode(__A     )

          return out_string



    def   lowerCamelCase_  (	self      :       Any     ):
          _lowerCamelCase				:   Optional[int]								=   {self.convert_ids_to_tokens(__A     ): i for i in range(self.vocab_size     )}
          vocab.update(self.added_tokens_encoder     )
          return vocab



    def   lowerCamelCase_  (	self      :       Optional[Any],__A      :       str,__A      :       Optional[str] = None     ):
          if not os.path.isdir(__A     ):
                logger.error(f'Vocabulary path ({save_directory}) should be a directory'     )
                return
          _lowerCamelCase				:   List[Any]								=   os.path.join(
              __A,(filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]     )

          if os.path.abspath(self.vocab_file     ) != os.path.abspath(__A     ) and os.path.isfile(self.vocab_file     ):
                copyfile(self.vocab_file,__A     )
          elif not os.path.isfile(self.vocab_file     ):
                with open(__A,"wb"     ) as fi:
                      _lowerCamelCase				:   str								=   self.sp_model.serialized_model_proto()
                      fi.write(__A     )

          return (out_vocab_file,)



    def   lowerCamelCase_  (	self      :       Optional[int],__A      :       Union[str, List[str]],__A      :       Union[str, bool] = False     ):

          if isinstance(__A,__A     ):
                _lowerCamelCase				:   List[Any]								=   self.preprocess_text(__A     )
                _lowerCamelCase				:   Optional[Any]								=   self.sp_model.encode(__A     )
          else:
                _lowerCamelCase				:   List[str]								=   [self.preprocess_text(__A     ) for t in text]
                _lowerCamelCase				:   int								=   self.sp_model.encode(__A     )

          if return_tensors is True or return_tensors == "pt":
                _lowerCamelCase				:   str								=   torch.tensor(__A     )

          return token_ids



    def   lowerCamelCase_  (	self      :       List[Any],__A      :       Union[int, List[int]]     ):
          return self.sp_model.decode(__A     )



    def   lowerCamelCase_  (	self      :       Optional[int],__A      :       "Conversation"     ):
          _lowerCamelCase				:   Any								=   [f'User: {text}' if is_user else f'Bot: {text}' for is_user, text in conversation.iter_texts()]
          _lowerCamelCase				:   Tuple								=   (
              f'{self.eos_token}{self.bos_token}' + f'{self.bos_token}'.join(__A     ) + f'{self.bos_token}Bot:'
          )
          return self.encode(text=__A     )