

import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
						import jax
						import jax.numpy as jnp

						from transformers.modeling_flax_pytorch_utils import (
						    convert_pytorch_state_dict_to_flax,
						    load_flax_weights_in_pytorch_model,
						)
						from transformers.models.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
						import torch

class 					__UpperCamelCase					:







	def __init__( self:			Dict	,			__UpperCamelCase:			Optional[Any]	,			__UpperCamelCase:			Optional[int]=14	,			__UpperCamelCase:			Any=7	,			__UpperCamelCase:			Any=True	,			__UpperCamelCase:			List[str]=True	,			__UpperCamelCase:			List[str]=False	,			__UpperCamelCase:			List[str]=True	,			__UpperCamelCase:			str=99	,			__UpperCamelCase:			List[Any]=32	,			__UpperCamelCase:			Dict=5	,			__UpperCamelCase:			Dict=4	,			__UpperCamelCase:			Optional[int]=37	,			__UpperCamelCase:			Union[str, Any]="gelu"	,			__UpperCamelCase:			Any=0.1	,			__UpperCamelCase:			List[str]=0.1	,			__UpperCamelCase:			Any=5_12	,			__UpperCamelCase:			List[str]=0.02	,			):

						'''simple docstring'''





						__magic_name__				=     parent
						__magic_name__				=     batch_size
						__magic_name__				=     seq_length
						__magic_name__				=     is_training
						__magic_name__				=     use_input_mask
						__magic_name__				=     use_token_type_ids
						__magic_name__				=     use_labels
						__magic_name__				=     vocab_size
						__magic_name__				=     hidden_size
						__magic_name__				=     num_hidden_layers
						__magic_name__				=     num_attention_heads
						__magic_name__				=     intermediate_size
						__magic_name__				=     hidden_act
						__magic_name__				=     hidden_dropout_prob
						__magic_name__				=     attention_probs_dropout_prob
						__magic_name__				=     max_position_embeddings
						__magic_name__				=     initializer_range
						__magic_name__				=     None
						__magic_name__				=     vocab_size - 1
						__magic_name__				=     vocab_size - 1
						__magic_name__				=     vocab_size - 1







	def 						_SCREAMING_SNAKE_CASE				( self:			str   ):

						'''simple docstring'''





						__magic_name__				=     ids_tensor([self.batch_size, self.seq_length]	,			self.vocab_size   )

						__magic_name__				=     None
						if self.use_input_mask:
											__magic_name__				=     random_attention_mask([self.batch_size, self.seq_length]   )

						__magic_name__				=     GPTaConfig(
						    vocab_size=self.vocab_size	,			n_embd=self.hidden_size	,			n_layer=self.num_hidden_layers	,			n_head=self.num_attention_heads	,			n_positions=self.max_position_embeddings	,			use_cache=__UpperCamelCase	,			bos_token_id=self.bos_token_id	,			eos_token_id=self.eos_token_id	,			pad_token_id=self.pad_token_id	,			)

						return (config, input_ids, input_mask)







	def 						_SCREAMING_SNAKE_CASE				( self:			List[str]   ):

						'''simple docstring'''





						__magic_name__				=     self.prepare_config_and_inputs()
						__magic_name__,						__magic_name__,						__magic_name__				=     config_and_inputs
						__magic_name__				=     {'input_ids': input_ids, 'attention_mask': attention_mask}
						return config, inputs_dict







	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[int]   ):

						'''simple docstring'''





						__magic_name__,						__magic_name__,						__magic_name__				=     self.prepare_config_and_inputs()

						__magic_name__				=     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
						__magic_name__				=     ids_tensor([self.batch_size, self.seq_length]	,			vocab_size=2   )

						return (
						    config,
						    input_ids,
						    attention_mask,
						    encoder_hidden_states,
						    encoder_attention_mask,
						)







	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[int]	,			__UpperCamelCase:			List[Any]	,			__UpperCamelCase:			Optional[int]	,			__UpperCamelCase:			Optional[int]	,			__UpperCamelCase:			List[Any]   ):

						'''simple docstring'''





						__magic_name__				=     20
						__magic_name__				=     model_class_name(__UpperCamelCase   )

						__magic_name__				=     model.init_cache(input_ids.shape[0]	,			__UpperCamelCase   )
						__magic_name__				=     jnp.ones((input_ids.shape[0], max_decoder_length)	,			dtype='i4'   )

						__magic_name__				=     jnp.broadcast_to(
						    jnp.arange(input_ids.shape[-1] - 1   )[None, :]	,			(input_ids.shape[0], input_ids.shape[-1] - 1)   )
						__magic_name__				=     model(
						    input_ids[:, :-1]	,			attention_mask=__UpperCamelCase	,			past_key_values=__UpperCamelCase	,			position_ids=__UpperCamelCase	,			)

						__magic_name__				=     jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]	,			dtype='i4'   )
						__magic_name__				=     model(
						    input_ids[:, -1:]	,			attention_mask=__UpperCamelCase	,			past_key_values=outputs_cache.past_key_values	,			position_ids=__UpperCamelCase	,			)

						__magic_name__				=     model(__UpperCamelCase   )

						__magic_name__				=     np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])   )   )
						self.parent.assertTrue(diff < 1E-3	,			msg=F'Max diff is {diff}'   )







	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[Any]	,			__UpperCamelCase:			Optional[Any]	,			__UpperCamelCase:			int	,			__UpperCamelCase:			Dict	,			__UpperCamelCase:			Dict   ):

						'''simple docstring'''





						__magic_name__				=     20
						__magic_name__				=     model_class_name(__UpperCamelCase   )

						__magic_name__				=     jnp.concatenate(
						    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])   )]	,			axis=-1	,			)

						__magic_name__				=     model.init_cache(input_ids.shape[0]	,			__UpperCamelCase   )
						__magic_name__				=     jnp.broadcast_to(
						    jnp.arange(input_ids.shape[-1] - 1   )[None, :]	,			(input_ids.shape[0], input_ids.shape[-1] - 1)   )

						__magic_name__				=     model(
						    input_ids[:, :-1]	,			attention_mask=__UpperCamelCase	,			past_key_values=__UpperCamelCase	,			position_ids=__UpperCamelCase	,			)
						__magic_name__				=     jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]	,			dtype='i4'   )
						__magic_name__				=     model(
						    input_ids[:, -1:]	,			past_key_values=outputs_cache.past_key_values	,			attention_mask=__UpperCamelCase	,			position_ids=__UpperCamelCase	,			)

						__magic_name__				=     model(__UpperCamelCase	,			attention_mask=__UpperCamelCase   )

						__magic_name__				=     np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])   )   )
						self.parent.assertTrue(diff < 1E-3	,			msg=F'Max diff is {diff}'   )




