


'''simple docstring'''



import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


A__		:					int												=  """▁"""

A__		:					Any												=  {"""vocab_file""": """spiece.model"""}

A__		:					Dict												=  {
    """vocab_file""": {"""google/pegasus-xsum""": """https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model"""}
}

A__		:					Union[str, Any]												=  {
    """google/pegasus-xsum""": 512,
}


A__		:					Union[str, Any]												=  logging.get_logger(__name__)





class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    lowerCamelCase  :					Any							     =			VOCAB_FILES_NAMES

    lowerCamelCase  :					List[Any]							     =			VOCAB_FILES_NAMES
    lowerCamelCase  :					Union[str, Any]							     =			PRETRAINED_VOCAB_FILES_MAP
    lowerCamelCase  :					Any							     =			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    lowerCamelCase  :					Tuple							     =			['input_ids', 'attention_mask']







    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_="<pad>"			,     SCREAMING_SNAKE_CASE_="</s>"			,     SCREAMING_SNAKE_CASE_="<unk>"			,     SCREAMING_SNAKE_CASE_="<mask_2>"			,     SCREAMING_SNAKE_CASE_="<mask_1>"			,     SCREAMING_SNAKE_CASE_=None			,     SCREAMING_SNAKE_CASE_=1_03			,     SCREAMING_SNAKE_CASE_ = None			,     **SCREAMING_SNAKE_CASE_			,     )				->     None:
         __lowerCamelCase   :	Dict								=		offset
         if additional_special_tokens is not None:
              if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	):
                   raise TypeError(
                       f'additional_special_tokens should be of type {type(SCREAMING_SNAKE_CASE_	)}, but is'
                       f' {type(SCREAMING_SNAKE_CASE_	)}'	)

              __lowerCamelCase   :	Dict								=		(
                  ([mask_token_sent] + additional_special_tokens)
                  if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
                  else additional_special_tokens
              )
              # fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
              additional_special_tokens_extended += [
                  f'<unk_{i}>' for i in range(len(SCREAMING_SNAKE_CASE_	)			,     self.offset - 1	)
              ]

              if len(set(SCREAMING_SNAKE_CASE_	)	) != len(SCREAMING_SNAKE_CASE_	):
                   raise ValueError(
                       'Please make sure that the provided additional_special_tokens do not contain an incorrectly'
                       f' shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}.'	)
              __lowerCamelCase   :	str								=		additional_special_tokens_extended
         else:
              __lowerCamelCase   :	List[str]								=		[mask_token_sent] if mask_token_sent is not None else []
              additional_special_tokens += [f'<unk_{i}>' for i in range(2			,     self.offset	)]

         __lowerCamelCase   :	int								=		{} if sp_model_kwargs is None else sp_model_kwargs

         super().__init__(
             eos_token=SCREAMING_SNAKE_CASE_			,     unk_token=SCREAMING_SNAKE_CASE_			,     mask_token=SCREAMING_SNAKE_CASE_			,     pad_token=SCREAMING_SNAKE_CASE_			,     mask_token_sent=SCREAMING_SNAKE_CASE_			,     offset=SCREAMING_SNAKE_CASE_			,     additional_special_tokens=SCREAMING_SNAKE_CASE_			,     sp_model_kwargs=self.sp_model_kwargs			,     **SCREAMING_SNAKE_CASE_			,     )
         __lowerCamelCase   :	Optional[Any]								=		mask_token_sent
         __lowerCamelCase   :	Optional[int]								=		vocab_file
         __lowerCamelCase   :	Tuple								=		spm.SentencePieceProcessor(**self.sp_model_kwargs	)
         self.sp_model.Load(SCREAMING_SNAKE_CASE_	)

         # add special tokens to encoder dict
         __lowerCamelCase   :	Dict[int, str]								=		{
             0: self.pad_token,
             1: self.eos_token,
         }

         if self.mask_token_sent is not None:
              self.encoder.update(
                  {
                      2: self.mask_token_sent,
                      3: self.mask_token,
                  }	)

         if self.offset > 0:
              # entries 2-104 are only used for pretraining and called <mask_1>, <mask_2>, unk_2, ...unk_102
              # mask_token_sent is already added to list -> so start at 1
              self.encoder.update({i + 3: additional_special_tokens[i] for i in range(1			,     self.offset - 1	)}	)

         __lowerCamelCase   :	Dict[str, int]								=		{v: k for k, v in self.encoder.items()}







    @property
    def       lowercase_				(							self	)				->     int:
         return len(self.sp_model	) + self.offset







    def       lowercase_				(							self	)				->     Dict[str, int]:
         __lowerCamelCase   :	Any								=		{self.convert_ids_to_tokens(SCREAMING_SNAKE_CASE_	): i for i in range(self.vocab_size	)}
         vocab.update(self.added_tokens_encoder	)
         return vocab







    def __getstate__(							self	)				->     List[Any]:
         __lowerCamelCase   :	List[str]								=		self.__dict__.copy()
         __lowerCamelCase   :	Optional[int]								=		None
         return state







    def __setstate__(							self			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase   :	List[Any]								=		d

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

         __lowerCamelCase   :	str								=		spm.SentencePieceProcessor(**self.sp_model_kwargs	)
         self.sp_model.Load(self.vocab_file	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     List[str]:
         return self.sp_model.encode(SCREAMING_SNAKE_CASE_			,     out_type=SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     int:
         if token in self.decoder:
              return self.decoder[token]
         elif token in self.added_tokens_decoder:
              return self.added_tokens_decoder[token]
         __lowerCamelCase   :	Tuple								=		self.sp_model.piece_to_id(SCREAMING_SNAKE_CASE_	)
         return sp_id + self.offset







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     str:
         if index in self.encoder:
              return self.encoder[index]
         elif index in self.added_tokens_encoder:
              return self.added_tokens_encoder[index]
         else:
              __lowerCamelCase   :	List[Any]								=		self.sp_model.IdToPiece(index - self.offset	)
         return token







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     Optional[Any]:
         __lowerCamelCase   :	Union[str, Any]								=		[]
         __lowerCamelCase   :	int								=		''
         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(SCREAMING_SNAKE_CASE_	) + token
                   __lowerCamelCase   :	Optional[Any]								=		[]
              else:
                   current_sub_tokens.append(SCREAMING_SNAKE_CASE_	)
         out_string += self.sp_model.decode(SCREAMING_SNAKE_CASE_	)
         return out_string.strip()







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=False	)				->     str:
         return 1







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     int:
         __lowerCamelCase   :	Union[str, Any]								=		set(self.all_special_ids	)  # call it once instead of inside list comp
         all_special_ids.remove(self.unk_token_id	)  # <unk> is only sometimes special

         return [1 if x in all_special_ids else 0 for x in seq]







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = None			,     SCREAMING_SNAKE_CASE_ = False	)				->     List[int]:
         if already_has_special_tokens:
              return self._special_token_mask(SCREAMING_SNAKE_CASE_	)
         elif token_ids_a is None:
              return self._special_token_mask(SCREAMING_SNAKE_CASE_	) + [1]
         else:
              return self._special_token_mask(token_ids_a + token_ids_a	) + [1]







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=None	)				->     List[int]:
         if token_ids_a is None:
              return token_ids_a + [self.eos_token_id]
         # We don't expect to process pairs, but leave the pair logic for API consistency
         return token_ids_a + token_ids_a + [self.eos_token_id]







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = None	)				->     Tuple[str]:
         if not os.path.isdir(SCREAMING_SNAKE_CASE_	):
              logger.error(f'Vocabulary path ({save_directory}) should be a directory'	)
              return
         __lowerCamelCase   :	List[str]								=		os.path.join(
             SCREAMING_SNAKE_CASE_			,     (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']	)

         if os.path.abspath(self.vocab_file	) != os.path.abspath(SCREAMING_SNAKE_CASE_	) and os.path.isfile(self.vocab_file	):
              copyfile(self.vocab_file			,     SCREAMING_SNAKE_CASE_	)
         elif not os.path.isfile(self.vocab_file	):
              with open(SCREAMING_SNAKE_CASE_			,     'wb'	) as fi:
                   __lowerCamelCase   :	int								=		self.sp_model.serialized_model_proto()
                   fi.write(SCREAMING_SNAKE_CASE_	)

         return (out_vocab_file,)










'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A__		:					List[Any]												=  {
    """configuration_megatron_bert""": ["""MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MegatronBertConfig"""],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A__		:					int												=  [
          """MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """MegatronBertForCausalLM""",
          """MegatronBertForMaskedLM""",
          """MegatronBertForMultipleChoice""",
          """MegatronBertForNextSentencePrediction""",
          """MegatronBertForPreTraining""",
          """MegatronBertForQuestionAnswering""",
          """MegatronBertForSequenceClassification""",
          """MegatronBertForTokenClassification""",
          """MegatronBertModel""",
          """MegatronBertPreTrainedModel""",
      ]

if TYPE_CHECKING:
      from .configuration_megatron_bert import MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, MegatronBertConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_megatron_bert import (
                MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                MegatronBertForCausalLM,
                MegatronBertForMaskedLM,
                MegatronBertForMultipleChoice,
                MegatronBertForNextSentencePrediction,
                MegatronBertForPreTraining,
                MegatronBertForQuestionAnswering,
                MegatronBertForSequenceClassification,
                MegatronBertForTokenClassification,
                MegatronBertModel,
                MegatronBertPreTrainedModel,
            )

else:
      import sys

      A__		:					List[Any]												=  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)






