



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





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



							'''simple docstring'''







							snake_case   :							Optional[Any]						  =parent
							snake_case   :							Optional[Any]						  =batch_size
							snake_case   :							Union[str, Any]						  =seq_length
							snake_case   :							Any						  =is_training
							snake_case   :							List[Any]						  =use_input_mask
							snake_case   :							int						  =use_token_type_ids
							snake_case   :							Dict						  =use_labels
							snake_case   :							int						  =vocab_size
							snake_case   :							Union[str, Any]						  =hidden_size
							snake_case   :							int						  =num_hidden_layers
							snake_case   :							Optional[Any]						  =num_attention_heads
							snake_case   :							int						  =intermediate_size
							snake_case   :							Optional[Any]						  =hidden_act
							snake_case   :							Dict						  =hidden_dropout_prob
							snake_case   :							str						  =attention_probs_dropout_prob
							snake_case   :							List[str]						  =max_position_embeddings
							snake_case   :							Tuple						  =initializer_range
							snake_case   :							Dict						  =None
							snake_case   :							Any						  =vocab_size - 1
							snake_case   :							Dict						  =vocab_size - 1
							snake_case   :							Tuple						  =vocab_size - 1





				def 					__snake_case       (    self    :    Tuple     ):



							'''simple docstring'''







							snake_case   :							Tuple						  =ids_tensor([self.batch_size, self.seq_length],    self.vocab_size     )

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

							snake_case   :							str						  =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=_snake_case,    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 					__snake_case       (    self    :    Dict     ):



							'''simple docstring'''







							snake_case   :							Optional[int]						  =self.prepare_config_and_inputs()
							snake_case   :							Optional[int]						  =config_and_inputs
							snake_case   :							Union[str, Any]						  ={'''input_ids''': input_ids, '''attention_mask''': attention_mask}
							return config, inputs_dict





				def 					__snake_case       (    self    :    List[Any]     ):



							'''simple docstring'''







							snake_case   :							List[str]						  =self.prepare_config_and_inputs()

							snake_case   :							Any						  =floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
							snake_case   :							List[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 					__snake_case       (    self    :    List[Any],    _snake_case    :    List[Any],    _snake_case    :    Optional[int],    _snake_case    :    Optional[int],    _snake_case    :    List[Any]     ):



							'''simple docstring'''







							snake_case   :							int						  =20
							snake_case   :							str						  =model_class_name(_snake_case     )

							snake_case   :							List[str]						  =model.init_cache(input_ids.shape[0],    _snake_case     )
							snake_case   :							Dict						  =jnp.ones((input_ids.shape[0], max_decoder_length),    dtype='''i4'''     )

							snake_case   :							Optional[int]						  =jnp.broadcast_to(
							    jnp.arange(input_ids.shape[-1] - 1     )[None, :],    (input_ids.shape[0], input_ids.shape[-1] - 1)     )
							snake_case   :							str						  =model(
							    input_ids[:, :-1],    attention_mask=_snake_case,    past_key_values=_snake_case,    position_ids=_snake_case,    )

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

							snake_case   :							Dict						  =model(_snake_case     )

							snake_case   :							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 					__snake_case       (    self    :    Any,    _snake_case    :    Tuple,    _snake_case    :    str,    _snake_case    :    Tuple,    _snake_case    :    List[Any]     ):



							'''simple docstring'''







							snake_case   :							Optional[Any]						  =20
							snake_case   :							Tuple						  =model_class_name(_snake_case     )

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

							snake_case   :							Dict						  =model.init_cache(input_ids.shape[0],    _snake_case     )
							snake_case   :							List[str]						  =jnp.broadcast_to(
							    jnp.arange(input_ids.shape[-1] - 1     )[None, :],    (input_ids.shape[0], input_ids.shape[-1] - 1)     )

							snake_case   :							str						  =model(
							    input_ids[:, :-1],    attention_mask=_snake_case,    past_key_values=_snake_case,    position_ids=_snake_case,    )
							snake_case   :							Dict						  =jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]],    dtype='''i4'''     )
							snake_case   :							Tuple						  =model(
							    input_ids[:, -1:],    past_key_values=outputs_cache.past_key_values,    attention_mask=_snake_case,    position_ids=_snake_case,    )

							snake_case   :							Any						  =model(_snake_case,    attention_mask=_snake_case     )

							snake_case   :							int						  =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						):
				__UpperCAmelCase							=	(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
				__UpperCAmelCase							=	(FlaxGPTaLMHeadModel,) if is_flax_available() else ()





				def 					__snake_case       (    self    :    List[Any]     ):



							'''simple docstring'''







							snake_case   :							List[str]						  =FlaxGPTaModelTester(self     )





				def 					__snake_case       (    self    :    Dict     ):



							'''simple docstring'''







							for model_class_name in self.all_model_classes:
										snake_case   :							Any						  =self.model_tester.prepare_config_and_inputs()
										self.model_tester.check_use_cache_forward(_snake_case,    _snake_case,    _snake_case,    _snake_case     )





				def 					__snake_case       (    self    :    Any     ):



							'''simple docstring'''







							for model_class_name in self.all_model_classes:
										snake_case   :							str						  =self.model_tester.prepare_config_and_inputs()
										self.model_tester.check_use_cache_forward_with_attn_mask(
										    _snake_case,    _snake_case,    _snake_case,    _snake_case     )





				@slow
				def 					__snake_case       (    self    :    Any     ):



							'''simple docstring'''







							snake_case   :							Optional[int]						  =GPTaTokenizer.from_pretrained('''gpt2''',    pad_token='''</s>''',    padding_side='''left'''     )
							snake_case   :							List[str]						  =tokenizer(['''Hello this is a long string''', '''Hey'''],    return_tensors='''np''',    padding=_snake_case,    truncation=_snake_case     )

							snake_case   :							Optional[int]						  =FlaxGPTaLMHeadModel.from_pretrained('''gpt2'''     )
							snake_case   :							Tuple						  =False
							snake_case   :							int						  =model.config.eos_token_id

							snake_case   :							List[str]						  =jax.jit(model.generate     )

							snake_case   :							Optional[int]						  =jit_generate(inputs['''input_ids'''],    attention_mask=inputs['''attention_mask''']     ).sequences

							snake_case   :							str						  =tokenizer.batch_decode(_snake_case,    skip_special_tokens=_snake_case     )

							snake_case   :							Union[str, 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(_snake_case,    _snake_case     )





				@is_pt_flax_cross_test
				def 					__snake_case       (    self    :    Optional[Any]     ):



							'''simple docstring'''







							snake_case   :							Union[str, 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
													snake_case   :							Tuple						  =self._prepare_for_class(_snake_case,    _snake_case     )
													snake_case   :							Dict						  ={k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

													# load corresponding PyTorch class
													snake_case   :							Any						  =model_class.__name__[4:]  # Skip the "Flax" at the beginning
													snake_case   :							int						  =getattr(_snake_case,    _snake_case     )

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

													snake_case   :							str						  =convert_pytorch_state_dict_to_flax(pt_model.state_dict(),    _snake_case     )
													snake_case   :							List[Any]						  =fx_state

													with torch.no_grad():
																snake_case   :							Union[str, Any]						  =pt_model(**_snake_case     ).to_tuple()

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

													with tempfile.TemporaryDirectory() as tmpdirname:
																pt_model.save_pretrained(_snake_case     )
																snake_case   :							int						  =model_class.from_pretrained(_snake_case,    from_pt=_snake_case     )

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





				@is_pt_flax_cross_test
				def 					__snake_case       (    self    :    str     ):



							'''simple docstring'''







							snake_case   :							Union[str, 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
													snake_case   :							Optional[int]						  =self._prepare_for_class(_snake_case,    _snake_case     )
													snake_case   :							str						  ={k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

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

													snake_case   :							Optional[int]						  =pt_model_class(_snake_case     ).eval()
													snake_case   :							Tuple						  =model_class(_snake_case,    dtype=jnp.floataa     )

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

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

													with torch.no_grad():
																snake_case   :							List[Any]						  =pt_model(**_snake_case     ).to_tuple()

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

													with tempfile.TemporaryDirectory() as tmpdirname:
																fx_model.save_pretrained(_snake_case     )
																snake_case   :							Tuple						  =pt_model_class.from_pretrained(_snake_case,    from_flax=_snake_case     )

													with torch.no_grad():
																snake_case   :							List[str]						  =pt_model_loaded(**_snake_case     ).to_tuple()

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






				@slow
				def 					__snake_case       (    self    :    Union[str, Any]     ):



							'''simple docstring'''







							for model_class_name in self.all_model_classes:
										snake_case   :							Optional[int]						  =model_class_name.from_pretrained('''gpt2''',    from_pt=_snake_case     )
										snake_case   :							Union[str, Any]						  =model(np.ones((1, 1)     )     )
										self.assertIsNotNone(_snake_case     )










'''simple docstring'''





import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
      import torch

      from transformers import GPTaLMHeadModel





@require_torch
@require_sentencepiece
@require_tokenizers
class 							lowerCAmelCase_		(	unittest.TestCase						):





    @slow
    @unittest.skip('''Model is not available.'''     )
    def 					__snake_case       (    self    :    List[Any]     ):



       '''simple docstring'''







       snake_case   :							Any						  ='''nvidia/megatron-gpt2-345m/'''
       if "MYDIR" in os.environ:
          snake_case   :							Any						  =os.path.join(os.environ['''MYDIR'''],    _snake_case     )
       snake_case   :							Optional[Any]						  =GPTaLMHeadModel.from_pretrained(_snake_case     )
       model.to(_snake_case     )
       model.half()

       snake_case   :							int						  =torch.tensor(
           [[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]],    device=_snake_case,    dtype=torch.long,    )

       with torch.no_grad():
          snake_case   :							Dict						  =model(_snake_case     ).logits

       snake_case   :							Dict						  =torch.Size((1, 9, 50_257)     )
       self.assertEqual(output.shape,    _snake_case     )

       snake_case   :							str						  =torch.tensor(
           [
               4.9414,
               -0.2920,
               -1.2148,
               -4.0273,
               -0.5161,
               -5.2109,
               -1.2412,
               -1.8301,
               -1.7734,
               -4.7148,
               -0.2317,
               -1.0811,
               -2.1777,
               0.4141,
               -3.7969,
               -4.0586,
               -2.5332,
               -3.3809,
               4.3867,
           ],    device=_snake_case,    dtype=torch.half,    )

       for i in range(19     ):
          snake_case  ,    snake_case   :							Tuple						  =8 * i // 17, 2_792 * i  # along the diagonal
          snake_case  ,    snake_case   :							List[str]						  =output[0, r, c], expected_diag[i]
          snake_case   :							Optional[int]						  =f'''row={r} col={c} computed={computed} expected={expected}'''
          self.assertAlmostEqual(_snake_case,    _snake_case,    delta=1E-4,    msg=_snake_case     )





