



'''simple docstring'''





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

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer


if TYPE_CHECKING:
 from ...tokenization_utils_base import TextInput
from ...utils import logging


_lowerCAmelCase      						=	logging.get_logger(__name__)

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

_lowerCAmelCase      						=	{
    '''vocab_file''': {
        '''t5-small''': '''https://huggingface.co/t5-small/resolve/main/spiece.model''',
        '''t5-base''': '''https://huggingface.co/t5-base/resolve/main/spiece.model''',
        '''t5-large''': '''https://huggingface.co/t5-large/resolve/main/spiece.model''',
        '''t5-3b''': '''https://huggingface.co/t5-3b/resolve/main/spiece.model''',
        '''t5-11b''': '''https://huggingface.co/t5-11b/resolve/main/spiece.model''',
    }
}


# TODO(PVP) - this should be removed in Transformers v5
_lowerCAmelCase      						=	{
    '''t5-small''': 512,
    '''t5-base''': 512,
    '''t5-large''': 512,
    '''t5-3b''': 512,
    '''t5-11b''': 512,
}

_lowerCAmelCase      						=	'''▁'''






class        lowerCAmelCase_(				A__     ):
  '''simple docstring'''




  __lowercase							:      Optional[Any]				=      VOCAB_FILES_NAMES
  __lowercase							:      str				=      PRETRAINED_VOCAB_FILES_MAP
  __lowercase							:      Any				=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  __lowercase							:      Optional[Any]				=      ['''input_ids''', '''attention_mask''']




  def __init__(    self  ,__UpperCAmelCase  ,__UpperCAmelCase="</s>"  ,__UpperCAmelCase="<unk>"  ,__UpperCAmelCase="<pad>"  ,__UpperCAmelCase=100  ,__UpperCAmelCase=None  ,__UpperCAmelCase = None  ,__UpperCAmelCase=True  ,**__UpperCAmelCase  ,)    ->		Any:






      if extra_ids > 0 and additional_special_tokens is None:
          lowerCAmelCase__ :  List[Any]							      =					[F"""<extra_id_{i}>""" for i in range(__UpperCAmelCase      )]
      elif extra_ids > 0 and additional_special_tokens is not None:
          # Check that we have the right number of extra_id special tokens
          lowerCAmelCase__ :  Dict							      =					len(set(filter(lambda __UpperCAmelCase      : bool("""extra_id""" in str(__UpperCAmelCase      )      )  ,__UpperCAmelCase      )      )      )
          if extra_tokens != extra_ids:
              raise ValueError(
                  F"""Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are"""
                  """ provided to T5Tokenizer. In this case the additional_special_tokens must include the extra_ids"""
                  """ tokens"""      )
      if legacy:
          logger.warning_once(
              F"""You are using the legacy behaviour of the {self.__class__}. This means that tokens that come after special tokens will not be properly handled. We recommend you to"""
              """ read the related pull request available at https://github.com/huggingface/transformers/pull/24565"""      )

      lowerCAmelCase__ :  List[Any]							      =					legacy
      lowerCAmelCase__ :  int							      =					{} if sp_model_kwargs is None else sp_model_kwargs

      super().__init__(
          eos_token=__UpperCAmelCase  ,unk_token=__UpperCAmelCase  ,pad_token=__UpperCAmelCase  ,extra_ids=__UpperCAmelCase  ,additional_special_tokens=__UpperCAmelCase  ,sp_model_kwargs=self.sp_model_kwargs  ,legacy=__UpperCAmelCase  ,**__UpperCAmelCase  ,)

      lowerCAmelCase__ :  List[str]							      =					vocab_file
      lowerCAmelCase__ :  Optional[Any]							      =					extra_ids

      lowerCAmelCase__ :  str							      =					spm.SentencePieceProcessor(**self.sp_model_kwargs      )
      self.sp_model.Load(__UpperCAmelCase      )




  @staticmethod
  def 							UpperCAmelCase_	(    __UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase      )    ->		Dict:






      if pretrained_model_name_or_path in TaTokenizer.max_model_input_sizes:
          lowerCAmelCase__ :  Tuple							      =					TaTokenizer.max_model_input_sizes[pretrained_model_name_or_path]
          if init_max_model_length is not None and init_max_model_length != max_model_length:
              return init_max_model_length
          elif init_max_model_length is None:
              warnings.warn(
                  """This tokenizer was incorrectly instantiated with a model max length of"""
                  F""" {deprecated_max_model_length} which will be corrected in Transformers v5.\nFor now, this"""
                  """ behavior is kept to avoid breaking backwards compatibility when padding/encoding with"""
                  """ `truncation is True`.\n- Be aware that you SHOULD NOT rely on"""
                  F""" {pretrained_model_name_or_path} automatically truncating your input to"""
                  F""" {deprecated_max_model_length} when padding/encoding.\n- If you want to encode/pad to sequences"""
                  F""" longer than {deprecated_max_model_length} you can either instantiate this tokenizer with"""
                  """ `model_max_length` or pass `max_length` when encoding/padding.\n- To avoid this warning, please"""
                  """ instantiate this tokenizer with `model_max_length` set to your preferred value."""  ,__UpperCAmelCase  ,)

      return max_model_length




  @property
  def 							UpperCAmelCase_	(    self      )    ->		List[str]:
      return self.sp_model.get_piece_size() + self._extra_ids




  def 							UpperCAmelCase_	(    self      )    ->		List[str]:
      lowerCAmelCase__ :  Optional[Any]							      =					{self.convert_ids_to_tokens(__UpperCAmelCase      ): i for i in range(self.vocab_size      )}
      vocab.update(self.added_tokens_encoder      )
      return vocab




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None  ,__UpperCAmelCase = False      )    ->		Any:






      if already_has_special_tokens:
          return super().get_special_tokens_mask(
              token_ids_a=__UpperCAmelCase  ,token_ids_a=__UpperCAmelCase  ,already_has_special_tokens=__UpperCAmelCase      )

      # normal case: some special tokens
      if token_ids_a is None:
          return ([0] * len(__UpperCAmelCase      )) + [1]
      return ([0] * len(__UpperCAmelCase      )) + [1] + ([0] * len(__UpperCAmelCase      )) + [1]




  def 							UpperCAmelCase_	(    self      )    ->		str:
      return list(
          set(filter(lambda __UpperCAmelCase      : bool(re.search(R"""<extra_id_\d+>"""  ,__UpperCAmelCase      )      ) is not None  ,self.additional_special_tokens      )      )      )




  def 							UpperCAmelCase_	(    self      )    ->		str:
      return [self._convert_token_to_id(__UpperCAmelCase      ) for token in self.get_sentinel_tokens()]




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Dict:






      if len(__UpperCAmelCase      ) > 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 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None      )    ->		Tuple:
      lowerCAmelCase__ :  Optional[int]							      =					[self.eos_token_id]

      if token_ids_a is None:
          return len(token_ids_a + eos      ) * [0]
      return len(token_ids_a + eos + token_ids_a + eos      ) * [0]




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None      )    ->		Union[str, Any]:
      lowerCAmelCase__ :  Dict							      =					self._add_eos_if_not_present(__UpperCAmelCase      )
      if token_ids_a is None:
          return token_ids_a
      else:
          lowerCAmelCase__ :  Optional[Any]							      =					self._add_eos_if_not_present(__UpperCAmelCase      )
          return token_ids_a + token_ids_a




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




  def __setstate__(    self  ,__UpperCAmelCase      )    ->		List[str]:
      lowerCAmelCase__ :  List[Any]							      =					d

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

      lowerCAmelCase__ :  int							      =					spm.SentencePieceProcessor(**self.sp_model_kwargs      )
      self.sp_model.Load(self.vocab_file      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,**__UpperCAmelCase      )    ->		List[str]:






      if not self.legacy:
          lowerCAmelCase__ :  Tuple							      =					SPIECE_UNDERLINE + text.replace(__UpperCAmelCase  ,""" """      )
      return super().tokenize(__UpperCAmelCase  ,**__UpperCAmelCase      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,**__UpperCAmelCase      )    ->		List[Any]:






      if not self.legacy:
          lowerCAmelCase__ :  List[str]							      =					text.startswith(__UpperCAmelCase      )
          if is_first:
              lowerCAmelCase__ :  List[str]							      =					text[1:]

      lowerCAmelCase__ :  Tuple							      =					self.sp_model.encode(__UpperCAmelCase  ,out_type=__UpperCAmelCase      )

      if not self.legacy and not is_first and not text.startswith(""" """      ) and tokens[0].startswith(__UpperCAmelCase      ):
          lowerCAmelCase__ :  Tuple							      =					([tokens[0][1:]] if len(tokens[0]      ) > 1 else []) + tokens[1:]
      return tokens




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Dict:






      if token.startswith("""<extra_id_"""      ):
          lowerCAmelCase__ :  str							      =					re.match(R"""<extra_id_(\d+)>"""  ,__UpperCAmelCase      )
          lowerCAmelCase__ :  List[str]							      =					int(match.group(1      )      )
          return self.vocab_size - num - 1
      return self.sp_model.piece_to_id(__UpperCAmelCase      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Union[str, Any]:






      if index < self.sp_model.get_piece_size():
          lowerCAmelCase__ :  Any							      =					self.sp_model.IdToPiece(__UpperCAmelCase      )
      else:
          lowerCAmelCase__ :  int							      =					F"""<extra_id_{self.vocab_size - 1 - index}>"""
      return token




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		int:
      lowerCAmelCase__ :  int							      =					[]
      lowerCAmelCase__ :  List[str]							      =					""
      lowerCAmelCase__ :  List[Any]							      =					False
      for token in tokens:
          # make sure that special tokens are not decoded using sentencepiece model
          if token in self.all_special_tokens:
              if not prev_is_special:
                  out_string += " "
              out_string += self.sp_model.decode(__UpperCAmelCase      ) + token
              lowerCAmelCase__ :  List[Any]							      =					True
              lowerCAmelCase__ :  Optional[Any]							      =					[]
          else:
              current_sub_tokens.append(__UpperCAmelCase      )
              lowerCAmelCase__ :  Optional[int]							      =					False
      out_string += self.sp_model.decode(__UpperCAmelCase      )
      return out_string.strip()






  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None      )    ->		Dict:






      if not os.path.isdir(__UpperCAmelCase      ):
          logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""      )
          return
      lowerCAmelCase__ :  Union[str, Any]							      =					os.path.join(
          __UpperCAmelCase  ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]      )

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

      return (out_vocab_file,)










