


'''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 ,    _UpperCamelCase ,    _UpperCamelCase=14 ,    _UpperCamelCase=7 ,    _UpperCamelCase=True ,    _UpperCamelCase=True ,    _UpperCamelCase=False ,    _UpperCamelCase=True ,    _UpperCamelCase=99 ,    _UpperCamelCase=32 ,    _UpperCamelCase=5 ,    _UpperCamelCase=4 ,    _UpperCamelCase=37 ,    _UpperCamelCase="gelu" ,    _UpperCamelCase=0.1 ,    _UpperCamelCase=0.1 ,    _UpperCamelCase=512 ,    _UpperCamelCase=0.02 ,    )->			Any:
            _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      )->			Any:
            _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=_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      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      )->			Tuple:
            _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 ,    _UpperCamelCase ,    _UpperCamelCase ,    _UpperCamelCase ,    _UpperCamelCase      )->			Union[str, Any]:
            _A				=       20
            _A				=       model_class_name(_UpperCamelCase      )

            _A				=       model.init_cache(input_ids.shape[0] ,    _UpperCamelCase      )
            _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=_UpperCamelCase ,    past_key_values=_UpperCamelCase ,    position_ids=_UpperCamelCase ,    )

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

            _A				=       model(_UpperCamelCase      )

            _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 ,    _UpperCamelCase ,    _UpperCamelCase ,    _UpperCamelCase ,    _UpperCamelCase      )->			Optional[Any]:
            _A				=       20
            _A				=       model_class_name(_UpperCamelCase      )

            _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] ,    _UpperCamelCase      )
            _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=_UpperCamelCase ,    past_key_values=_UpperCamelCase ,    position_ids=_UpperCamelCase ,    )
            _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=_UpperCamelCase ,    position_ids=_UpperCamelCase ,    )

            _A				=       model(_UpperCamelCase ,    attention_mask=_UpperCamelCase      )

            _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_     (    UpperCAmelCase			,       UpperCAmelCase			,       unittest.TestCase  ):
       __UpperCAmelCase				=(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
       __UpperCAmelCase				=(FlaxGPTaLMHeadModel,) if is_flax_available() else ()






       def      UpperCamelCase	(   self      )->			Tuple:
            _A				=       FlaxGPTaModelTester(self      )






       def      UpperCamelCase	(   self      )->			Union[str, 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(_UpperCamelCase ,    _UpperCamelCase ,    _UpperCamelCase ,    _UpperCamelCase      )






       def      UpperCamelCase	(   self      )->			Dict:
            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(
                     _UpperCamelCase ,    _UpperCamelCase ,    _UpperCamelCase ,    _UpperCamelCase      )






       @slow
       def      UpperCamelCase	(   self      )->			Dict:
            _A				=       GPTaTokenizer.from_pretrained('gpt2' ,    pad_token='</s>' ,    padding_side='left'      )
            _A				=       tokenizer(['Hello this is a long string', 'Hey'] ,    return_tensors='np' ,    padding=_UpperCamelCase ,    truncation=_UpperCamelCase      )

            _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(_UpperCamelCase ,    skip_special_tokens=_UpperCamelCase      )

            _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(_UpperCamelCase ,    _UpperCamelCase      )






       @is_pt_flax_cross_test
       def      UpperCamelCase	(   self      )->			str:
            _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(_UpperCamelCase ,    _UpperCamelCase      )
                      _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(_UpperCamelCase ,    _UpperCamelCase      )

                      _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(_UpperCamelCase      ):
                           _A				=       0
                           _A				=       1
                           _A				=       0
                           _A				=       1
                      _A				=       pt_model_class(_UpperCamelCase      ).eval()
                      _A				=       model_class(_UpperCamelCase ,    dtype=jnp.floataa      )

                      _A				=       convert_pytorch_state_dict_to_flax(pt_model.state_dict() ,    _UpperCamelCase      )
                      _A				=       fx_state

                      with torch.no_grad():
                           _A				=       pt_model(**_UpperCamelCase      ).to_tuple()

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

                      _A				=       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      UpperCamelCase	(   self      )->			Optional[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(_UpperCamelCase ,    _UpperCamelCase      )
                      _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(_UpperCamelCase ,    _UpperCamelCase      )

                      _A				=       pt_model_class(_UpperCamelCase      ).eval()
                      _A				=       model_class(_UpperCamelCase ,    dtype=jnp.floataa      )

                      _A				=       load_flax_weights_in_pytorch_model(_UpperCamelCase ,    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(_UpperCamelCase      ):
                           _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(**_UpperCamelCase      ).to_tuple()

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

                      with torch.no_grad():
                           _A				=       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      UpperCamelCase	(   self      )->			int:
            for model_class_name in self.all_model_classes:
                 _A				=       model_class_name.from_pretrained('gpt2' ,    from_pt=_UpperCamelCase      )
                 _A				=       model(np.ones((1, 1)      )      )
                 self.assertIsNotNone(_UpperCamelCase      )





'''simple docstring'''
from __future__ import annotations

def 				lowerCamelCase_   (   __UpperCamelCase				:					list[int | float]		)				->  int | float:







     """simple docstring"""
     if len(__UpperCamelCase		) == 0:
          raise ValueError('find_max() arg is an empty sequence'		)
     _A				=       nums[0]
     for x in nums:
          if x > max_num:
               _A				=       x
     return max_num


if __name__ == "__main__":
    import doctest

    doctest.testmod(verbose=True)





