

"""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


__UpperCAmelCase =				'▁'

__UpperCAmelCase =				{'vocab_file': 'spiece.model'}

__UpperCAmelCase =				{
    'vocab_file': {'google/pegasus-xsum': 'https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model'}
}

__UpperCAmelCase =				{
    'google/pegasus-xsum': 512,
}


__UpperCAmelCase =				logging.get_logger(__name__)







class      __lowercase			(    __lowerCamelCase					):
  snake_case_				       =       VOCAB_FILES_NAMES

  snake_case_				       =       VOCAB_FILES_NAMES
  snake_case_				       =       PRETRAINED_VOCAB_FILES_MAP
  snake_case_				       =       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  snake_case_				       =       ["""input_ids""", """attention_mask"""]

  def __init__(  self		:			Optional[int]					,A		:			List[str]					,A		:			List[str]="<pad>"					,A		:			Tuple="</s>"					,A		:			Tuple="<unk>"					,A		:			str="<mask_2>"					,A		:			Dict="<mask_1>"					,A		:			Tuple=None					,A		:			Optional[int]=103					,A		:			Optional[Dict[str, Any]] = None					,**A		:			str					,):






   '''simple docstring'''
   UpperCAmelCase__					:  Tuple					   =					offset
   if additional_special_tokens is not None:
    if not isinstance(A					,A    ):
     raise TypeError(
         f"additional_special_tokens should be of type {type(A    )}, but is"
         f" {type(A    )}"    )

    UpperCAmelCase__					:  Optional[int]					   =					(
        ([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(A    )					,self.offset - 1    )
    ]

    if len(set(A    )    ) != len(A    ):
     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}."    )
    UpperCAmelCase__					:  Optional[Any]					   =					additional_special_tokens_extended
   else:
    UpperCAmelCase__					:  Optional[int]					   =					[mask_token_sent] if mask_token_sent is not None else []
    additional_special_tokens += [f"<unk_{i}>" for i in range(2					,self.offset    )]

   UpperCAmelCase__					:  Dict					   =					{} if sp_model_kwargs is None else sp_model_kwargs

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

   # add special tokens to encoder dict
   UpperCAmelCase__					:  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    )}    )

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

  @property
  def    __lowercase       (  self		:			List[Any]    ):






   '''simple docstring'''
   return len(self.sp_model    ) + self.offset

  def    __lowercase       (  self		:			Tuple    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Union[str, Any]					   =					{self.convert_ids_to_tokens(A    ): i for i in range(self.vocab_size    )}
   vocab.update(self.added_tokens_encoder    )
   return vocab

  def __getstate__(  self		:			List[str]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Tuple					   =					self.__dict__.copy()
   UpperCAmelCase__					:  Optional[int]					   =					None
   return state

  def __setstate__(  self		:			str					,A		:			Optional[int]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Dict					   =					d

   # for backward compatibility
   if not hasattr(self					,"""sp_model_kwargs"""    ):
    UpperCAmelCase__					:  Any					   =					{}

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

  def    __lowercase       (  self		:			Dict					,A		:			str    ):






   '''simple docstring'''
   return self.sp_model.encode(A					,out_type=A    )

  def    __lowercase       (  self		:			str					,A		:			str    ):






   '''simple docstring'''
   if token in self.decoder:
    return self.decoder[token]
   elif token in self.added_tokens_decoder:
    return self.added_tokens_decoder[token]
   UpperCAmelCase__					:  str					   =					self.sp_model.piece_to_id(A    )
   return sp_id + self.offset

  def    __lowercase       (  self		:			Optional[Any]					,A		:			int    ):






   '''simple docstring'''
   if index in self.encoder:
    return self.encoder[index]
   elif index in self.added_tokens_encoder:
    return self.added_tokens_encoder[index]
   else:
    UpperCAmelCase__					:  int					   =					self.sp_model.IdToPiece(index - self.offset    )
   return token

  def    __lowercase       (  self		:			List[str]					,A		:			Optional[int]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					[]
   UpperCAmelCase__					:  Optional[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(A    ) + token
     UpperCAmelCase__					:  Optional[Any]					   =					[]
    else:
     current_sub_tokens.append(A    )
   out_string += self.sp_model.decode(A    )
   return out_string.strip()

  def    __lowercase       (  self		:			Tuple					,A		:			List[Any]=False    ):






   '''simple docstring'''
   return 1

  def    __lowercase       (  self		:			Dict					,A		:			List[str]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[int]					   =					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		:			Union[str, Any]					,A		:			List					,A		:			Optional[List] = None					,A		:			bool = False    ):






   '''simple docstring'''
   if already_has_special_tokens:
    return self._special_token_mask(A    )
   elif token_ids_a is None:
    return self._special_token_mask(A    ) + [1]
   else:
    return self._special_token_mask(token_ids_a + token_ids_a    ) + [1]

  def    __lowercase       (  self		:			Dict					,A		:			List[str]					,A		:			List[Any]=None    ):






   '''simple docstring'''
   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		:			Union[str, Any]					,A		:			str					,A		:			Optional[str] = None    ):






   '''simple docstring'''
   if not os.path.isdir(A    ):
    logger.error(f"Vocabulary path ({save_directory}) should be a directory"    )
    return
   UpperCAmelCase__					:  Optional[int]					   =					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:
     UpperCAmelCase__					:  Optional[Any]					   =					self.sp_model.serialized_model_proto()
     fi.write(A    )

   return (out_vocab_file,)








"""simple docstring"""




from typing import List, Optional, TypeVar

from .arrow_dataset import Dataset, _concatenate_map_style_datasets, _interleave_map_style_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .info import DatasetInfo
from .iterable_dataset import IterableDataset, _concatenate_iterable_datasets, _interleave_iterable_datasets
from .splits import NamedSplit
from .utils import logging
from .utils.py_utils import Literal


__UpperCAmelCase =				logging.get_logger(__name__)


__UpperCAmelCase =				TypeVar('DatasetType', Dataset, IterableDataset)
def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = "first_exhausted"							,    ):
 '''simple docstring'''





 from .arrow_dataset import Dataset
 from .iterable_dataset import IterableDataset

 if not datasets:
  raise ValueError("""Unable to interleave an empty list of datasets."""      )
 for i, dataset in enumerate(__UpperCamelCase      ):
  if not isinstance(__UpperCamelCase							,    (Dataset, IterableDataset)      ):
   if isinstance(__UpperCamelCase							,    (DatasetDict, IterableDatasetDict)      ):
    if not dataset:
     raise ValueError(
         F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
         """is an empty dataset dictionary."""      )
    raise ValueError(
        F"Dataset at position {i} has at least one split: {list(__UpperCamelCase      )}\n"
        F"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(__UpperCamelCase      )      )}']"      )
   raise ValueError(
       F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(__UpperCamelCase      ).__name__}."      )
  if i == 0:
   UpperCAmelCase__   ,      UpperCAmelCase__					:  List[str]					   =					(
       (Dataset, IterableDataset) if isinstance(__UpperCamelCase							,    __UpperCamelCase      ) else (IterableDataset, Dataset)
   )
  elif not isinstance(__UpperCamelCase							,    __UpperCamelCase      ):
   raise ValueError(
       F"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."      )
 if stopping_strategy not in ["first_exhausted", "all_exhausted"]:
  raise ValueError(F"{stopping_strategy} is not supported. Please enter a valid stopping_strategy."      )
 if dataset_type is Dataset:
  return _interleave_map_style_datasets(
      __UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    stopping_strategy=__UpperCamelCase      )
 else:
  return _interleave_iterable_datasets(
      __UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    stopping_strategy=__UpperCamelCase      )




def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = 0							,    ):
 '''simple docstring'''





 if not dsets:
  raise ValueError("""Unable to concatenate an empty list of datasets."""      )
 for i, dataset in enumerate(__UpperCamelCase      ):
  if not isinstance(__UpperCamelCase							,    (Dataset, IterableDataset)      ):
   if isinstance(__UpperCamelCase							,    (DatasetDict, IterableDatasetDict)      ):
    if not dataset:
     raise ValueError(
         F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
         """is an empty dataset dictionary."""      )
    raise ValueError(
        F"Dataset at position {i} has at least one split: {list(__UpperCamelCase      )}\n"
        F"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(__UpperCamelCase      )      )}']"      )
   raise ValueError(
       F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(__UpperCamelCase      ).__name__}."      )
  if i == 0:
   UpperCAmelCase__   ,      UpperCAmelCase__					:  Optional[int]					   =					(
       (Dataset, IterableDataset) if isinstance(__UpperCamelCase							,    __UpperCamelCase      ) else (IterableDataset, Dataset)
   )
  elif not isinstance(__UpperCamelCase							,    __UpperCamelCase      ):
   raise ValueError(
       F"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."      )
 if dataset_type is Dataset:
  return _concatenate_map_style_datasets(__UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    axis=__UpperCamelCase      )
 else:
  return _concatenate_iterable_datasets(__UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    axis=__UpperCamelCase      )