import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam





class 			__a  (    datasets.BeamBasedBuilder       ):







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      return datasets.DatasetInfo(
          features=datasets.Features({"content": datasets.Value("string"						)}						)							,   supervised_keys=SCREAMING_SNAKE_CASE							,   )







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


      '''simple docstring'''






      return [datasets.SplitGenerator(name=datasets.Split.TRAIN							,   gen_kwargs={"examples": get_test_dummy_examples()}						)]







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   int						):


      '''simple docstring'''






      import apache_beam as beam

      return pipeline | "Load Examples" >> beam.Create(SCREAMING_SNAKE_CASE						)







class 			__a  (    datasets.BeamBasedBuilder       ):







 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






      return datasets.DatasetInfo(
          features=datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"						)}						)}						)							,   supervised_keys=SCREAMING_SNAKE_CASE							,   )







 def 					__lowercase					(		self					:   Any							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   List[Any]						):


      '''simple docstring'''






      return [
          datasets.SplitGenerator(name=datasets.Split.TRAIN							,   gen_kwargs={"examples": get_test_nested_examples()}						)
      ]







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   Any						):


      '''simple docstring'''






      import apache_beam as beam

      return pipeline | "Load Examples" >> beam.Create(SCREAMING_SNAKE_CASE						)







def 		SCREAMING_SNAKE_CASE							(					) ->	List[str]:
     return [(i, {"content": content}) for i, content in enumerate(["foo", "bar", "foobar"]			)]