@require_flax
class 					__UpperCamelCase					(       SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE	,    unittest.TestCase      ):
	_lowercase	: int												=	(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
	_lowercase	: Optional[Any]												=	(FlaxGPTaLMHeadModel,) if is_flax_available() else ()







	def 						_SCREAMING_SNAKE_CASE				( self:			List[Any]   ):

						'''simple docstring'''





						__magic_name__				=     FlaxGPTaModelTester(self   )







	def 						_SCREAMING_SNAKE_CASE				( self:			str   ):

						'''simple docstring'''





						for model_class_name in self.all_model_classes:
											__magic_name__,						__magic_name__,						__magic_name__				=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.check_use_cache_forward(__UpperCamelCase	,			__UpperCamelCase	,			__UpperCamelCase	,			__UpperCamelCase   )







	def 						_SCREAMING_SNAKE_CASE				( self:			int   ):

						'''simple docstring'''





						for model_class_name in self.all_model_classes:
											__magic_name__,						__magic_name__,						__magic_name__				=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.check_use_cache_forward_with_attn_mask(
											    __UpperCamelCase	,			__UpperCamelCase	,			__UpperCamelCase	,			__UpperCamelCase   )







	@slow
	def 						_SCREAMING_SNAKE_CASE				( self:			List[Any]   ):

						'''simple docstring'''





						__magic_name__				=     GPTaTokenizer.from_pretrained('gpt2'	,			pad_token='</s>'	,			padding_side='left'   )
						__magic_name__				=     tokenizer(['Hello this is a long string', 'Hey']	,			return_tensors='np'	,			padding=__UpperCamelCase	,			truncation=__UpperCamelCase   )

						__magic_name__				=     FlaxGPTaLMHeadModel.from_pretrained('gpt2'   )
						__magic_name__				=     False
						__magic_name__				=     model.config.eos_token_id

						__magic_name__				=     jax.jit(model.generate   )

						__magic_name__				=     jit_generate(inputs['input_ids']	,			attention_mask=inputs['attention_mask']   ).sequences

						__magic_name__				=     tokenizer.batch_decode(__UpperCamelCase	,			skip_special_tokens=__UpperCamelCase   )

						__magic_name__				=     [
						    'Hello this is a long string of words. I\'m going to start with the first one.\n',
						    'Hey, I\'m not sure if I\'m going to be able to do',
						]

						self.assertListEqual(__UpperCamelCase	,			__UpperCamelCase   )







	@is_pt_flax_cross_test
	def 						_SCREAMING_SNAKE_CASE				( self:			Tuple   ):

						'''simple docstring'''





						__magic_name__,						__magic_name__				=     self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
											with self.subTest(model_class.__name__   ):
																# prepare inputs
																__magic_name__				=     self._prepare_for_class(__UpperCamelCase	,			__UpperCamelCase   )
																__magic_name__				=     {k: torch.tensor(v.tolist()   ) for k, v in prepared_inputs_dict.items()}

																# load corresponding PyTorch class
																__magic_name__				=     model_class.__name__[4:]  # Skip the "Flax" at the beginning
																__magic_name__				=     getattr(__UpperCamelCase	,			__UpperCamelCase   )

																__magic_name__,						__magic_name__				=     pt_inputs['input_ids'].shape
																__magic_name__				=     np.random.randint(0	,			seq_length - 1	,			size=(batch_size,)   )
																for batch_idx, start_index in enumerate(__UpperCamelCase   ):
																					__magic_name__				=     0
																					__magic_name__				=     1
																					__magic_name__				=     0
																					__magic_name__				=     1
																__magic_name__				=     pt_model_class(__UpperCamelCase   ).eval()
																__magic_name__				=     model_class(__UpperCamelCase	,			dtype=jnp.floataa   )

																__magic_name__				=     convert_pytorch_state_dict_to_flax(pt_model.state_dict()	,			__UpperCamelCase   )
																__magic_name__				=     fx_state

																with torch.no_grad():
																					__magic_name__				=     pt_model(**__UpperCamelCase   ).to_tuple()

																__magic_name__				=     fx_model(**__UpperCamelCase   ).to_tuple()
																self.assertEqual(len(__UpperCamelCase   )	,			len(__UpperCamelCase   )	,			'Output lengths differ between Flax and PyTorch'   )
																for fx_output, pt_output in zip(__UpperCamelCase	,			__UpperCamelCase   ):
																					self.assert_almost_equals(fx_output[:, -1]	,			pt_output[:, -1].numpy()	,			4E-2   )

																with tempfile.TemporaryDirectory() as tmpdirname:
																					pt_model.save_pretrained(__UpperCamelCase   )
																					__magic_name__				=     model_class.from_pretrained(__UpperCamelCase	,			from_pt=__UpperCamelCase   )

																__magic_name__				=     fx_model_loaded(**__UpperCamelCase   ).to_tuple()
																self.assertEqual(
																    len(__UpperCamelCase   )	,			len(__UpperCamelCase   )	,			'Output lengths differ between Flax and PyTorch'   )
																for fx_output_loaded, pt_output in zip(__UpperCamelCase	,			__UpperCamelCase   ):
																					self.assert_almost_equals(fx_output_loaded[:, -1]	,			pt_output[:, -1].numpy()	,			4E-2   )







	@is_pt_flax_cross_test
	def 						_SCREAMING_SNAKE_CASE				( self:			Union[str, Any]   ):

						'''simple docstring'''





						__magic_name__,						__magic_name__				=     self.model_tester.prepare_config_and_inputs_for_common()
						for model_class in self.all_model_classes:
											with self.subTest(model_class.__name__   ):
																# prepare inputs
																__magic_name__				=     self._prepare_for_class(__UpperCamelCase	,			__UpperCamelCase   )
																__magic_name__				=     {k: torch.tensor(v.tolist()   ) for k, v in prepared_inputs_dict.items()}

																# load corresponding PyTorch class
																__magic_name__				=     model_class.__name__[4:]  # Skip the "Flax" at the beginning
																__magic_name__				=     getattr(__UpperCamelCase	,			__UpperCamelCase   )

																__magic_name__				=     pt_model_class(__UpperCamelCase   ).eval()
																__magic_name__				=     model_class(__UpperCamelCase	,			dtype=jnp.floataa   )

																__magic_name__				=     load_flax_weights_in_pytorch_model(__UpperCamelCase	,			fx_model.params   )
																__magic_name__,						__magic_name__				=     pt_inputs['input_ids'].shape
																__magic_name__				=     np.random.randint(0	,			seq_length - 1	,			size=(batch_size,)   )
																for batch_idx, start_index in enumerate(__UpperCamelCase   ):
																					__magic_name__				=     0
																					__magic_name__				=     1
																					__magic_name__				=     0
																					__magic_name__				=     1

																# make sure weights are tied in PyTorch
																pt_model.tie_weights()

																with torch.no_grad():
																					__magic_name__				=     pt_model(**__UpperCamelCase   ).to_tuple()

																__magic_name__				=     fx_model(**__UpperCamelCase   ).to_tuple()
																self.assertEqual(len(__UpperCamelCase   )	,			len(__UpperCamelCase   )	,			'Output lengths differ between Flax and PyTorch'   )
																for fx_output, pt_output in zip(__UpperCamelCase	,			__UpperCamelCase   ):
																					self.assert_almost_equals(fx_output[:, -1]	,			pt_output[:, -1].numpy()	,			4E-2   )

																with tempfile.TemporaryDirectory() as tmpdirname:
																					fx_model.save_pretrained(__UpperCamelCase   )
																					__magic_name__				=     pt_model_class.from_pretrained(__UpperCamelCase	,			from_flax=__UpperCamelCase   )

																with torch.no_grad():
																					__magic_name__				=     pt_model_loaded(**__UpperCamelCase   ).to_tuple()

																self.assertEqual(
																    len(__UpperCamelCase   )	,			len(__UpperCamelCase   )	,			'Output lengths differ between Flax and PyTorch'   )
																for fx_output, pt_output in zip(__UpperCamelCase	,			__UpperCamelCase   ):
																					self.assert_almost_equals(fx_output[:, -1]	,			pt_output[:, -1].numpy()	,			4E-2   )







	@slow
	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[Any]   ):

						'''simple docstring'''





						for model_class_name in self.all_model_classes:
											__magic_name__				=     model_class_name.from_pretrained('gpt2'	,			from_pt=__UpperCamelCase   )
											__magic_name__				=     model(np.ones((1, 1)   )   )
											self.assertIsNotNone(__UpperCamelCase   )







