






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


UpperCamelCase    :      Optional[Any]  			=		logging.get_logger(__name__)

UpperCamelCase    :      Tuple  			=		{'vocab_file': 'spiece.model'}

UpperCamelCase    :      Dict  			=		{
    '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
UpperCamelCase    :      Union[str, Any]  			=		{
    't5-small': 5_12,
    't5-base': 5_12,
    't5-large': 5_12,
    't5-3b': 5_12,
    't5-11b': 5_12,
}

UpperCamelCase    :      Union[str, Any]  			=		'▁'




class 		UpperCamelCase__    (a			):




		'''simple docstring'''





		_UpperCamelCase				  =					VOCAB_FILES_NAMES
		_UpperCamelCase				  =					PRETRAINED_VOCAB_FILES_MAP
		_UpperCamelCase				  =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
		_UpperCamelCase				  =					['input_ids', 'attention_mask']






		def __init__(      self   ,_lowerCAmelCase   ,_lowerCAmelCase="</s>"   ,_lowerCAmelCase="<unk>"   ,_lowerCAmelCase="<pad>"   ,_lowerCAmelCase=1_00   ,_lowerCAmelCase=None   ,_lowerCAmelCase = None   ,_lowerCAmelCase=True   ,**_lowerCAmelCase   ,):
								# Add extra_ids to the special token list
								if extra_ids > 0 and additional_special_tokens is None:
														lowerCamelCase__         =  [F'''<extra_id_{i}>''' for i in range(_lowerCAmelCase       )]
								elif extra_ids > 0 and additional_special_tokens is not None:
														# Check that we have the right number of extra_id special tokens
														lowerCamelCase__         =  len(set(filter(lambda _lowerCAmelCase       : bool("""extra_id""" in str(_lowerCAmelCase       )       )   ,_lowerCAmelCase       )       )       )
														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__         =  legacy
								lowerCamelCase__         =  {} if sp_model_kwargs is None else sp_model_kwargs

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

								lowerCamelCase__         =  vocab_file
								lowerCamelCase__         =  extra_ids

								lowerCamelCase__         =  spm.SentencePieceProcessor(**self.sp_model_kwargs       )
								self.sp_model.Load(_lowerCAmelCase       )






		@staticmethod
		def 			UpperCamelCase_					(      _lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase       ):
								if pretrained_model_name_or_path in TaTokenizer.max_model_input_sizes:
														lowerCamelCase__         =  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."""   ,_lowerCAmelCase   ,)

								return max_model_length






		@property
		def 			UpperCamelCase_					(      self       ):
								return self.sp_model.get_piece_size() + self._extra_ids






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  {self.convert_ids_to_tokens(_lowerCAmelCase       ): i for i in range(self.vocab_size       )}
								vocab.update(self.added_tokens_encoder       )
								return vocab






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase   ,_lowerCAmelCase = None   ,_lowerCAmelCase = False       ):
								if already_has_special_tokens:
														return super().get_special_tokens_mask(
														    token_ids_a=_lowerCAmelCase   ,token_ids_a=_lowerCAmelCase   ,already_has_special_tokens=_lowerCAmelCase       )

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






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






		def 			UpperCamelCase_					(      self       ):
								return [self._convert_token_to_id(_lowerCAmelCase       ) for token in self.get_sentinel_tokens()]






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase       ):
								if len(_lowerCAmelCase       ) > 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   ,_lowerCAmelCase   ,_lowerCAmelCase = None       ):
								lowerCamelCase__         =  [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   ,_lowerCAmelCase   ,_lowerCAmelCase = None       ):
								lowerCamelCase__         =  self._add_eos_if_not_present(_lowerCAmelCase       )
								if token_ids_a is None:
														return token_ids_a
								else:
														lowerCamelCase__         =  self._add_eos_if_not_present(_lowerCAmelCase       )
														return token_ids_a + token_ids_a






		def __getstate__(      self       ):
								lowerCamelCase__         =  self.__dict__.copy()
								lowerCamelCase__         =  None
								return state






		def __setstate__(      self   ,_lowerCAmelCase       ):
								lowerCamelCase__         =  d

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

								lowerCamelCase__         =  spm.SentencePieceProcessor(**self.sp_model_kwargs       )
								self.sp_model.Load(self.vocab_file       )






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase   ,**_lowerCAmelCase       ):
								# Replace the SPIECE_UNDERLINE with a space to make sure SPIECE_UNDERLINE is only used at
								# the beginning of the text
								if not self.legacy:
														lowerCamelCase__         =  SPIECE_UNDERLINE + text.replace(_lowerCAmelCase   ,""" """       )
								return super().tokenize(_lowerCAmelCase   ,**_lowerCAmelCase       )






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase   ,**_lowerCAmelCase       ):
								if not self.legacy:
														lowerCamelCase__         =  text.startswith(_lowerCAmelCase       )
														if is_first:
																				lowerCamelCase__         =  text[1:]

								lowerCamelCase__         =  self.sp_model.encode(_lowerCAmelCase   ,out_type=_lowerCAmelCase       )

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






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase       ):
								if token.startswith("""<extra_id_"""       ):
														lowerCamelCase__         =  re.match(R"""<extra_id_(\d+)>"""   ,_lowerCAmelCase       )
														lowerCamelCase__         =  int(match.group(1       )       )
														return self.vocab_size - num - 1
								return self.sp_model.piece_to_id(_lowerCAmelCase       )






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase       ):
								if index < self.sp_model.get_piece_size():
														lowerCamelCase__         =  self.sp_model.IdToPiece(_lowerCAmelCase       )
								else:
														lowerCamelCase__         =  F'''<extra_id_{self.vocab_size - 1 - index}>'''
								return token






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase       ):
								lowerCamelCase__         =  []
								lowerCamelCase__         =  """"""
								lowerCamelCase__         =  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(_lowerCAmelCase       ) + token
																				lowerCamelCase__         =  True
																				lowerCamelCase__         =  []
														else:
																				current_sub_tokens.append(_lowerCAmelCase       )
																				lowerCamelCase__         =  False
								out_string += self.sp_model.decode(_lowerCAmelCase       )
								return out_string.strip()






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase   ,_lowerCAmelCase = None       ):
								if not os.path.isdir(_lowerCAmelCase       ):
														logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''       )
														return
								lowerCamelCase__         =  os.path.join(
								    _lowerCAmelCase   ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]       )

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

								return (out_vocab_file,)









'''simple docstring'''







import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


UpperCamelCase    :      Union[str, Any]  			=		logging.get_logger(__name__)
def 						A__			(				__lowerCAmelCase		:			int						):
						lowerCamelCase__         =  R"""\w+[.]\d+"""
						lowerCamelCase__         =  re.findall(__lowerCAmelCase   ,    __lowerCAmelCase						)
						for pat in pats:
												lowerCamelCase__         =  key.replace(__lowerCAmelCase   ,    """_""".join(pat.split("""."""						)						)						)
						return key
def 						A__			(				__lowerCAmelCase		:			Union[str, Any]   ,    __lowerCAmelCase		:			Union[str, Any]   ,    __lowerCAmelCase		:			Tuple						):
						lowerCamelCase__         =  pt_tuple_key[:-1] + ("""scale""",)
						if (
						    any("""norm""" in str_ for str_ in pt_tuple_key						)
						    and (pt_tuple_key[-1] == "bias")
						    and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
						    and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
						):
												lowerCamelCase__         =  pt_tuple_key[:-1] + ("""scale""",)
												return renamed_pt_tuple_key, pt_tensor
						elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
												lowerCamelCase__         =  pt_tuple_key[:-1] + ("""scale""",)
												return renamed_pt_tuple_key, pt_tensor

						# embedding
						if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
												lowerCamelCase__         =  pt_tuple_key[:-1] + ("""embedding""",)
												return renamed_pt_tuple_key, pt_tensor

						# conv layer
						lowerCamelCase__         =  pt_tuple_key[:-1] + ("""kernel""",)
						if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
												lowerCamelCase__         =  pt_tensor.transpose(2   ,    3   ,    1   ,    0						)
												return renamed_pt_tuple_key, pt_tensor

						# linear layer
						lowerCamelCase__         =  pt_tuple_key[:-1] + ("""kernel""",)
						if pt_tuple_key[-1] == "weight":
												lowerCamelCase__         =  pt_tensor.T
												return renamed_pt_tuple_key, pt_tensor

						# old PyTorch layer norm weight
						lowerCamelCase__         =  pt_tuple_key[:-1] + ("""weight""",)
						if pt_tuple_key[-1] == "gamma":
												return renamed_pt_tuple_key, pt_tensor

						# old PyTorch layer norm bias
						lowerCamelCase__         =  pt_tuple_key[:-1] + ("""bias""",)
						if pt_tuple_key[-1] == "beta":
												return renamed_pt_tuple_key, pt_tensor

						return pt_tuple_key, pt_tensor





def 						A__			(				__lowerCAmelCase		:			int   ,    __lowerCAmelCase		:			Union[str, Any]   ,    __lowerCAmelCase		:			Tuple=42						):
						# Step 1: Convert pytorch tensor to numpy
						lowerCamelCase__         =  {k: v.numpy() for k, v in pt_state_dict.items()}

						# Step 2: Since the model is stateless, get random Flax params
						lowerCamelCase__         =  flax_model.init_weights(PRNGKey(__lowerCAmelCase						)						)

						lowerCamelCase__         =  flatten_dict(__lowerCAmelCase						)
						lowerCamelCase__         =  {}

						# Need to change some parameters name to match Flax names
						for pt_key, pt_tensor in pt_state_dict.items():
												lowerCamelCase__         =  rename_key(__lowerCAmelCase						)
												lowerCamelCase__         =  tuple(renamed_pt_key.split("""."""						)						)

												# Correctly rename weight parameters
												lowerCamelCase__       ,			lowerCamelCase__         =  rename_key_and_reshape_tensor(__lowerCAmelCase   ,    __lowerCAmelCase   ,    __lowerCAmelCase						)

												if flax_key in random_flax_state_dict:
																		if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																								raise ValueError(
																								    F'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
																								    F'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''						)

        # also add unexpected weight so that warning is thrown
												lowerCamelCase__         =  jnp.asarray(__lowerCAmelCase						)

						return unflatten_dict(__lowerCAmelCase						)

