






import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
							import os

							import jax.numpy as jnp
							from jax import jit

							from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
							from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

							UpperCAmelCase__						:			List[Any]									='''0.12'''  # assumed parallelism: 8


if is_torch_available():
							import torch





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase=None				)						->  Dict:
				if rng is None:
								lowerCamelCase	 =random.Random()

				lowerCamelCase	 =1
				for dim in shape:
								total_dims *= dim

				lowerCamelCase	 =[]
				for _ in range(_UpperCAmelCase				):
								values.append(rng.randint(0			,      vocab_size - 1				)				)

				lowerCamelCase	 =np.array(_UpperCAmelCase			,      dtype=jnp.intaa				).reshape(_UpperCAmelCase				)

				return output





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase=None				)						->  Optional[Any]:
				lowerCamelCase	 =ids_tensor(_UpperCAmelCase			,      vocab_size=2			,      rng=_UpperCAmelCase				)
				# make sure that at least one token is attended to for each batch
				lowerCamelCase	 =1
				return attn_mask

@require_flax
class       __A     :
				__A      =				None
				__A      =				()


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()

								# cut to half length & take max batch_size 3
								lowerCamelCase	 =2
								lowerCamelCase	 =inputs["""input_ids"""].shape[-1] // 2
								lowerCamelCase	 =inputs["""input_ids"""][:max_batch_size, :sequence_length]

								lowerCamelCase	 =jnp.ones_like(UpperCAmelCase_							)
								lowerCamelCase	 =attention_mask[:max_batch_size, :sequence_length]

								# generate max 5 tokens
								lowerCamelCase	 =input_ids.shape[-1] + 5
								if config.eos_token_id is not None and config.pad_token_id is None:
												# hack to allow generate for models such as GPT2 as is done in `generate()`
												lowerCamelCase	 =config.eos_token_id
								return config, input_ids, attention_mask, max_length


				@is_pt_flax_cross_test
				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =self._get_input_ids_and_config()
								lowerCamelCase	 =False
								lowerCamelCase	 =max_length
								lowerCamelCase	 =0

								for model_class in self.all_generative_model_classes:
												lowerCamelCase	 =model_class(UpperCAmelCase_							)

												lowerCamelCase	 =model_class.__name__[4:]  # Skip the "Flax" at the beginning
												lowerCamelCase	 =getattr(UpperCAmelCase_   , UpperCAmelCase_							)
												lowerCamelCase	 =pt_model_class(UpperCAmelCase_							).eval()
												lowerCamelCase	 =load_flax_weights_in_pytorch_model(UpperCAmelCase_   , flax_model.params							)

												lowerCamelCase	 =flax_model.generate(UpperCAmelCase_							).sequences
												lowerCamelCase	 =pt_model.generate(torch.tensor(UpperCAmelCase_   , dtype=torch.long							)							)

												if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
																lowerCamelCase	 =flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

												self.assertListEqual(pt_generation_outputs.numpy().tolist()   , flax_generation_outputs.tolist()							)


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =self._get_input_ids_and_config()
								lowerCamelCase	 =False
								lowerCamelCase	 =max_length

								for model_class in self.all_generative_model_classes:
												lowerCamelCase	 =model_class(UpperCAmelCase_							)

												lowerCamelCase	 =model.generate(UpperCAmelCase_							).sequences
												self.assertEqual(generation_outputs.shape[-1]   , UpperCAmelCase_							)

												lowerCamelCase	 =jit(model.generate							)
												lowerCamelCase	 =jit_generate(UpperCAmelCase_							).sequences

												self.assertListEqual(generation_outputs.tolist()   , jit_generation_outputs.tolist()							)


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =self._get_input_ids_and_config()
								lowerCamelCase	 =True
								lowerCamelCase	 =max_length

								for model_class in self.all_generative_model_classes:
												lowerCamelCase	 =model_class(UpperCAmelCase_							)

												lowerCamelCase	 =model.generate(UpperCAmelCase_							).sequences
												self.assertEqual(generation_outputs.shape[-1]   , UpperCAmelCase_							)

												lowerCamelCase	 =jit(model.generate							)
												lowerCamelCase	 =jit_generate(UpperCAmelCase_							).sequences

												self.assertListEqual(generation_outputs.tolist()   , jit_generation_outputs.tolist()							)


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =self._get_input_ids_and_config()
								lowerCamelCase	 =False
								lowerCamelCase	 =max_length
								lowerCamelCase	 =2

								for model_class in self.all_generative_model_classes:
												lowerCamelCase	 =model_class(UpperCAmelCase_							)

												lowerCamelCase	 =model.generate(UpperCAmelCase_							).sequences
												self.assertEqual(generation_outputs.shape[-1]   , UpperCAmelCase_							)

												lowerCamelCase	 =jit(model.generate							)
												lowerCamelCase	 =jit_generate(UpperCAmelCase_							).sequences

												self.assertListEqual(generation_outputs.tolist()   , jit_generation_outputs.tolist()							)


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =self._get_input_ids_and_config()
								lowerCamelCase	 =False
								lowerCamelCase	 =max_length
								lowerCamelCase	 =2
								lowerCamelCase	 =2

								for model_class in self.all_generative_model_classes:
												lowerCamelCase	 =model_class(UpperCAmelCase_							)

												lowerCamelCase	 =model.generate(UpperCAmelCase_							).sequences
												self.assertEqual(generation_outputs.shape[0]   , input_ids.shape[0] * config.num_return_sequences							)


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =self._get_input_ids_and_config()
								lowerCamelCase	 =True
								lowerCamelCase	 =max_length
								lowerCamelCase	 =0.8
								lowerCamelCase	 =10
								lowerCamelCase	 =0.3
								lowerCamelCase	 =1
								lowerCamelCase	 =8
								lowerCamelCase	 =9

								for model_class in self.all_generative_model_classes:
												lowerCamelCase	 =model_class(UpperCAmelCase_							)

												lowerCamelCase	 =model.generate(UpperCAmelCase_							).sequences
												self.assertEqual(generation_outputs.shape[-1]   , UpperCAmelCase_							)

												lowerCamelCase	 =jit(model.generate							)
												lowerCamelCase	 =jit_generate(UpperCAmelCase_							).sequences

												self.assertListEqual(generation_outputs.tolist()   , jit_generation_outputs.tolist()							)


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =self._get_input_ids_and_config()
								lowerCamelCase	 =max_length
								lowerCamelCase	 =1
								lowerCamelCase	 =8
								lowerCamelCase	 =9

								for model_class in self.all_generative_model_classes:
												lowerCamelCase	 =model_class(UpperCAmelCase_							)

												lowerCamelCase	 =model.generate(UpperCAmelCase_							).sequences
												self.assertEqual(generation_outputs.shape[-1]   , UpperCAmelCase_							)

												lowerCamelCase	 =jit(model.generate							)
												lowerCamelCase	 =jit_generate(UpperCAmelCase_							).sequences

												self.assertListEqual(generation_outputs.tolist()   , jit_generation_outputs.tolist()							)


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =self._get_input_ids_and_config()
								lowerCamelCase	 =max_length
								lowerCamelCase	 =2
								lowerCamelCase	 =1
								lowerCamelCase	 =8
								lowerCamelCase	 =9

								for model_class in self.all_generative_model_classes:
												lowerCamelCase	 =model_class(UpperCAmelCase_							)

												lowerCamelCase	 =model.generate(UpperCAmelCase_							).sequences
												self.assertEqual(generation_outputs.shape[-1]   , UpperCAmelCase_							)

												lowerCamelCase	 =jit(model.generate							)
												lowerCamelCase	 =jit_generate(UpperCAmelCase_							).sequences

												self.assertListEqual(generation_outputs.tolist()   , jit_generation_outputs.tolist()							)


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =self._get_input_ids_and_config()

								# pad attention mask on the left
								lowerCamelCase	 =attention_mask.at[(0, 0)].set(0							)

								lowerCamelCase	 =False
								lowerCamelCase	 =max_length

								for model_class in self.all_generative_model_classes:
												lowerCamelCase	 =model_class(UpperCAmelCase_							)

												lowerCamelCase	 =model.generate(UpperCAmelCase_   , attention_mask=UpperCAmelCase_							).sequences
												self.assertEqual(generation_outputs.shape[-1]   , UpperCAmelCase_							)

												lowerCamelCase	 =jit(model.generate							)
												lowerCamelCase	 =jit_generate(UpperCAmelCase_   , attention_mask=UpperCAmelCase_							).sequences

												self.assertListEqual(generation_outputs.tolist()   , jit_generation_outputs.tolist()							)


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =self._get_input_ids_and_config()

								# pad attention mask on the left
								lowerCamelCase	 =attention_mask.at[(0, 0)].set(0							)

								lowerCamelCase	 =True
								lowerCamelCase	 =max_length

								for model_class in self.all_generative_model_classes:
												lowerCamelCase	 =model_class(UpperCAmelCase_							)

												lowerCamelCase	 =model.generate(UpperCAmelCase_   , attention_mask=UpperCAmelCase_							).sequences
												self.assertEqual(generation_outputs.shape[-1]   , UpperCAmelCase_							)

												lowerCamelCase	 =jit(model.generate							)
												lowerCamelCase	 =jit_generate(UpperCAmelCase_   , attention_mask=UpperCAmelCase_							).sequences

												self.assertListEqual(generation_outputs.tolist()   , jit_generation_outputs.tolist()							)


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =self._get_input_ids_and_config()

								# pad attention mask on the left
								lowerCamelCase	 =attention_mask.at[(0, 0)].set(0							)

								lowerCamelCase	 =2
								lowerCamelCase	 =max_length

								for model_class in self.all_generative_model_classes:
												lowerCamelCase	 =model_class(UpperCAmelCase_							)

												lowerCamelCase	 =model.generate(UpperCAmelCase_   , attention_mask=UpperCAmelCase_							).sequences
												self.assertEqual(generation_outputs.shape[-1]   , UpperCAmelCase_							)

												lowerCamelCase	 =jit(model.generate							)
												lowerCamelCase	 =jit_generate(UpperCAmelCase_   , attention_mask=UpperCAmelCase_							).sequences

												self.assertListEqual(generation_outputs.tolist()   , jit_generation_outputs.tolist()							)




