

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






      '''simple docstring'''






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



          """simple docstring"""





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




      def 	__UpperCAmelCase			(					self :		int   ):



          """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 : Union[str, Any]												=		random_attention_mask([self.batch_size, self.seq_length]   )

          _snake_case : 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=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 :		int   ):



          """simple docstring"""





          _snake_case : str												=		self.prepare_config_and_inputs()
          _snake_case : Dict												=		config_and_inputs
          _snake_case : List[str]												=		{"input_ids": input_ids, "attention_mask": attention_mask}
          return config, inputs_dict




      def 	__UpperCAmelCase			(					self :		Optional[int]   ):



          """simple docstring"""





          _snake_case : str												=		self.prepare_config_and_inputs()

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



          """simple docstring"""





          _snake_case : Optional[Any]												=		20
          _snake_case : Union[str, Any]												=		model_class_name(lowerCamelCase__   )

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

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

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

          _snake_case : List[Any]												=		model(lowerCamelCase__   )

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



          """simple docstring"""





          _snake_case : Union[str, Any]												=		20
          _snake_case : List[str]												=		model_class_name(lowerCamelCase__   )

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

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

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

          _snake_case : Union[str, Any]												=		model(lowerCamelCase__     ,					attention_mask=lowerCamelCase__   )

          _snake_case : List[str]												=		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 			SCREAMING_SNAKE_CASE					(							UpperCAmelCase__      ,UpperCAmelCase__      ,unittest.TestCase					):






      '''simple docstring'''






      SCREAMING_SNAKE_CASE__ :				Optional[int]    =				(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
      SCREAMING_SNAKE_CASE__ :				Union[str, Any]    =				(FlaxGPTaLMHeadModel,) if is_flax_available() else ()




      def 	__UpperCAmelCase			(					self :		Optional[Any]   ):



          """simple docstring"""





          _snake_case : Any												=		FlaxGPTaModelTester(self   )




      def 	__UpperCAmelCase			(					self :		Union[str, Any]   ):



          """simple docstring"""





          for model_class_name in self.all_model_classes:
              _snake_case : Dict												=		self.model_tester.prepare_config_and_inputs()
              self.model_tester.check_use_cache_forward(lowerCamelCase__     ,					lowerCamelCase__     ,					lowerCamelCase__     ,					lowerCamelCase__   )




      def 	__UpperCAmelCase			(					self :		List[Any]   ):



          """simple docstring"""





          for model_class_name in self.all_model_classes:
              _snake_case : Optional[Any]												=		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 :		Any   ):



          """simple docstring"""





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

          _snake_case : List[Any]												=		FlaxGPTaLMHeadModel.from_pretrained('gpt2'   )
          _snake_case : List[str]												=		False
          _snake_case : int												=		model.config.eos_token_id

          _snake_case : Optional[int]												=		jax.jit(model.generate   )

          _snake_case : Any												=		jit_generate(inputs['input_ids']     ,					attention_mask=inputs['attention_mask']   ).sequences

          _snake_case : Any												=		tokenizer.batch_decode(lowerCamelCase__     ,					skip_special_tokens=lowerCamelCase__   )

          _snake_case : 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(lowerCamelCase__     ,					lowerCamelCase__   )




      @is_pt_flax_cross_test
      def 	__UpperCAmelCase			(					self :		Union[str, 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 : Any												=		self._prepare_for_class(lowerCamelCase__     ,					lowerCamelCase__   )
                  _snake_case : str												=		{k: torch.tensor(v.tolist()   ) for k, v in prepared_inputs_dict.items()}

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

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

                  _snake_case : int												=		convert_pytorch_state_dict_to_flax(pt_model.state_dict()     ,					lowerCamelCase__   )
                  _snake_case : Optional[int]												=		fx_state

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

                  _snake_case : List[str]												=		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__   )
                      _snake_case : Dict												=		model_class.from_pretrained(lowerCamelCase__     ,					from_pt=lowerCamelCase__   )

                  _snake_case : List[Any]												=		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 :		List[Any]   ):



          """simple docstring"""





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

                  # load corresponding PyTorch class
                  _snake_case : List[str]												=		model_class.__name__[4:]  # Skip the "Flax" at the beginning
                  _snake_case : Dict												=		getattr(lowerCamelCase__     ,					lowerCamelCase__   )

                  _snake_case : Dict												=		pt_model_class(lowerCamelCase__   ).eval()
                  _snake_case : Optional[Any]												=		model_class(lowerCamelCase__     ,					dtype=jnp.floataa   )

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

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

                  with torch.no_grad():
                      _snake_case : Any												=		pt_model(**lowerCamelCase__   ).to_tuple()

                  _snake_case : Union[str, Any]												=		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__   )
                      _snake_case : str												=		pt_model_class.from_pretrained(lowerCamelCase__     ,					from_flax=lowerCamelCase__   )

                  with torch.no_grad():
                      _snake_case : int												=		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 :		Any   ):



          """simple docstring"""





          for model_class_name in self.all_model_classes:
              _snake_case : Tuple												=		model_class_name.from_pretrained('gpt2'     ,					from_pt=lowerCamelCase__   )
              _snake_case : int												=		model(np.ones((1, 1)   )   )
              self.assertIsNotNone(lowerCamelCase__   )






'''simple docstring'''



from __future__ import annotations

def   lowerCamelCase__       (      a__)		->    int | float:



    """simple docstring"""







    _snake_case : int												=		sorted(a__)
    _snake_case : List[str]												=		len(a__)
    _snake_case : Union[str, Any]												=		length >> 1
    return (
        (sorted_list[mid_index] + sorted_list[mid_index - 1]) / 2
        if length % 2 == 0
        else sorted_list[mid_index]
    )

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



    """simple docstring"""







    import doctest

    doctest.testmod()


if __name__ == "__main__":
   main()