from hashlib import mda, shaaaa

A__							=				(shaaaa, mda)

class 					__UpperCamelCase					:







	def __init__( self:			Optional[int]	,			__UpperCamelCase:			int = 8   ):

						'''simple docstring'''





						__magic_name__				=     0b0
						__magic_name__				=     size







	def 						_SCREAMING_SNAKE_CASE				( self:			Tuple	,			__UpperCamelCase:			str   ):

						'''simple docstring'''





						__magic_name__				=     self.hash_(__UpperCamelCase   )
						self.bitarray |= h







	def 						_SCREAMING_SNAKE_CASE				( self:			List[str]	,			__UpperCamelCase:			str   ):

						'''simple docstring'''





						__magic_name__				=     self.hash_(__UpperCamelCase   )
						return (h & self.bitarray) == h







	def __contains__( self:			Optional[int]	,			__UpperCamelCase:			str   ):

						'''simple docstring'''





						return self.exists(__UpperCamelCase   )







	def 						_SCREAMING_SNAKE_CASE				( self:			int	,			__UpperCamelCase:			int   ):

						'''simple docstring'''





						__magic_name__				=     bin(__UpperCamelCase   )[2:]
						return res.zfill(self.size   )







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			str   ):

						'''simple docstring'''





						return self.format_bin(self.bitarray   )







	def 						_SCREAMING_SNAKE_CASE				( self:			str	,			__UpperCamelCase:			str   ):

						'''simple docstring'''





						__magic_name__				=     0b0
						for func in HASH_FUNCTIONS:
											__magic_name__				=     (
											    int.from_bytes(func(value.encode()   ).digest()	,			'little'   ) % self.size
											)
											res |= 2**position
						return res







	def 						_SCREAMING_SNAKE_CASE				( self:			Dict	,			__UpperCamelCase:			str   ):

						'''simple docstring'''





						return self.format_bin(self.hash_(__UpperCamelCase   )   )







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[int]   ):

						'''simple docstring'''





						__magic_name__				=     bin(self.bitarray   ).count('1'   )
						return (n_ones / self.size) ** len(__UpperCamelCase   )




