



'''simple docstring'''




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

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import SPIECE_UNDERLINE, logging


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

a   :			Dict					  =    {'vocab_file': 'spiece.model'}

a   :			List[Any]					  =    {
    'vocab_file': {
        'xlnet-base-cased': 'https://huggingface.co/xlnet-base-cased/resolve/main/spiece.model',
        'xlnet-large-cased': 'https://huggingface.co/xlnet-large-cased/resolve/main/spiece.model',
    }
}

a   :			Optional[Any]					  =    {
    'xlnet-base-cased': None,
    'xlnet-large-cased': None,
}

# Segments (not really needed)
a   :			Dict					  =    0
a   :			Dict					  =    1
a   :			str					  =    2
a   :			str					  =    3
a   :			int					  =    4

class 			UpperCamelCase_ (							_snake_case			):
	lowercase          = VOCAB_FILES_NAMES
	lowercase          = PRETRAINED_VOCAB_FILES_MAP
	lowercase          = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	lowercase          = """left"""





	def __init__(  self		,		A		,		A=False		,		A=True		,		A=False		,		A="<s>"		,		A="</s>"		,		A="<unk>"		,		A="<sep>"		,		A="<pad>"		,		A="<cls>"		,		A="<mask>"		,		A=["<eop>", "<eod>"]		,		A = None		,		**A		,		)			->			List[str]:
								# Mask token behave like a normal word, i.e. include the space before it
								UpperCAmelCase     :		Union[str, Any]							     =		AddedToken(__snake_case		,		lstrip=__snake_case		,		rstrip=__snake_case    ) if isinstance(__snake_case		,		__snake_case    ) else mask_token

								UpperCAmelCase     :		Tuple							     =		{} if sp_model_kwargs is None else sp_model_kwargs

								super().__init__(
								    do_lower_case=__snake_case		,		remove_space=__snake_case		,		keep_accents=__snake_case		,		bos_token=__snake_case		,		eos_token=__snake_case		,		unk_token=__snake_case		,		sep_token=__snake_case		,		pad_token=__snake_case		,		cls_token=__snake_case		,		mask_token=__snake_case		,		additional_special_tokens=__snake_case		,		sp_model_kwargs=self.sp_model_kwargs		,		**__snake_case		,		)

								UpperCAmelCase     :		Union[str, Any]							     =		3

								UpperCAmelCase     :		Dict							     =		do_lower_case
								UpperCAmelCase     :		Tuple							     =		remove_space
								UpperCAmelCase     :		List[str]							     =		keep_accents
								UpperCAmelCase     :		int							     =		vocab_file

								UpperCAmelCase     :		str							     =		spm.SentencePieceProcessor(**self.sp_model_kwargs    )
								self.sp_model.Load(__snake_case    )





	@property
	def  _lowercase(  self    )			->			Optional[int]:
								return len(self.sp_model    )





	def  _lowercase(  self    )			->			str:
								UpperCAmelCase     :		Optional[int]							     =		{self.convert_ids_to_tokens(__snake_case    ): i for i in range(self.vocab_size    )}
								vocab.update(self.added_tokens_encoder    )
								return vocab





	def __getstate__(  self    )			->			int:
								UpperCAmelCase     :		str							     =		self.__dict__.copy()
								UpperCAmelCase     :		Union[str, Any]							     =		None
								return state





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

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

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





	def  _lowercase(  self		,		A    )			->			Any:
								if self.remove_space:
															UpperCAmelCase     :		int							     =		""" """.join(inputs.strip().split()    )
								else:
															UpperCAmelCase     :		Optional[Any]							     =		inputs
								UpperCAmelCase     :		List[Any]							     =		outputs.replace("""``"""		,		"""\""""    ).replace("""''"""		,		"""\""""    )

								if not self.keep_accents:
															UpperCAmelCase     :		Any							     =		unicodedata.normalize("""NFKD"""		,		__snake_case    )
															UpperCAmelCase     :		Tuple							     =		"""""".join([c for c in outputs if not unicodedata.combining(__snake_case    )]    )
								if self.do_lower_case:
															UpperCAmelCase     :		Dict							     =		outputs.lower()

								return outputs





	def  _lowercase(  self		,		A    )			->			Union[str, Any]:
								UpperCAmelCase     :		Optional[Any]							     =		self.preprocess_text(__snake_case    )
								UpperCAmelCase     :		str							     =		self.sp_model.encode(__snake_case		,		out_type=__snake_case    )
								UpperCAmelCase     :		Optional[Any]							     =		[]
								for piece in pieces:
															if len(__snake_case    ) > 1 and piece[-1] == str(""","""    ) and piece[-2].isdigit():
																						UpperCAmelCase     :		Any							     =		self.sp_model.EncodeAsPieces(piece[:-1].replace(__snake_case		,		""""""    )    )
																						if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
																													if len(cur_pieces[0]    ) == 1:
																																				UpperCAmelCase     :		Dict							     =		cur_pieces[1:]
																													else:
																																				UpperCAmelCase     :		int							     =		cur_pieces[0][1:]
																						cur_pieces.append(piece[-1]    )
																						new_pieces.extend(__snake_case    )
															else:
																						new_pieces.append(__snake_case    )

								return new_pieces





	def  _lowercase(  self		,		A    )			->			Optional[Any]:
								return self.sp_model.PieceToId(__snake_case    )





	def  _lowercase(  self		,		A    )			->			int:
								return self.sp_model.IdToPiece(__snake_case    )





	def  _lowercase(  self		,		A    )			->			List[Any]:
								UpperCAmelCase     :		Optional[int]							     =		"""""".join(__snake_case    ).replace(__snake_case		,		""" """    ).strip()
								return out_string





	def  _lowercase(  self		,		A		,		A = False		,		A = None		,		A = True		,		**A		,		)			->			List[str]:
								UpperCAmelCase     :		List[str]							     =		kwargs.pop("""use_source_tokenizer"""		,		__snake_case    )

								UpperCAmelCase     :		Union[str, Any]							     =		self.convert_ids_to_tokens(__snake_case		,		skip_special_tokens=__snake_case    )

								# To avoid mixing byte-level and unicode for byte-level BPT
								# we need to build string separately for added tokens and byte-level tokens
								# cf. https://github.com/huggingface/transformers/issues/1133
								UpperCAmelCase     :		Union[str, Any]							     =		[]
								UpperCAmelCase     :		str							     =		[]
								for token in filtered_tokens:
															if skip_special_tokens and token in self.all_special_ids:
																						continue
															if token in self.added_tokens_encoder:
																						if current_sub_text:
																													sub_texts.append(self.convert_tokens_to_string(__snake_case    )    )
																													UpperCAmelCase     :		List[Any]							     =		[]
																						sub_texts.append(__snake_case    )
															else:
																						current_sub_text.append(__snake_case    )
								if current_sub_text:
															sub_texts.append(self.convert_tokens_to_string(__snake_case    )    )

								# Mimic the behavior of the Rust tokenizer:
								# By default, there are no spaces between special tokens
								UpperCAmelCase     :		Dict							     =		"""""".join(__snake_case    )

								UpperCAmelCase     :		str							     =		(
								    clean_up_tokenization_spaces
								    if clean_up_tokenization_spaces is not None
								    else self.clean_up_tokenization_spaces
								)
								if clean_up_tokenization_spaces:
															UpperCAmelCase     :		Any							     =		self.clean_up_tokenization(__snake_case    )
															return clean_text
								else:
															return text





	def  _lowercase(  self		,		A		,		A = None    )			->			Tuple:
								UpperCAmelCase     :		str							     =		[self.sep_token_id]
								UpperCAmelCase     :		Any							     =		[self.cls_token_id]
								if token_ids_a is None:
															return token_ids_a + sep + cls
								return token_ids_a + sep + token_ids_a + sep + cls





	def  _lowercase(  self		,		A		,		A = None		,		A = False    )			->			int:

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

								if token_ids_a is not None:
															return ([0] * len(__snake_case    )) + [1] + ([0] * len(__snake_case    )) + [1, 1]
								return ([0] * len(__snake_case    )) + [1, 1]





	def  _lowercase(  self		,		A		,		A = None    )			->			Tuple:
								UpperCAmelCase     :		str							     =		[self.sep_token_id]
								UpperCAmelCase     :		Any							     =		[2]

								if token_ids_a is None:
															return len(token_ids_a + sep    ) * [0] + cls_segment_id
								return len(token_ids_a + sep    ) * [0] + len(token_ids_a + sep    ) * [1] + cls_segment_id






	def  _lowercase(  self		,		A		,		A = None    )			->			Dict:
								if not os.path.isdir(__snake_case    ):
															logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''    )
															return
								UpperCAmelCase     :		int							     =		os.path.join(
								    __snake_case		,		(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]    )

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

								return (out_vocab_file,)





