



'''simple docstring'''







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		:  Dict   ,		__A		:  int   ,		__A		:  Tuple=14   ,		__A		:  int=7   ,		__A		:  Optional[Any]=True   ,		__A		:  List[Any]=True   ,		__A		:  Any=False   ,		__A		:  Optional[int]=True   ,		__A		:  Dict=99   ,		__A		:  Any=32   ,		__A		:  int=5   ,		__A		:  int=4   ,		__A		:  Optional[Any]=37   ,		__A		:  Optional[int]="gelu"   ,		__A		:  int=0.1   ,		__A		:  int=0.1   ,		__A		:  str=512   ,		__A		:  Union[str, Any]=0.0_2   ,		)     ->			Optional[int]:




								'''simple docstring'''





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




	def      lowercase__							(  self		:  Dict					)     ->			Union[str, Any]:




								'''simple docstring'''





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

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

								lowerCAmelCase__  =   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      lowercase__							(  self		:  Dict					)     ->			List[str]:




								'''simple docstring'''





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




	def      lowercase__							(  self		:  str					)     ->			Optional[int]:




								'''simple docstring'''





								lowerCAmelCase__  ,lowerCAmelCase__  ,lowerCAmelCase__  =   self.prepare_config_and_inputs()

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

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




	def      lowercase__							(  self		:  Tuple   ,		__A		:  List[str]   ,		__A		:  str   ,		__A		:  List[str]   ,		__A		:  Optional[Any]					)     ->			Optional[int]:




								'''simple docstring'''





								lowerCAmelCase__  =   20
								lowerCAmelCase__  =   model_class_name(__A					)

								lowerCAmelCase__  =   model.init_cache(input_ids.shape[0]   ,		__A					)
								lowerCAmelCase__  =   jnp.ones((input_ids.shape[0], max_decoder_length)   ,		dtype="""i4"""					)

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

								lowerCAmelCase__  =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]   ,		dtype="""i4"""					)
								lowerCAmelCase__  =   model(
								    input_ids[:, -1:]   ,		attention_mask=__A   ,		past_key_values=outputs_cache.past_key_values   ,		position_ids=__A   ,		)

								lowerCAmelCase__  =   model(__A					)

								lowerCAmelCase__  =   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      lowercase__							(  self		:  Dict   ,		__A		:  str   ,		__A		:  Optional[Any]   ,		__A		:  Tuple   ,		__A		:  Tuple					)     ->			Optional[int]:




								'''simple docstring'''





								lowerCAmelCase__  =   20
								lowerCAmelCase__  =   model_class_name(__A					)

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

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

								lowerCAmelCase__  =   model(
								    input_ids[:, :-1]   ,		attention_mask=__A   ,		past_key_values=__A   ,		position_ids=__A   ,		)
								lowerCAmelCase__  =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]   ,		dtype="""i4"""					)
								lowerCAmelCase__  =   model(
								    input_ids[:, -1:]   ,		past_key_values=outputs_cache.past_key_values   ,		attention_mask=__A   ,		position_ids=__A   ,		)

								lowerCAmelCase__  =   model(__A   ,		attention_mask=__A					)

								lowerCAmelCase__  =   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__     (     _A,     _A,     unittest.TestCase				):






	'''simple docstring'''

	A__		=	(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
	A__		=	(FlaxGPTaLMHeadModel,) if is_flax_available() else ()




	def      lowercase__							(  self		:  List[str]					)     ->			int:




								'''simple docstring'''





								lowerCAmelCase__  =   FlaxGPTaModelTester(self					)




	def      lowercase__							(  self		:  Optional[int]					)     ->			Any:




								'''simple docstring'''





								for model_class_name in self.all_model_classes:
															lowerCAmelCase__  ,lowerCAmelCase__  ,lowerCAmelCase__  =   self.model_tester.prepare_config_and_inputs()
															self.model_tester.check_use_cache_forward(__A   ,		__A   ,		__A   ,		__A					)




	def      lowercase__							(  self		:  Any					)     ->			Union[str, Any]:




								'''simple docstring'''





								for model_class_name in self.all_model_classes:
															lowerCAmelCase__  ,lowerCAmelCase__  ,lowerCAmelCase__  =   self.model_tester.prepare_config_and_inputs()
															self.model_tester.check_use_cache_forward_with_attn_mask(
															    __A   ,		__A   ,		__A   ,		__A					)




	@slow
	def      lowercase__							(  self		:  Any					)     ->			List[str]:




								'''simple docstring'''





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

								lowerCAmelCase__  =   FlaxGPTaLMHeadModel.from_pretrained("""gpt2"""					)
								lowerCAmelCase__  =   False
								lowerCAmelCase__  =   model.config.eos_token_id

								lowerCAmelCase__  =   jax.jit(model.generate					)

								lowerCAmelCase__  =   jit_generate(inputs["""input_ids"""]   ,		attention_mask=inputs["""attention_mask"""]					).sequences

								lowerCAmelCase__  =   tokenizer.batch_decode(__A   ,		skip_special_tokens=__A					)

								lowerCAmelCase__  =   [
								    """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      lowercase__							(  self		:  Union[str, Any]					)     ->			Union[str, Any]:




								'''simple docstring'''





								lowerCAmelCase__  ,lowerCAmelCase__  =   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
																						lowerCAmelCase__  =   self._prepare_for_class(__A   ,		__A					)
																						lowerCAmelCase__  =   {k: torch.tensor(v.tolist()					) for k, v in prepared_inputs_dict.items()}

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

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

																						lowerCAmelCase__  =   convert_pytorch_state_dict_to_flax(pt_model.state_dict()   ,		__A					)
																						lowerCAmelCase__  =   fx_state

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

																						lowerCAmelCase__  =   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					)
																													lowerCAmelCase__  =   model_class.from_pretrained(__A   ,		from_pt=__A					)

																						lowerCAmelCase__  =   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      lowercase__							(  self		:  Union[str, Any]					)     ->			List[Any]:




								'''simple docstring'''





								lowerCAmelCase__  ,lowerCAmelCase__  =   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
																						lowerCAmelCase__  =   self._prepare_for_class(__A   ,		__A					)
																						lowerCAmelCase__  =   {k: torch.tensor(v.tolist()					) for k, v in prepared_inputs_dict.items()}

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

																						lowerCAmelCase__  =   pt_model_class(__A					).eval()
																						lowerCAmelCase__  =   model_class(__A   ,		dtype=jnp.floataa					)

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

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

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

																						lowerCAmelCase__  =   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					)
																													lowerCAmelCase__  =   pt_model_class.from_pretrained(__A   ,		from_flax=__A					)

																						with torch.no_grad():
																													lowerCAmelCase__  =   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      lowercase__							(  self		:  Tuple					)     ->			Any:




								'''simple docstring'''





								for model_class_name in self.all_model_classes:
															lowerCAmelCase__  =   model_class_name.from_pretrained("""gpt2"""   ,		from_pt=__A					)
															lowerCAmelCase__  =   model(np.ones((1, 1)					)					)
															self.assertIsNotNone(__A					)










'''simple docstring'''







def 			_lowerCAmelCase(      UpperCAmelCase_      :	list					, UpperCAmelCase_      :	int					, UpperCAmelCase_      :	int					, UpperCAmelCase_      :	int						) ->							int:
							lowerCAmelCase__  =   int((lower_bound + upper_bound) // 2						)
							if array[r] == value:
														return r
							if lower_bound >= upper_bound:
														return -1
							if array[r] < value:
														return binary_search(UpperCAmelCase_					, r + 1					, UpperCAmelCase_					, UpperCAmelCase_						)
							else:
														return binary_search(UpperCAmelCase_					, UpperCAmelCase_					, r - 1					, UpperCAmelCase_						)
def 			_lowerCAmelCase(      UpperCAmelCase_      :	int					, UpperCAmelCase_      :	list						) ->							list:
							lowerCAmelCase__  =   0
							if matrix[index][0] == value:
														return [index, 0]
							while index < len(UpperCAmelCase_						) and matrix[index][0] < value:
														lowerCAmelCase__  =   binary_search(matrix[index]					, 0					, len(matrix[index]						) - 1					, UpperCAmelCase_						)
														if r != -1:
																					return [index, r]
														index += 1
							return [-1, -1]


if __name__ == "__main__":
		import doctest

		doctest.testmod()





