






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 				_lowerCAmelCase:





       """simple docstring"""







       def __init__(					self  ,			_lowerCamelCase  ,			_lowerCamelCase=1_4  ,			_lowerCamelCase=7  ,			_lowerCamelCase=True  ,			_lowerCamelCase=True  ,			_lowerCamelCase=False  ,			_lowerCamelCase=True  ,			_lowerCamelCase=9_9  ,			_lowerCamelCase=3_2  ,			_lowerCamelCase=5  ,			_lowerCamelCase=4  ,			_lowerCamelCase=3_7  ,			_lowerCamelCase="gelu"  ,			_lowerCamelCase=0.1  ,			_lowerCamelCase=0.1  ,			_lowerCamelCase=5_1_2  ,			_lowerCamelCase=0.0_2  ,			):
         UpperCamelCase_:					str          =				parent
         UpperCamelCase_:					Dict          =				batch_size
         UpperCamelCase_:					int          =				seq_length
         UpperCamelCase_:					List[str]          =				is_training
         UpperCamelCase_:					Tuple          =				use_input_mask
         UpperCamelCase_:					Optional[int]          =				use_token_type_ids
         UpperCamelCase_:					Optional[int]          =				use_labels
         UpperCamelCase_:					List[Any]          =				vocab_size
         UpperCamelCase_:					Dict          =				hidden_size
         UpperCamelCase_:					Dict          =				num_hidden_layers
         UpperCamelCase_:					List[Any]          =				num_attention_heads
         UpperCamelCase_:					int          =				intermediate_size
         UpperCamelCase_:					Optional[Any]          =				hidden_act
         UpperCamelCase_:					Union[str, Any]          =				hidden_dropout_prob
         UpperCamelCase_:					Optional[int]          =				attention_probs_dropout_prob
         UpperCamelCase_:					Any          =				max_position_embeddings
         UpperCamelCase_:					Any          =				initializer_range
         UpperCamelCase_:					List[str]          =				None
         UpperCamelCase_:					List[str]          =				vocab_size - 1
         UpperCamelCase_:					Tuple          =				vocab_size - 1
         UpperCamelCase_:					int          =				vocab_size - 1





       def 			_a    (					self							):
         UpperCamelCase_:					List[Any]          =				ids_tensor([self.batch_size, self.seq_length]  ,			self.vocab_size							)

         UpperCamelCase_:					Optional[int]          =				None
         if self.use_input_mask:
           UpperCamelCase_:					Optional[Any]          =				random_attention_mask([self.batch_size, self.seq_length]							)

         UpperCamelCase_:					Tuple          =				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=__A  ,			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 			_a    (					self							):
         UpperCamelCase_:					int          =				self.prepare_config_and_inputs()
         UpperCamelCase_      ,UpperCamelCase_      ,UpperCamelCase_:					int          =				config_and_inputs
         UpperCamelCase_:					Tuple          =				{'input_ids': input_ids, 'attention_mask': attention_mask}
         return config, inputs_dict





       def 			_a    (					self							):
         UpperCamelCase_      ,UpperCamelCase_      ,UpperCamelCase_:					List[Any]          =				self.prepare_config_and_inputs()

         UpperCamelCase_:					int          =				floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
         UpperCamelCase_:					int          =				ids_tensor([self.batch_size, self.seq_length]  ,			vocab_size=2							)

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





       def 			_a    (					self  ,			_lowerCamelCase  ,			_lowerCamelCase  ,			_lowerCamelCase  ,			_lowerCamelCase							):
         UpperCamelCase_:					int          =				2_0
         UpperCamelCase_:					List[Any]          =				model_class_name(__A							)

         UpperCamelCase_:					Tuple          =				model.init_cache(input_ids.shape[0]  ,			__A							)
         UpperCamelCase_:					int          =				jnp.ones((input_ids.shape[0], max_decoder_length)  ,			dtype='i4'							)

         UpperCamelCase_:					Optional[Any]          =				jnp.broadcast_to(
             jnp.arange(input_ids.shape[-1] - 1							)[None, :]  ,			(input_ids.shape[0], input_ids.shape[-1] - 1)							)
         UpperCamelCase_:					int          =				model(
             input_ids[:, :-1]  ,			attention_mask=__A  ,			past_key_values=__A  ,			position_ids=__A  ,			)

         UpperCamelCase_:					Union[str, Any]          =				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]  ,			dtype='i4'							)
         UpperCamelCase_:					Optional[Any]          =				model(
             input_ids[:, -1:]  ,			attention_mask=__A  ,			past_key_values=outputs_cache.past_key_values  ,			position_ids=__A  ,			)

         UpperCamelCase_:					Any          =				model(__A							)

         UpperCamelCase_:					Dict          =				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 			_a    (					self  ,			_lowerCamelCase  ,			_lowerCamelCase  ,			_lowerCamelCase  ,			_lowerCamelCase							):
         UpperCamelCase_:					List[Any]          =				2_0
         UpperCamelCase_:					List[Any]          =				model_class_name(__A							)

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

         UpperCamelCase_:					Optional[int]          =				model.init_cache(input_ids.shape[0]  ,			__A							)
         UpperCamelCase_:					List[Any]          =				jnp.broadcast_to(
             jnp.arange(input_ids.shape[-1] - 1							)[None, :]  ,			(input_ids.shape[0], input_ids.shape[-1] - 1)							)

         UpperCamelCase_:					Optional[int]          =				model(
             input_ids[:, :-1]  ,			attention_mask=__A  ,			past_key_values=__A  ,			position_ids=__A  ,			)
         UpperCamelCase_:					str          =				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]  ,			dtype='i4'							)
         UpperCamelCase_:					Union[str, Any]          =				model(
             input_ids[:, -1:]  ,			past_key_values=outputs_cache.past_key_values  ,			attention_mask=__A  ,			position_ids=__A  ,			)

         UpperCamelCase_:					Tuple          =				model(__A  ,			attention_mask=__A							)

         UpperCamelCase_:					str          =				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 				_lowerCAmelCase(       UpperCAmelCase_							,	UpperCAmelCase_							,	unittest.TestCase					):





       """simple docstring"""







       a				:		Optional[Any]									=(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
       a				:		str									=(FlaxGPTaLMHeadModel,) if is_flax_available() else ()





       def 			_a    (					self							):
         UpperCamelCase_:					Optional[int]          =				FlaxGPTaModelTester(self							)





       def 			_a    (					self							):
         for model_class_name in self.all_model_classes:
           UpperCamelCase_      ,UpperCamelCase_      ,UpperCamelCase_:					Tuple          =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.check_use_cache_forward(__A  ,			__A  ,			__A  ,			__A							)





       def 			_a    (					self							):
         for model_class_name in self.all_model_classes:
           UpperCamelCase_      ,UpperCamelCase_      ,UpperCamelCase_:					str          =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.check_use_cache_forward_with_attn_mask(
               __A  ,			__A  ,			__A  ,			__A							)





       @slow
       def 			_a    (					self							):
         UpperCamelCase_:					Optional[Any]          =				GPTaTokenizer.from_pretrained('gpt2'  ,			pad_token='</s>'  ,			padding_side='left'							)
         UpperCamelCase_:					Optional[Any]          =				tokenizer(['Hello this is a long string', 'Hey']  ,			return_tensors='np'  ,			padding=__A  ,			truncation=__A							)

         UpperCamelCase_:					Dict          =				FlaxGPTaLMHeadModel.from_pretrained('gpt2'							)
         UpperCamelCase_:					Optional[int]          =				False
         UpperCamelCase_:					List[Any]          =				model.config.eos_token_id

         UpperCamelCase_:					Union[str, Any]          =				jax.jit(model.generate							)

         UpperCamelCase_:					Dict          =				jit_generate(inputs['input_ids']  ,			attention_mask=inputs['attention_mask']							).sequences

         UpperCamelCase_:					int          =				tokenizer.batch_decode(__A  ,			skip_special_tokens=__A							)

         UpperCamelCase_:					Optional[Any]          =				[
             '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(__A  ,			__A							)





       @is_pt_flax_cross_test
       def 			_a    (					self							):
         UpperCamelCase_      ,UpperCamelCase_:					Any          =				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
             UpperCamelCase_:					Dict          =				self._prepare_for_class(__A  ,			__A							)
             UpperCamelCase_:					str          =				{k: torch.tensor(v.tolist()							) for k, v in prepared_inputs_dict.items()}

             # load corresponding PyTorch class
             UpperCamelCase_:					Optional[int]          =				model_class.__name__[4:]  # Skip the "Flax" at the beginning
             UpperCamelCase_:					Union[str, Any]          =				getattr(__A  ,			__A							)

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

             UpperCamelCase_:					Union[str, Any]          =				convert_pytorch_state_dict_to_flax(pt_model.state_dict()  ,			__A							)
             UpperCamelCase_:					Optional[int]          =				fx_state

             with torch.no_grad():
               UpperCamelCase_:					Optional[int]          =				pt_model(**__A							).to_tuple()

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

             with tempfile.TemporaryDirectory() as tmpdirname:
               pt_model.save_pretrained(__A							)
               UpperCamelCase_:					str          =				model_class.from_pretrained(__A  ,			from_pt=__A							)

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





       @is_pt_flax_cross_test
       def 			_a    (					self							):
         UpperCamelCase_      ,UpperCamelCase_:					Any          =				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
             UpperCamelCase_:					str          =				self._prepare_for_class(__A  ,			__A							)
             UpperCamelCase_:					str          =				{k: torch.tensor(v.tolist()							) for k, v in prepared_inputs_dict.items()}

             # load corresponding PyTorch class
             UpperCamelCase_:					Union[str, Any]          =				model_class.__name__[4:]  # Skip the "Flax" at the beginning
             UpperCamelCase_:					List[Any]          =				getattr(__A  ,			__A							)

             UpperCamelCase_:					Optional[int]          =				pt_model_class(__A							).eval()
             UpperCamelCase_:					Optional[int]          =				model_class(__A  ,			dtype=jnp.floataa							)

             UpperCamelCase_:					Any          =				load_flax_weights_in_pytorch_model(__A  ,			fx_model.params							)
             UpperCamelCase_      ,UpperCamelCase_:					Optional[int]          =				pt_inputs['input_ids'].shape
             UpperCamelCase_:					Any          =				np.random.randint(0  ,			seq_length - 1  ,			size=(batch_size,)							)
             for batch_idx, start_index in enumerate(__A							):
               UpperCamelCase_:					Tuple          =				0
               UpperCamelCase_:					int          =				1
               UpperCamelCase_:					Optional[Any]          =				0
               UpperCamelCase_:					str          =				1

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

             with torch.no_grad():
               UpperCamelCase_:					str          =				pt_model(**__A							).to_tuple()

             UpperCamelCase_:					Union[str, Any]          =				fx_model(**__A							).to_tuple()
             self.assertEqual(len(__A							)  ,			len(__A							)  ,			'Output lengths differ between Flax and PyTorch'							)
             for fx_output, pt_output in zip(__A  ,			__A							):
               self.assert_almost_equals(fx_output[:, -1]  ,			pt_output[:, -1].numpy()  ,			4e-2							)

             with tempfile.TemporaryDirectory() as tmpdirname:
               fx_model.save_pretrained(__A							)
               UpperCamelCase_:					int          =				pt_model_class.from_pretrained(__A  ,			from_flax=__A							)

             with torch.no_grad():
               UpperCamelCase_:					Dict          =				pt_model_loaded(**__A							).to_tuple()

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





       @slow
       def 			_a    (					self							):
         for model_class_name in self.all_model_classes:
           UpperCamelCase_:					Any          =				model_class_name.from_pretrained('gpt2'  ,			from_pt=__A							)
           UpperCamelCase_:					List[str]          =				model(np.ones((1, 1)							)							)
           self.assertIsNotNone(__A							)







class   SCREAMING_SNAKE_CASE :





							"""simple docstring"""

							def __init__(     self:      Any      )			-> None:
												# numbermap keys are the number and its values are lists of indexes sorted
												# in ascending order
												_A									=       {}
												# indexmap keys are an index and it's values are the number at that index
												_A									=       {}

							def 					__A       (     self:      Any						,   __A:      list | str | range						,   __A:      int      )			-> list[int]:
												if isinstance(__A						,   (range, str)      ):
																	_A									=       list(__A      )
												elif not isinstance(__A						,   __A      ):
																	raise TypeError(
																	    '''binary_search_delete() only accepts either a list, range or str'''      )

												_A									=       0
												_A									=       len(__A      ) - 1

												while low <= high:
																	_A									=       (low + high) // 2
																	if array[mid] == item:
																						array.pop(__A      )
																						return array
																	elif array[mid] < item:
																						_A									=       mid + 1
																	else:
																						_A									=       mid - 1
												raise ValueError(
												    '''Either the item is not in the array or the array was unsorted'''      )

							def 					__A       (     self:      int						,   __A:      list | str | range						,   __A:      int      )			-> list[int]:
												if isinstance(__A						,   (range, str)      ):
																	_A									=       list(__A      )
												elif not isinstance(__A						,   __A      ):
																	raise TypeError(
																	    '''binary_search_insert() only accepts either a list, range or str'''      )

												_A									=       0
												_A									=       len(__A      ) - 1

												while low <= high:
																	_A									=       (low + high) // 2
																	if array[mid] == index:
																						# If the item already exists in the array,
																						# insert it after the existing item
																						array.insert(mid + 1						,   __A      )
																						return array
																	elif array[mid] < index:
																						_A									=       mid + 1
																	else:
																						_A									=       mid - 1

        # If the item doesn't exist in the array, insert it at the appropriate position
												array.insert(__A						,   __A      )
												return array

							def 					__A       (     self:      Union[str, Any]						,   __A:      int						,   __A:      int      )			-> None:
												# Remove previous index
												if index in self.indexmap:
																	_A									=       self.indexmap[index]
																	if len(self.numbermap[n]      ) == 1:
																						del self.numbermap[n]
																	else:
																						_A									=       self.binary_search_delete(self.numbermap[n]						,   __A      )

        # Set new index
												_A									=       number

												# Number not seen before or empty so insert number value
												if number not in self.numbermap:
																	_A									=       [index]

												# Here we need to perform a binary search insertion in order to insert
												# The item in the correct place
												else:
																	_A									=       self.binary_search_insert(
																	    self.numbermap[number]						,   __A      )



							def 					__A       (     self:      List[Any]						,   __A:      int      )			-> int:
												return self.numbermap.get(__A						,   [-1]      )[0]


if __name__ == "__main__":
				import doctest

				doctest.testmod()



