




'''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  A		:




				'''simple docstring'''





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

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

							__UpperCamelCase				:      Union[str, Any]									=       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 	a_	(self      )     ->	List[str]:
							__UpperCamelCase				:      str									=       self.prepare_config_and_inputs()
							__UpperCamelCase		, __UpperCamelCase		, __UpperCamelCase				:      List[str]									=       config_and_inputs
							__UpperCamelCase				:      Optional[Any]									=       {"input_ids": input_ids, "attention_mask": attention_mask}
							return config, inputs_dict
				def 	a_	(self      )     ->	Dict:
							__UpperCamelCase		, __UpperCamelCase		, __UpperCamelCase				:      Dict									=       self.prepare_config_and_inputs()

							__UpperCamelCase				:      int									=       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
							__UpperCamelCase				:      Any									=       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			,  _UpperCAmelCase			,  _UpperCAmelCase			,  _UpperCAmelCase			,  _UpperCAmelCase      )     ->	Optional[int]:
							__UpperCamelCase				:      int									=       2_0
							__UpperCamelCase				:      int									=       model_class_name(_UpperCAmelCase      )

							__UpperCamelCase				:      Any									=       model.init_cache(input_ids.shape[0]			,  _UpperCAmelCase      )
							__UpperCamelCase				:      Optional[Any]									=       jnp.ones((input_ids.shape[0], max_decoder_length)			,  dtype="i4"      )

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

							__UpperCamelCase				:      Tuple									=       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]			,  dtype="i4"      )
							__UpperCamelCase				:      Optional[int]									=       model(
							    input_ids[:, -1:]			,  attention_mask=_UpperCAmelCase			,  past_key_values=outputs_cache.past_key_values			,  position_ids=_UpperCAmelCase			,  )

							__UpperCamelCase				:      Optional[int]									=       model(_UpperCAmelCase      )

							__UpperCamelCase				:      Optional[Any]									=       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			,  _UpperCAmelCase			,  _UpperCAmelCase			,  _UpperCAmelCase			,  _UpperCAmelCase      )     ->	int:
							__UpperCamelCase				:      List[str]									=       2_0
							__UpperCamelCase				:      Any									=       model_class_name(_UpperCAmelCase      )

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

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

							__UpperCamelCase				:      Any									=       model(
							    input_ids[:, :-1]			,  attention_mask=_UpperCAmelCase			,  past_key_values=_UpperCAmelCase			,  position_ids=_UpperCAmelCase			,  )
							__UpperCamelCase				:      List[str]									=       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]			,  dtype="i4"      )
							__UpperCamelCase				:      List[str]									=       model(
							    input_ids[:, -1:]			,  past_key_values=outputs_cache.past_key_values			,  attention_mask=_UpperCAmelCase			,  position_ids=_UpperCAmelCase			,  )

							__UpperCamelCase				:      Optional[int]									=       model(_UpperCAmelCase			,  attention_mask=_UpperCAmelCase      )

							__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}"      )







@require_flax
class  A		(			SCREAMING_SNAKE_CASE__      ,							SCREAMING_SNAKE_CASE__      ,							unittest.TestCase		):




				'''simple docstring'''





				A				     =    (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
				A				     =    (FlaxGPTaLMHeadModel,) if is_flax_available() else ()
				def 	a_	(self      )     ->	List[str]:
							__UpperCamelCase				:      Optional[int]									=       FlaxGPTaModelTester(self      )
				def 	a_	(self      )     ->	Dict:
							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(_UpperCAmelCase			,  _UpperCAmelCase			,  _UpperCAmelCase			,  _UpperCAmelCase      )
				def 	a_	(self      )     ->	List[Any]:
							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_with_attn_mask(
										    _UpperCAmelCase			,  _UpperCAmelCase			,  _UpperCAmelCase			,  _UpperCAmelCase      )
				@slow
				def 	a_	(self      )     ->	int:
							__UpperCamelCase				:      Optional[Any]									=       GPTaTokenizer.from_pretrained("gpt2"			,  pad_token="</s>"			,  padding_side="left"      )
							__UpperCamelCase				:      List[str]									=       tokenizer(["Hello this is a long string", "Hey"]			,  return_tensors="np"			,  padding=_UpperCAmelCase			,  truncation=_UpperCAmelCase      )

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

							__UpperCamelCase				:      List[Any]									=       jax.jit(model.generate      )

							__UpperCamelCase				:      Any									=       jit_generate(inputs["input_ids"]			,  attention_mask=inputs["attention_mask"]      ).sequences

							__UpperCamelCase				:      Tuple									=       tokenizer.batch_decode(_UpperCAmelCase			,  skip_special_tokens=_UpperCAmelCase      )

							__UpperCamelCase				:      Dict									=       [
							    "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 	a_	(self      )     ->	Tuple:
							__UpperCamelCase		, __UpperCamelCase				:      int									=       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				:      List[Any]									=       self._prepare_for_class(_UpperCAmelCase			,  _UpperCAmelCase      )
													__UpperCamelCase				:      Tuple									=       {k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

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

													__UpperCamelCase		, __UpperCamelCase				:      Dict									=       pt_inputs["input_ids"].shape
													__UpperCamelCase				:      List[Any]									=       np.random.randint(0			,  seq_length - 1			,  size=(batch_size,)      )
													for batch_idx, start_index in enumerate(_UpperCAmelCase      ):
																__UpperCamelCase				:      List[Any]									=       0
																__UpperCamelCase				:      Optional[int]									=       1
																__UpperCamelCase				:      Tuple									=       0
																__UpperCamelCase				:      Any									=       1
													__UpperCamelCase				:      List[str]									=       pt_model_class(_UpperCAmelCase      ).eval()
													__UpperCamelCase				:      Optional[Any]									=       model_class(_UpperCAmelCase			,  dtype=jnp.floataa      )

													__UpperCamelCase				:      List[str]									=       convert_pytorch_state_dict_to_flax(pt_model.state_dict()			,  _UpperCAmelCase      )
													__UpperCamelCase				:      int									=       fx_state

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

													__UpperCamelCase				:      str									=       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      )
																__UpperCamelCase				:      str									=       model_class.from_pretrained(_UpperCAmelCase			,  from_pt=_UpperCAmelCase      )

													__UpperCamelCase				:      Any									=       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 	a_	(self      )     ->	int:
							__UpperCamelCase		, __UpperCamelCase				:      Dict									=       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				:      Optional[Any]									=       self._prepare_for_class(_UpperCAmelCase			,  _UpperCAmelCase      )
													__UpperCamelCase				:      List[str]									=       {k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

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

													__UpperCamelCase				:      Any									=       pt_model_class(_UpperCAmelCase      ).eval()
													__UpperCamelCase				:      Optional[int]									=       model_class(_UpperCAmelCase			,  dtype=jnp.floataa      )

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

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

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

													__UpperCamelCase				:      Dict									=       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      )
																__UpperCamelCase				:      str									=       pt_model_class.from_pretrained(_UpperCAmelCase			,  from_flax=_UpperCAmelCase      )

													with torch.no_grad():
																__UpperCamelCase				:      Any									=       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 	a_	(self      )     ->	List[str]:
							for model_class_name in self.all_model_classes:
										__UpperCamelCase				:      Dict									=       model_class_name.from_pretrained("gpt2"			,  from_pt=_UpperCAmelCase      )
										__UpperCamelCase				:      Union[str, Any]									=       model(np.ones((1, 1)      )      )
										self.assertIsNotNone(_UpperCAmelCase      )









'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_lowerCAmelCase				     = {'''configuration_visual_bert''': ['''VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''VisualBertConfig''']}

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_lowerCAmelCase				     = [
				    '''VISUAL_BERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
				    '''VisualBertForMultipleChoice''',
				    '''VisualBertForPreTraining''',
				    '''VisualBertForQuestionAnswering''',
				    '''VisualBertForRegionToPhraseAlignment''',
				    '''VisualBertForVisualReasoning''',
				    '''VisualBertLayer''',
				    '''VisualBertModel''',
				    '''VisualBertPreTrainedModel''',
				]


if TYPE_CHECKING:
				from .configuration_visual_bert import VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, VisualBertConfig

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_visual_bert import (
								    VISUAL_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
								    VisualBertForMultipleChoice,
								    VisualBertForPreTraining,
								    VisualBertForQuestionAnswering,
								    VisualBertForRegionToPhraseAlignment,
								    VisualBertForVisualReasoning,
								    VisualBertLayer,
								    VisualBertModel,
								    VisualBertPreTrainedModel,
								)


else:
				import sys

				_lowerCAmelCase				     = _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