def 		SCREAMING_SNAKE_CASE							(					) ->	Optional[Any]:
     return [(i, {"a": {"b": [content]}}) for i, content in enumerate(["foo", "bar", "foobar"]			)]





class 			__a  (    A__       ):







 @require_beam
 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  len(get_test_dummy_examples()						)
      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   List[Any]       				=  DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE							,   beam_runner="DirectRunner"						)
           builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train.arrow'						)						)						)
           self.assertDictEqual(builder.info.features							,   datasets.Features({"content": datasets.Value("string"						)}						)						)
           UpperCamelCase__			:   Optional[Any]       				=  builder.as_dataset()
           self.assertEqual(dset["train"].num_rows							,   SCREAMING_SNAKE_CASE						)
           self.assertEqual(dset["train"].info.splits["train"].num_examples							,   SCREAMING_SNAKE_CASE						)
           self.assertDictEqual(dset["train"][0]							,   get_test_dummy_examples()[0][1]						)
           self.assertDictEqual(
               dset["train"][expected_num_examples - 1]							,   get_test_dummy_examples()[expected_num_examples - 1][1]						)
           self.assertTrue(
               os.path.exists(os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   "dataset_info.json"						)						)						)
           del dset







 @require_beam
 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






      import apache_beam as beam

      UpperCamelCase__			:   Any       				=  beam.io.parquetio.WriteToParquet

      UpperCamelCase__			:   Tuple       				=  len(get_test_dummy_examples()						)
      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   str       				=  DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE							,   beam_runner="DirectRunner"						)
           with patch("apache_beam.io.parquetio.WriteToParquet"						) as write_parquet_mock:
                UpperCamelCase__			:   Optional[Any]       				=  partial(SCREAMING_SNAKE_CASE							,   num_shards=2						)
                builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(
                       SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train-00000-of-00002.arrow'						)						)						)
           self.assertTrue(
               os.path.exists(
                   os.path.join(
                       SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train-00000-of-00002.arrow'						)						)						)
           self.assertDictEqual(builder.info.features							,   datasets.Features({"content": datasets.Value("string"						)}						)						)
           UpperCamelCase__			:   List[Any]       				=  builder.as_dataset()
           self.assertEqual(dset["train"].num_rows							,   SCREAMING_SNAKE_CASE						)
           self.assertEqual(dset["train"].info.splits["train"].num_examples							,   SCREAMING_SNAKE_CASE						)
           # Order is not preserved when sharding, so we just check that all the elements are there
           self.assertListEqual(sorted(dset["train"]["content"]						)							,   sorted(["foo", "bar", "foobar"]						)						)
           self.assertTrue(
               os.path.exists(os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   "dataset_info.json"						)						)						)
           del dset







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


      '''simple docstring'''






      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   int       				=  DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE						)
           self.assertRaises(datasets.builder.MissingBeamOptions							,   builder.download_and_prepare						)







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


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  len(get_test_nested_examples()						)
      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   Any       				=  NestedBeamDataset(cache_dir=SCREAMING_SNAKE_CASE							,   beam_runner="DirectRunner"						)
           builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train.arrow'						)						)						)
           self.assertDictEqual(
               builder.info.features							,   datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"						)}						)}						)						)
           UpperCamelCase__			:   List[Any]       				=  builder.as_dataset()
           self.assertEqual(dset["train"].num_rows							,   SCREAMING_SNAKE_CASE						)
           self.assertEqual(dset["train"].info.splits["train"].num_examples							,   SCREAMING_SNAKE_CASE						)
           self.assertDictEqual(dset["train"][0]							,   get_test_nested_examples()[0][1]						)
           self.assertDictEqual(
               dset["train"][expected_num_examples - 1]							,   get_test_nested_examples()[expected_num_examples - 1][1]						)
           self.assertTrue(
               os.path.exists(os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   "dataset_info.json"						)						)						)
           del dset