
'''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       :			Tuple   ,     lowerCAmelCase       :			Any   ,     lowerCAmelCase       :			Optional[Any]=14   ,     lowerCAmelCase       :			Any=7   ,     lowerCAmelCase       :			List[str]=True   ,     lowerCAmelCase       :			int=True   ,     lowerCAmelCase       :			str=False   ,     lowerCAmelCase       :			str=True   ,     lowerCAmelCase       :			Optional[int]=99   ,     lowerCAmelCase       :			Optional[int]=32   ,     lowerCAmelCase       :			Dict=5   ,     lowerCAmelCase       :			int=4   ,     lowerCAmelCase       :			List[str]=37   ,     lowerCAmelCase       :			List[Any]="gelu"   ,     lowerCAmelCase       :			Any=0.1   ,     lowerCAmelCase       :			Optional[int]=0.1   ,     lowerCAmelCase       :			List[Any]=5_12   ,     lowerCAmelCase       :			List[str]=0.0_2   ,     ):
													A_ 						=		parent
													A_ 						=		batch_size
													A_ 						=		seq_length
													A_ 						=		is_training
													A_ 						=		use_input_mask
													A_ 						=		use_token_type_ids
													A_ 						=		use_labels
													A_ 						=		vocab_size
													A_ 						=		hidden_size
													A_ 						=		num_hidden_layers
													A_ 						=		num_attention_heads
													A_ 						=		intermediate_size
													A_ 						=		hidden_act
													A_ 						=		hidden_dropout_prob
													A_ 						=		attention_probs_dropout_prob
													A_ 						=		max_position_embeddings
													A_ 						=		initializer_range
													A_ 						=		None
													A_ 						=		vocab_size - 1
													A_ 						=		vocab_size - 1
													A_ 						=		vocab_size - 1






						def 							_UpperCAmelCase					(     self       :			str  ):
													A_ 						=		ids_tensor([self.batch_size, self.seq_length]   ,     self.vocab_size  )

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

													A_ 						=		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=lowerCAmelCase   ,     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 							_UpperCAmelCase					(     self       :			List[Any]  ):
													A_ 						=		self.prepare_config_and_inputs()
													A_						,				A_						,				A_ 						=		config_and_inputs
													A_ 						=		{"input_ids": input_ids, "attention_mask": attention_mask}
													return config, inputs_dict






						def 							_UpperCAmelCase					(     self       :			Any  ):
													A_						,				A_						,				A_ 						=		self.prepare_config_and_inputs()

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

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






						def 							_UpperCAmelCase					(     self       :			Optional[Any]   ,     lowerCAmelCase       :			Dict   ,     lowerCAmelCase       :			Optional[Any]   ,     lowerCAmelCase       :			Optional[Any]   ,     lowerCAmelCase       :			Tuple  ):
													A_ 						=		20
													A_ 						=		model_class_name(lowerCAmelCase  )

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

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

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

													A_ 						=		model(lowerCAmelCase  )

													A_ 						=		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 							_UpperCAmelCase					(     self       :			List[Any]   ,     lowerCAmelCase       :			Optional[Any]   ,     lowerCAmelCase       :			Optional[Any]   ,     lowerCAmelCase       :			Union[str, Any]   ,     lowerCAmelCase       :			Any  ):
													A_ 						=		20
													A_ 						=		model_class_name(lowerCAmelCase  )

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

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

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

													A_ 						=		model(lowerCAmelCase   ,     attention_mask=lowerCAmelCase  )

													A_ 						=		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      (      lowercase							, lowercase							, unittest.TestCase	):

						"""simple docstring"""






						_UpperCAmelCase		:					List[Any]      	=(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
						_UpperCAmelCase		:					Optional[int]      	=(FlaxGPTaLMHeadModel,) if is_flax_available() else ()






						def 							_UpperCAmelCase					(     self       :			Union[str, Any]  ):
													A_ 						=		FlaxGPTaModelTester(self  )






						def 							_UpperCAmelCase					(     self       :			Optional[Any]  ):
													for model_class_name in self.all_model_classes:
																				A_						,				A_						,				A_ 						=		self.model_tester.prepare_config_and_inputs()
																				self.model_tester.check_use_cache_forward(lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase  )






						def 							_UpperCAmelCase					(     self       :			Any  ):
													for model_class_name in self.all_model_classes:
																				A_						,				A_						,				A_ 						=		self.model_tester.prepare_config_and_inputs()
																				self.model_tester.check_use_cache_forward_with_attn_mask(
																				    lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase  )






						@slow
						def 							_UpperCAmelCase					(     self       :			str  ):
													A_ 						=		GPTaTokenizer.from_pretrained("gpt2"   ,     pad_token="</s>"   ,     padding_side="left"  )
													A_ 						=		tokenizer(["Hello this is a long string", "Hey"]   ,     return_tensors="np"   ,     padding=lowerCAmelCase   ,     truncation=lowerCAmelCase  )

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

													A_ 						=		jax.jit(model.generate  )

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

													A_ 						=		tokenizer.batch_decode(lowerCAmelCase   ,     skip_special_tokens=lowerCAmelCase  )

													A_ 						=		[
													    "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(lowerCAmelCase   ,     lowerCAmelCase  )






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

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

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

																											A_ 						=		convert_pytorch_state_dict_to_flax(pt_model.state_dict()   ,     lowerCAmelCase  )
																											A_ 						=		fx_state

																											with torch.no_grad():
																																		A_ 						=		pt_model(**lowerCAmelCase  ).to_tuple()

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

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

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






						@is_pt_flax_cross_test
						def 							_UpperCAmelCase					(     self       :			Optional[int]  ):
													A_						,				A_ 						=		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
																											A_ 						=		self._prepare_for_class(lowerCAmelCase   ,     lowerCAmelCase  )
																											A_ 						=		{k: torch.tensor(v.tolist()  ) for k, v in prepared_inputs_dict.items()}

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

																											A_ 						=		pt_model_class(lowerCAmelCase  ).eval()
																											A_ 						=		model_class(lowerCAmelCase   ,     dtype=jnp.floataa  )

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

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

																											with torch.no_grad():
																																		A_ 						=		pt_model(**lowerCAmelCase  ).to_tuple()

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

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

																											with torch.no_grad():
																																		A_ 						=		pt_model_loaded(**lowerCAmelCase  ).to_tuple()

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






						@slow
						def 							_UpperCAmelCase					(     self       :			str  ):
													for model_class_name in self.all_model_classes:
																				A_ 						=		model_class_name.from_pretrained("gpt2"   ,     from_pt=lowerCAmelCase  )
																				A_ 						=		model(np.ones((1, 1)  )  )
																				self.assertIsNotNone(lowerCAmelCase  )









'''simple docstring'''
import unittest

import numpy as np

from transformers import is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
							import tensorflow as tf

							from transformers.activations_tf import get_tf_activation






@require_tf
class 				__lowerCAmelCase      (      unittest.TestCase	):

						"""simple docstring"""






						def 							_UpperCAmelCase					(     self       :			List[Any]  ):
													A_ 						=		tf.constant([-1_00, -1.0, -0.1, 0, 0.1, 1.0, 1_0_0.0]  )
													A_ 						=		get_tf_activation("gelu"  )
													A_ 						=		get_tf_activation("gelu_10"  )

													A_ 						=		gelu(lowerCAmelCase  )
													A_ 						=		geluaa(lowerCAmelCase  )

													A_ 						=		tf.where(y_gelu_aa < 1_0.0   ,     1.0   ,     0.0  )

													self.assertEqual(tf.math.reduce_max(lowerCAmelCase  ).numpy().item()   ,     1_0.0  )
													self.assertTrue(np.allclose(y_gelu * clipped_mask   ,     y_gelu_aa * clipped_mask  )  )






						def 							_UpperCAmelCase					(     self       :			Union[str, Any]  ):
													get_tf_activation("gelu"  )
													get_tf_activation("gelu_10"  )
													get_tf_activation("gelu_fast"  )
													get_tf_activation("gelu_new"  )
													get_tf_activation("glu"  )
													get_tf_activation("mish"  )
													get_tf_activation("quick_gelu"  )
													get_tf_activation("relu"  )
													get_tf_activation("sigmoid"  )
													get_tf_activation("silu"  )
													get_tf_activation("swish"  )
													get_tf_activation("tanh"  )
													with self.assertRaises(lowerCAmelCase  ):
																				get_tf_activation("bogus"  )
													with self.assertRaises(lowerCAmelCase  ):
																				get_tf_activation(lowerCAmelCase  )





