





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






    '''simple docstring'''





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

      """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 						A		( self    :		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=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 						A		( self    :		Tuple				):

      """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 						A		( self    :		Optional[Any]				):

      """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 						A		( self    :		int      ,    snake_case    :		Tuple      ,    snake_case    :		Tuple      ,    snake_case    :		int      ,    snake_case    :		Union[str, Any]				):

      """simple docstring"""





      __lowerCAmelCase								=					20
      __lowerCAmelCase								=					model_class_name(snake_case				)

      __lowerCAmelCase								=					model.init_cache(input_ids.shape[0]      ,    snake_case				)
      __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=snake_case      ,    past_key_values=snake_case      ,    position_ids=snake_case      ,    )

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

      __lowerCAmelCase								=					model(snake_case				)

      __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 						A		( self    :		List[str]      ,    snake_case    :		List[Any]      ,    snake_case    :		Union[str, Any]      ,    snake_case    :		Any      ,    snake_case    :		int				):

      """simple docstring"""





      __lowerCAmelCase								=					20
      __lowerCAmelCase								=					model_class_name(snake_case				)

      __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]      ,    snake_case				)
      __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=snake_case      ,    past_key_values=snake_case      ,    position_ids=snake_case      ,    )
      __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=snake_case      ,    position_ids=snake_case      ,    )

      __lowerCAmelCase								=					model(snake_case      ,    attention_mask=snake_case				)

      __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 	__snake_case					(      _a							,					_a							,					unittest.TestCase ):






    '''simple docstring'''
    _lowerCAmelCase  =						(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
    _lowerCAmelCase  =						(FlaxGPTaLMHeadModel,) if is_flax_available() else ()





    def 						A		( self    :		Optional[int]				):

      """simple docstring"""





      __lowerCAmelCase								=					FlaxGPTaModelTester(self				)





    def 						A		( self    :		Optional[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(snake_case      ,    snake_case      ,    snake_case      ,    snake_case				)





    def 						A		( self    :		List[str]				):

      """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(
            snake_case      ,    snake_case      ,    snake_case      ,    snake_case				)





    @slow
    def 						A		( self    :		List[Any]				):

      """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=snake_case      ,    truncation=snake_case				)

      __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(snake_case      ,    skip_special_tokens=snake_case				)

      __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(snake_case      ,    snake_case				)





    @is_pt_flax_cross_test
    def 						A		( self    :		int				):

      """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(snake_case      ,    snake_case				)
          __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(snake_case      ,    snake_case				)

          __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(snake_case				):
            __lowerCAmelCase								=					0
            __lowerCAmelCase								=					1
            __lowerCAmelCase								=					0
            __lowerCAmelCase								=					1
          __lowerCAmelCase								=					pt_model_class(snake_case				).eval()
          __lowerCAmelCase								=					model_class(snake_case      ,    dtype=jnp.floataa				)

          __lowerCAmelCase								=					convert_pytorch_state_dict_to_flax(pt_model.state_dict()      ,    snake_case				)
          __lowerCAmelCase								=					fx_state

          with torch.no_grad():
            __lowerCAmelCase								=					pt_model(**snake_case				).to_tuple()

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

          __lowerCAmelCase								=					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 						A		( self    :		Optional[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(snake_case      ,    snake_case				)
          __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(snake_case      ,    snake_case				)

          __lowerCAmelCase								=					pt_model_class(snake_case				).eval()
          __lowerCAmelCase								=					model_class(snake_case      ,    dtype=jnp.floataa				)

          __lowerCAmelCase								=					load_flax_weights_in_pytorch_model(snake_case      ,    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(snake_case				):
            __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(**snake_case				).to_tuple()

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

          with torch.no_grad():
            __lowerCAmelCase								=					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 						A		( self    :		Any				):

      """simple docstring"""





      for model_class_name in self.all_model_classes:
        __lowerCAmelCase								=					model_class_name.from_pretrained('''gpt2'''      ,    from_pt=snake_case				)
        __lowerCAmelCase								=					model(np.ones((1, 1)				)				)
        self.assertIsNotNone(snake_case				)







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 	__snake_case					(      unittest.TestCase ):






    '''simple docstring'''





    @slow
    @unittest.skip('''Model is not available.'''				)
    def 						A		( self    :		str				):

      """simple docstring"""





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

      __lowerCAmelCase								=					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():
        __lowerCAmelCase								=					model(snake_case				).logits

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

      __lowerCAmelCase								=					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				):
        __lowerCAmelCase						,    __lowerCAmelCase								=					8 * i // 17, 2_792 * i  # along the diagonal
        __lowerCAmelCase						,    __lowerCAmelCase								=					output[0, r, c], expected_diag[i]
        __lowerCAmelCase								=					f'''row={r} col={c} computed={computed} expected={expected}'''
        self.assertAlmostEqual(snake_case      ,    snake_case      ,    delta=1e-4      ,    msg=snake_case				)
