


'''simple docstring'''


import unittest

from transformers import SPIECE_UNDERLINE, ReformerTokenizer, ReformerTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, require_torch, slow
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin


UpperCAmelCase_		:   Optional[Any]							   =		get_tests_dir('fixtures/test_sentencepiece.model')





@require_sentencepiece
@require_tokenizers
class     lowercase__  (						_snake_case     ,     unittest.TestCase					):





							'''simple docstring'''


							A_							:      str      	=  ReformerTokenizer
							A_							:      Optional[int]      	=  ReformerTokenizerFast
							A_							:      str      	=  True
							A_							:      Tuple      	=  False
							A_							:      str      	=  True






							def UpperCAmelCase_	(    self  ):
											super().setUp()

											_SCREAMING_SNAKE_CASE :       int     		=      ReformerTokenizer(__snake_case    ,					keep_accents=__snake_case  )
											tokenizer.save_pretrained(self.tmpdirname  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      """<s>"""
											_SCREAMING_SNAKE_CASE :       Dict     		=      1

											self.assertEqual(self.get_tokenizer()._convert_token_to_id(__snake_case  )    ,					__snake_case  )
											self.assertEqual(self.get_tokenizer()._convert_id_to_token(__snake_case  )    ,					__snake_case  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      list(self.get_tokenizer().get_vocab().keys()  )

											self.assertEqual(vocab_keys[0]    ,					"""<unk>"""  )
											self.assertEqual(vocab_keys[1]    ,					"""<s>"""  )
											self.assertEqual(vocab_keys[-1]    ,					"""j"""  )
											self.assertEqual(len(__snake_case  )    ,					1000  )






							def UpperCAmelCase_	(    self  ):
											self.assertEqual(self.get_tokenizer().vocab_size    ,					1000  )






							def UpperCAmelCase_	(    self  ):
											if not self.test_rust_tokenizer:
															return

											_SCREAMING_SNAKE_CASE :       Dict     		=      self.get_tokenizer()
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      self.get_rust_tokenizer()

											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      """I was born in 92000, and this is falsé."""

											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      tokenizer.tokenize(__snake_case  )
											_SCREAMING_SNAKE_CASE :       str     		=      rust_tokenizer.tokenize(__snake_case  )
											self.assertListEqual(__snake_case    ,					__snake_case  )

											_SCREAMING_SNAKE_CASE :       str     		=      tokenizer.encode(__snake_case    ,					add_special_tokens=__snake_case  )
											_SCREAMING_SNAKE_CASE :       int     		=      rust_tokenizer.encode(__snake_case    ,					add_special_tokens=__snake_case  )
											self.assertListEqual(__snake_case    ,					__snake_case  )

											_SCREAMING_SNAKE_CASE :       List[str]     		=      self.get_rust_tokenizer()
											_SCREAMING_SNAKE_CASE :       Dict     		=      tokenizer.encode(__snake_case  )
											_SCREAMING_SNAKE_CASE :       Dict     		=      rust_tokenizer.encode(__snake_case  )
											self.assertListEqual(__snake_case    ,					__snake_case  )






							def UpperCAmelCase_	(    self    ,					__snake_case=15  ):
											for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
															with self.subTest(f"""{tokenizer.__class__.__name__} ({pretrained_name})"""  ):
																			_SCREAMING_SNAKE_CASE :       Any     		=      self.rust_tokenizer_class.from_pretrained(__snake_case    ,					**__snake_case  )

																			# Simple input
																			_SCREAMING_SNAKE_CASE :       Tuple     		=      """This is a simple input"""
																			_SCREAMING_SNAKE_CASE :       List[str]     		=      ["""This is a simple input 1""", """This is a simple input 2"""]
																			_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      ("""This is a simple input""", """This is a pair""")
																			_SCREAMING_SNAKE_CASE :       List[str]     		=      [
																			    ("""This is a simple input 1""", """This is a simple input 2"""),
																			    ("""This is a simple pair 1""", """This is a simple pair 2"""),
																			]

																			# Simple input tests
																			self.assertRaises(__snake_case    ,					tokenizer_r.encode    ,					__snake_case    ,					max_length=__snake_case    ,					padding="""max_length"""  )

																			# Simple input
																			self.assertRaises(__snake_case    ,					tokenizer_r.encode_plus    ,					__snake_case    ,					max_length=__snake_case    ,					padding="""max_length"""  )

																			# Simple input
																			self.assertRaises(
																			    __snake_case    ,					tokenizer_r.batch_encode_plus    ,					__snake_case    ,					max_length=__snake_case    ,					padding="""max_length"""    ,					)

																			# Pair input
																			self.assertRaises(__snake_case    ,					tokenizer_r.encode    ,					__snake_case    ,					max_length=__snake_case    ,					padding="""max_length"""  )

																			# Pair input
																			self.assertRaises(__snake_case    ,					tokenizer_r.encode_plus    ,					__snake_case    ,					max_length=__snake_case    ,					padding="""max_length"""  )

																			# Pair input
																			self.assertRaises(
																			    __snake_case    ,					tokenizer_r.batch_encode_plus    ,					__snake_case    ,					max_length=__snake_case    ,					padding="""max_length"""    ,					)






							def UpperCAmelCase_	(    self  ):
											pass






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      ReformerTokenizer(__snake_case    ,					keep_accents=__snake_case  )

											_SCREAMING_SNAKE_CASE :       List[str]     		=      tokenizer.tokenize("""This is a test"""  )
											self.assertListEqual(__snake_case    ,					["""▁This""", """▁is""", """▁a""", """▁t""", """est"""]  )

											self.assertListEqual(
											    tokenizer.convert_tokens_to_ids(__snake_case  )    ,					[285, 46, 10, 170, 382]    ,					)

											_SCREAMING_SNAKE_CASE :       List[str]     		=      tokenizer.tokenize("""I was born in 92000, and this is falsé."""  )
											self.assertListEqual(
											    __snake_case    ,					[
											        SPIECE_UNDERLINE + """I""",
											        SPIECE_UNDERLINE + """was""",
											        SPIECE_UNDERLINE + """b""",
											        """or""",
											        """n""",
											        SPIECE_UNDERLINE + """in""",
											        SPIECE_UNDERLINE + """""",
											        """9""",
											        """2""",
											        """0""",
											        """0""",
											        """0""",
											        """,""",
											        SPIECE_UNDERLINE + """and""",
											        SPIECE_UNDERLINE + """this""",
											        SPIECE_UNDERLINE + """is""",
											        SPIECE_UNDERLINE + """f""",
											        """al""",
											        """s""",
											        """é""",
											        """.""",
											    ]    ,					)
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      tokenizer.convert_tokens_to_ids(__snake_case  )
											self.assertListEqual(
											    __snake_case    ,					[8, 21, 84, 55, 24, 19, 7, 0, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 0, 4]    ,					)

											_SCREAMING_SNAKE_CASE :       List[Any]     		=      tokenizer.convert_ids_to_tokens(__snake_case  )
											self.assertListEqual(
											    __snake_case    ,					[
											        SPIECE_UNDERLINE + """I""",
											        SPIECE_UNDERLINE + """was""",
											        SPIECE_UNDERLINE + """b""",
											        """or""",
											        """n""",
											        SPIECE_UNDERLINE + """in""",
											        SPIECE_UNDERLINE + """""",
											        """<unk>""",
											        """2""",
											        """0""",
											        """0""",
											        """0""",
											        """,""",
											        SPIECE_UNDERLINE + """and""",
											        SPIECE_UNDERLINE + """this""",
											        SPIECE_UNDERLINE + """is""",
											        SPIECE_UNDERLINE + """f""",
											        """al""",
											        """s""",
											        """<unk>""",
											        """.""",
											    ]    ,					)






							@cached_property
							def UpperCAmelCase_	(    self  ):
											return ReformerTokenizer.from_pretrained("""google/reformer-crime-and-punishment"""  )






							@slow
							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       str     		=      """Hello World!"""
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      [126, 32, 262, 152, 38, 72, 287]

											self.assertListEqual(__snake_case    ,					self.big_tokenizer.encode(__snake_case  )  )






							@slow
							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      (
											    """This is a very long text with a lot of weird characters, such as: . , ~ ? ( ) \" [ ] ! : - . Also we will"""
											    """ add words that should not exsist and be tokenized to <unk>, such as saoneuhaoesuth"""
											)
											_SCREAMING_SNAKE_CASE :       int     		=      [
											    108,
											    265,
											    24,
											    111,
											    4,
											    258,
											    156,
											    35,
											    28,
											    275,
											    3,
											    259,
											    297,
											    260,
											    84,
											    4,
											    35,
											    110,
											    44,
											    8,
											    259,
											    91,
											    268,
											    21,
											    11,
											    209,
											    274,
											    109,
											    266,
											    277,
											    117,
											    86,
											    93,
											    315,
											    258,
											    278,
											    258,
											    277,
											    258,
											    0,
											    258,
											    288,
											    258,
											    319,
											    258,
											    0,
											    258,
											    0,
											    258,
											    0,
											    258,
											    0,
											    258,
											    287,
											    258,
											    315,
											    258,
											    289,
											    258,
											    278,
											    99,
											    269,
											    266,
											    262,
											    8,
											    259,
											    241,
											    4,
											    217,
											    230,
											    268,
											    266,
											    55,
											    168,
											    106,
											    75,
											    193,
											    266,
											    223,
											    27,
											    49,
											    26,
											    282,
											    25,
											    264,
											    299,
											    19,
											    26,
											    0,
											    258,
											    277,
											    117,
											    86,
											    93,
											    176,
											    183,
											    270,
											    11,
											    262,
											    42,
											    61,
											    265,
											]

											self.assertListEqual(__snake_case    ,					self.big_tokenizer.encode(__snake_case  )  )






							@require_torch
							@slow
							def UpperCAmelCase_	(    self  ):
											import torch

											from transformers import ReformerConfig, ReformerModel

											# Build sequence
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      list(self.big_tokenizer.get_vocab().keys()  )[:10]
											_SCREAMING_SNAKE_CASE :       int     		=      """ """.join(__snake_case  )
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      self.big_tokenizer.encode_plus(__snake_case    ,					return_tensors="""pt"""  )
											_SCREAMING_SNAKE_CASE :       Dict     		=      self.big_tokenizer.batch_encode_plus([sequence, sequence]    ,					return_tensors="""pt"""  )

											_SCREAMING_SNAKE_CASE :       int     		=      ReformerConfig()
											# The input gets padded during training so adjust the axial position encodings from the pretrained model value of (512, 1024)
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      encoded_sequence["""input_ids"""].shape
											_SCREAMING_SNAKE_CASE :       Tuple     		=      ReformerModel(__snake_case  )

											# Reformer has config.vocab_size == tokenizer.vocab_size == len(tokenizer) - 1 = 320; len(tokenizer) is 321 (including a pad token with id 320)
											assert model.get_input_embeddings().weight.shape[0] >= self.big_tokenizer.vocab_size

											with torch.no_grad():
															model(**__snake_case  )
															model(**__snake_case  )






							@slow
							def UpperCAmelCase_	(    self  ):
											# fmt: off
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      {"""input_ids""": [[108, 265, 24, 111, 4, 258, 156, 7, 51, 279, 58, 7, 76, 25, 69, 278], [140, 243, 264, 134, 17, 267, 77, 263, 22, 262, 297, 258, 304, 177, 279, 266, 14, 89, 13, 35, 261, 299, 272, 137, 275, 278]], """attention_mask""": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}  # noqa: E501
											# fmt: on

											# This tokenizer does not know some characters like ")".
											# That is the reason why we use very simple texts here.
											# Also see https://github.com/huggingface/transformers/pull/11737#issuecomment-850769064
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      [
											    """This is a very simple sentence.""",
											    """The quick brown fox jumps over the lazy dog.""",
											]

											self.tokenizer_integration_test_util(
											    expected_encoding=__snake_case    ,					model_name="""google/reformer-crime-and-punishment"""    ,					revision="""0e6c3decb8211d49bf881013425dc8b0448b3f5a"""    ,					padding=__snake_case    ,					sequences=__snake_case    ,					)








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




