

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


_SCREAMING_SNAKE_CASE  :						Optional[Any]         =    logging.get_logger(__name__)
def        UpperCAmelCase_     (		_A   ,	_A       ):





     '''simple docstring'''




     try:
          with open(__UpperCamelCase   ,	'''rb'''       ) as flax_state_f:
               SCREAMING_SNAKE_CASE__					=					from_bytes(__UpperCamelCase   ,	flax_state_f.read()       )
     except UnpicklingError as e:
          try:
               with open(__UpperCamelCase       ) 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(__UpperCamelCase   ,	__UpperCamelCase       )







def        UpperCAmelCase_     (		_A   ,	_A       ):





     '''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__					=					flatten_dict(jax.tree_util.tree_map(lambda _A       : x.dtype == jnp.bfloataa   ,	__UpperCamelCase       )       ).values()
     if any(__UpperCamelCase       ):
          # 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__					=					jax.tree_util.tree_map(
              lambda _A       : params.astype(np.floataa       ) if params.dtype == jnp.bfloataa else params   ,	__UpperCamelCase       )

     SCREAMING_SNAKE_CASE__					=					''''''

     SCREAMING_SNAKE_CASE__					=					flatten_dict(__UpperCamelCase   ,	sep='''.'''       )
     SCREAMING_SNAKE_CASE__					=					pt_model.state_dict()

     # keep track of unexpected & missing keys
     SCREAMING_SNAKE_CASE__					=					[]
     SCREAMING_SNAKE_CASE__					=					set(pt_model_dict.keys()       )

     for flax_key_tuple, flax_tensor in flax_state_dict.items():
          SCREAMING_SNAKE_CASE__					=					flax_key_tuple.split('''.'''       )

          if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
               SCREAMING_SNAKE_CASE__					=					flax_key_tuple_array[:-1] + ['''weight''']
               SCREAMING_SNAKE_CASE__					=					jnp.transpose(__UpperCamelCase   ,	(3, 2, 0, 1)       )
          elif flax_key_tuple_array[-1] == "kernel":
               SCREAMING_SNAKE_CASE__					=					flax_key_tuple_array[:-1] + ['''weight''']
               SCREAMING_SNAKE_CASE__					=					flax_tensor.T
          elif flax_key_tuple_array[-1] == "scale":
               SCREAMING_SNAKE_CASE__					=					flax_key_tuple_array[:-1] + ['''weight''']

          if "time_embedding" not in flax_key_tuple_array:
               for i, flax_key_tuple_string in enumerate(__UpperCamelCase       ):
                    SCREAMING_SNAKE_CASE__					=					(
                        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__					=					'''.'''.join(__UpperCamelCase       )

          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__					=					np.asarray(__UpperCamelCase       ) if not isinstance(__UpperCamelCase   ,	np.ndarray       ) else flax_tensor
                    SCREAMING_SNAKE_CASE__					=					torch.from_numpy(__UpperCamelCase       )
                    # remove from missing keys
                    missing_keys.remove(__UpperCamelCase       )
          else:
               # weight is not expected by PyTorch model
               unexpected_keys.append(__UpperCamelCase       )

     pt_model.load_state_dict(__UpperCamelCase       )

     # re-transform missing_keys to list
     SCREAMING_SNAKE_CASE__					=					list(__UpperCamelCase       )

     if len(__UpperCamelCase       ) > 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(__UpperCamelCase       ) > 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





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


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

_SCREAMING_SNAKE_CASE  :						int         =    {'''vocab_file''': '''spiece.model'''}

_SCREAMING_SNAKE_CASE  :						int         =    {
    '''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''',
    }
}

_SCREAMING_SNAKE_CASE  :						List[Any]         =    {
    '''xlnet-base-cased''': None,
    '''xlnet-large-cased''': None,
}

# Segments (not really needed)
_SCREAMING_SNAKE_CASE  :						Optional[int]         =    0
_SCREAMING_SNAKE_CASE  :						Optional[Any]         =    1
_SCREAMING_SNAKE_CASE  :						int         =    2
_SCREAMING_SNAKE_CASE  :						str         =    3
_SCREAMING_SNAKE_CASE  :						str         =    4






class  UpperCAmelCase__	( A__ ):






  """simple docstring"""


  a     =      VOCAB_FILES_NAMES
  a     =      PRETRAINED_VOCAB_FILES_MAP
  a     =      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  a     =      "left"






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

       SCREAMING_SNAKE_CASE__					=					{} if sp_model_kwargs is None else sp_model_kwargs

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

       SCREAMING_SNAKE_CASE__					=					3

       SCREAMING_SNAKE_CASE__					=					do_lower_case
       SCREAMING_SNAKE_CASE__					=					remove_space
       SCREAMING_SNAKE_CASE__					=					keep_accents
       SCREAMING_SNAKE_CASE__					=					vocab_file

       SCREAMING_SNAKE_CASE__					=					spm.SentencePieceProcessor(**self.sp_model_kwargs      )
       self.sp_model.Load(__lowerCamelCase      )






  @property
  def    lowercase_	(							self		:	Optional[Any]      )	->							Union[str, Any]:
       return len(self.sp_model      )






  def    lowercase_	(							self		:	Tuple      )	->							Dict:
       SCREAMING_SNAKE_CASE__					=					{self.convert_ids_to_tokens(__lowerCamelCase      ): i for i in range(self.vocab_size      )}
       vocab.update(self.added_tokens_encoder      )
       return vocab






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






  def __setstate__(							self		:	Any			,     __lowerCamelCase		:	Tuple      )	->							Tuple:
       SCREAMING_SNAKE_CASE__					=					d

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

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






  def    lowercase_	(							self		:	Union[str, Any]			,     __lowerCamelCase		:	Union[str, Any]      )	->							List[str]:
       if self.remove_space:
            SCREAMING_SNAKE_CASE__					=					''' '''.join(inputs.strip().split()      )
       else:
            SCREAMING_SNAKE_CASE__					=					inputs
       SCREAMING_SNAKE_CASE__					=					outputs.replace('''``'''			,     '''"'''      ).replace('''\'\''''			,     '''"'''      )

       if not self.keep_accents:
            SCREAMING_SNAKE_CASE__					=					unicodedata.normalize('''NFKD'''			,     __lowerCamelCase      )
            SCREAMING_SNAKE_CASE__					=					''''''.join([c for c in outputs if not unicodedata.combining(__lowerCamelCase      )]      )
       if self.do_lower_case:
            SCREAMING_SNAKE_CASE__					=					outputs.lower()

       return outputs






  def    lowercase_	(							self		:	Any			,     __lowerCamelCase		:	str      )	->							List[str]:
       SCREAMING_SNAKE_CASE__					=					self.preprocess_text(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					self.sp_model.encode(__lowerCamelCase			,     out_type=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					[]
       for piece in pieces:
            if len(__lowerCamelCase      ) > 1 and piece[-1] == str(''','''      ) and piece[-2].isdigit():
                 SCREAMING_SNAKE_CASE__					=					self.sp_model.EncodeAsPieces(piece[:-1].replace(__lowerCamelCase			,     ''''''      )      )
                 if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
                      if len(cur_pieces[0]      ) == 1:
                           SCREAMING_SNAKE_CASE__					=					cur_pieces[1:]
                      else:
                           SCREAMING_SNAKE_CASE__					=					cur_pieces[0][1:]
                 cur_pieces.append(piece[-1]      )
                 new_pieces.extend(__lowerCamelCase      )
            else:
                 new_pieces.append(__lowerCamelCase      )

       return new_pieces






  def    lowercase_	(							self		:	List[Any]			,     __lowerCamelCase		:	Optional[int]      )	->							Optional[Any]:
       return self.sp_model.PieceToId(__lowerCamelCase      )






  def    lowercase_	(							self		:	int			,     __lowerCamelCase		:	Tuple      )	->							List[Any]:
       return self.sp_model.IdToPiece(__lowerCamelCase      )






  def    lowercase_	(							self		:	int			,     __lowerCamelCase		:	Tuple      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					''''''.join(__lowerCamelCase      ).replace(__lowerCamelCase			,     ''' '''      ).strip()
       return out_string






  def    lowercase_	(							self		:	Union[str, Any]			,     __lowerCamelCase		:	List[int]			,     __lowerCamelCase		:	bool = False			,     __lowerCamelCase		:	bool = None			,     __lowerCamelCase		:	bool = True			,     **__lowerCamelCase		:	int			,     )	->							str:
       SCREAMING_SNAKE_CASE__					=					kwargs.pop('''use_source_tokenizer'''			,     __lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					self.convert_ids_to_tokens(__lowerCamelCase			,     skip_special_tokens=__lowerCamelCase      )

       # 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
       SCREAMING_SNAKE_CASE__					=					[]
       SCREAMING_SNAKE_CASE__					=					[]
       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(__lowerCamelCase      )      )
                      SCREAMING_SNAKE_CASE__					=					[]
                 sub_texts.append(__lowerCamelCase      )
            else:
                 current_sub_text.append(__lowerCamelCase      )
       if current_sub_text:
            sub_texts.append(self.convert_tokens_to_string(__lowerCamelCase      )      )

       # Mimic the behavior of the Rust tokenizer:
       # By default, there are no spaces between special tokens
       SCREAMING_SNAKE_CASE__					=					''''''.join(__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					(
           clean_up_tokenization_spaces
           if clean_up_tokenization_spaces is not None
           else self.clean_up_tokenization_spaces
       )
       if clean_up_tokenization_spaces:
            SCREAMING_SNAKE_CASE__					=					self.clean_up_tokenization(__lowerCamelCase      )
            return clean_text
       else:
            return text






  def    lowercase_	(							self		:	Dict			,     __lowerCamelCase		:	List[int]			,     __lowerCamelCase		:	Optional[List[int]] = None      )	->							List[int]:
       SCREAMING_SNAKE_CASE__					=					[self.sep_token_id]
       SCREAMING_SNAKE_CASE__					=					[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		:	List[Any]			,     __lowerCamelCase		:	List[int]			,     __lowerCamelCase		:	Optional[List[int]] = None			,     __lowerCamelCase		:	bool = False      )	->							List[int]:

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

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






  def    lowercase_	(							self		:	List[str]			,     __lowerCamelCase		:	List[int]			,     __lowerCamelCase		:	Optional[List[int]] = None      )	->							List[int]:
       SCREAMING_SNAKE_CASE__					=					[self.sep_token_id]
       SCREAMING_SNAKE_CASE__					=					[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		:	str			,     __lowerCamelCase		:	str			,     __lowerCamelCase		:	Optional[str] = None      )	->							Tuple[str]:
       if not os.path.isdir(__lowerCamelCase      ):
            logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''      )
            return
       SCREAMING_SNAKE_CASE__					=					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:
                 SCREAMING_SNAKE_CASE__					=					self.sp_model.serialized_model_proto()
                 fi.write(__lowerCamelCase      )

       return (out_vocab_file,)