'''simple docstring'''


from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


UpperCAmelCase_		:   List[Any]							   =		logging.get_logger(__name__)







def 			snake_case_   (  SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				try:
								with open(SCREAMING_SNAKE_CASE__							,      """rb"""     ) as flax_state_f:
												_SCREAMING_SNAKE_CASE :       Dict     		=      from_bytes(SCREAMING_SNAKE_CASE__							,      flax_state_f.read()     )
				except UnpicklingError as e:
								try:
												with open(SCREAMING_SNAKE_CASE__     ) as f:
																if f.read().startswith("""version"""     ):
																				raise OSError(
																				    """You seem to have cloned a repository without having git-lfs installed. Please"""
																				    """ install git-lfs and run `git lfs install` followed by `git lfs pull` in the"""
																				    """ folder you cloned."""     )
																else:
																				raise ValueError from e
								except (UnicodeDecodeError, ValueError):
												raise EnvironmentError(f"""Unable to convert {model_file} to Flax deserializable object. """     )

				return load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     )







def 			snake_case_   (  SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				try:
								import torch  # noqa: F401
				except ImportError:
								logger.error(
								    """Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see"""
								    """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
								    """ instructions."""     )
								raise

				# check if we have bf16 weights
				_SCREAMING_SNAKE_CASE :       List[Any]     		=      flatten_dict(jax.tree_util.tree_map(lambda SCREAMING_SNAKE_CASE__     : x.dtype == jnp.bfloataa							,      SCREAMING_SNAKE_CASE__     )     ).values()
				if any(SCREAMING_SNAKE_CASE__     ):
								# convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

								# and bf16 is not fully supported in PT yet.
								logger.warning(
								    """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
								    """before loading those in PyTorch model."""     )
								_SCREAMING_SNAKE_CASE :       Dict     		=      jax.tree_util.tree_map(
								    lambda SCREAMING_SNAKE_CASE__     : params.astype(np.floataa     ) if params.dtype == jnp.bfloataa else params							,      SCREAMING_SNAKE_CASE__     )

				_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      """"""

				_SCREAMING_SNAKE_CASE :       str     		=      flatten_dict(SCREAMING_SNAKE_CASE__							,      sep="""."""     )
				_SCREAMING_SNAKE_CASE :       str     		=      pt_model.state_dict()

				# keep track of unexpected & missing keys
				_SCREAMING_SNAKE_CASE :       Tuple     		=      []
				_SCREAMING_SNAKE_CASE :       int     		=      set(pt_model_dict.keys()     )

				for flax_key_tuple, flax_tensor in flax_state_dict.items():
								_SCREAMING_SNAKE_CASE :       Any     		=      flax_key_tuple.split("""."""     )

								if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
												_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      flax_key_tuple_array[:-1] + ["""weight"""]
												_SCREAMING_SNAKE_CASE :       List[str]     		=      jnp.transpose(SCREAMING_SNAKE_CASE__							,      (3, 2, 0, 1)     )
								elif flax_key_tuple_array[-1] == "kernel":
												_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      flax_key_tuple_array[:-1] + ["""weight"""]
												_SCREAMING_SNAKE_CASE :       Any     		=      flax_tensor.T
								elif flax_key_tuple_array[-1] == "scale":
												_SCREAMING_SNAKE_CASE :       Optional[int]     		=      flax_key_tuple_array[:-1] + ["""weight"""]

								if "time_embedding" not in flax_key_tuple_array:
												for i, flax_key_tuple_string in enumerate(SCREAMING_SNAKE_CASE__     ):
																_SCREAMING_SNAKE_CASE :       Optional[int]     		=      (
																    flax_key_tuple_string.replace("""_0"""							,      """.0"""     )
																    .replace("""_1"""							,      """.1"""     )
																    .replace("""_2"""							,      """.2"""     )
																    .replace("""_3"""							,      """.3"""     )
																    .replace("""_4"""							,      """.4"""     )
																    .replace("""_5"""							,      """.5"""     )
																    .replace("""_6"""							,      """.6"""     )
																    .replace("""_7"""							,      """.7"""     )
																    .replace("""_8"""							,      """.8"""     )
																    .replace("""_9"""							,      """.9"""     )
																)

								_SCREAMING_SNAKE_CASE :       Tuple     		=      """.""".join(SCREAMING_SNAKE_CASE__     )

								if flax_key in pt_model_dict:
												if flax_tensor.shape != pt_model_dict[flax_key].shape:
																raise ValueError(
																    f"""Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected """
																    f"""to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."""     )
												else:
																# add weight to pytorch dict
																_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      np.asarray(SCREAMING_SNAKE_CASE__     ) if not isinstance(SCREAMING_SNAKE_CASE__							,      np.ndarray     ) else flax_tensor
																_SCREAMING_SNAKE_CASE :       int     		=      torch.from_numpy(SCREAMING_SNAKE_CASE__     )
																# remove from missing keys
																missing_keys.remove(SCREAMING_SNAKE_CASE__     )
								else:
												# weight is not expected by PyTorch model
												unexpected_keys.append(SCREAMING_SNAKE_CASE__     )

				pt_model.load_state_dict(SCREAMING_SNAKE_CASE__     )

				# re-transform missing_keys to list
				_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      list(SCREAMING_SNAKE_CASE__     )

				if len(SCREAMING_SNAKE_CASE__     ) > 0:
								logger.warning(
								    """Some weights of the Flax model were not used when initializing the PyTorch model"""
								    f""" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"""
								    f""" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"""
								    """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
								    f""" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"""
								    """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
								    """ FlaxBertForSequenceClassification model)."""     )
				if len(SCREAMING_SNAKE_CASE__     ) > 0:
								logger.warning(
								    f"""Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"""
								    f""" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"""
								    """ use it for predictions and inference."""     )

				return pt_model