@require_flax
class       __A     (						unittest.TestCase    ):


				def _snake_case       (		self							):
								lowerCamelCase	 =AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-bert"""							)
								lowerCamelCase	 =FlaxAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""							)

								lowerCamelCase	 ="""Hello world"""
								lowerCamelCase	 =tokenizer(UpperCAmelCase_   , return_tensors="""np"""							).input_ids

								# typos are quickly detected (the correct argument is `do_sample`)
								with self.assertRaisesRegex(UpperCAmelCase_   , """do_samples"""							):
												model.generate(UpperCAmelCase_   , do_samples=UpperCAmelCase_							)

								# arbitrary arguments that will not be used anywhere are also not accepted
								with self.assertRaisesRegex(UpperCAmelCase_   , """foo"""							):
												lowerCamelCase	 ={"""foo""": """bar"""}
												model.generate(UpperCAmelCase_   , **UpperCAmelCase_							)












from typing import List

from .keymap import KEYMAP, get_character





def      _lowercase					(						_UpperCAmelCase				)						->  Tuple:

				def decorator(_UpperCAmelCase				):
								lowerCamelCase	 =getattr(_UpperCAmelCase			,      """handle_key"""			,      []				)
								handle += [key]
								setattr(_UpperCAmelCase			,      """handle_key"""			,      _UpperCAmelCase				)
								return func

				return decorator





def      _lowercase					(						*_UpperCAmelCase				)						->  Tuple:

				def decorator(_UpperCAmelCase				):
								lowerCamelCase	 =getattr(_UpperCAmelCase			,      """handle_key"""			,      []				)
								handle += keys
								setattr(_UpperCAmelCase			,      """handle_key"""			,      _UpperCAmelCase				)
								return func

				return decorator

class       __A     (						a    ):


				def __new__(		cls   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							):
								lowerCamelCase	 =super().__new__(cls   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							)
								if not hasattr(UpperCAmelCase_   , """key_handler"""							):
												setattr(UpperCAmelCase_   , """key_handler"""   , {}							)
								setattr(UpperCAmelCase_   , """handle_input"""   , KeyHandler.handle_input							)

								for value in attrs.values():
												lowerCamelCase	 =getattr(UpperCAmelCase_   , """handle_key"""   , []							)
												for key in handled_keys:
																lowerCamelCase	 =value
								return new_cls


				@staticmethod
				def _snake_case       (		cls							):
								lowerCamelCase	 =get_character()
								if char != KEYMAP["undefined"]:
												lowerCamelCase	 =ord(UpperCAmelCase_							)
								lowerCamelCase	 =cls.key_handler.get(UpperCAmelCase_							)
								if handler:
												lowerCamelCase	 =char
												return handler(cls							)
								else:
												return None





def      _lowercase					(						cls				)						->  List[str]:
				return KeyHandler(cls.__name__			,      cls.__bases__			,      cls.__dict__.copy()				)





