'''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					,      lowerCAmelCase__					,      lowerCAmelCase__=14					,      lowerCAmelCase__=7					,      lowerCAmelCase__=True					,      lowerCAmelCase__=True					,      lowerCAmelCase__=False					,      lowerCAmelCase__=True					,      lowerCAmelCase__=99					,      lowerCAmelCase__=32					,      lowerCAmelCase__=5					,      lowerCAmelCase__=4					,      lowerCAmelCase__=37					,      lowerCAmelCase__="gelu"					,      lowerCAmelCase__=0.1					,      lowerCAmelCase__=0.1					,      lowerCAmelCase__=512					,      lowerCAmelCase__=0.02					,      )				->     Tuple:
        """simple docstring"""





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



   def  _UpperCamelCase		(    self     )				->     Dict:
        """simple docstring"""





        _UpperCamelCase  :Dict       =ids_tensor([self.batch_size, self.seq_length]					,      self.vocab_size     )

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

        _UpperCamelCase  :Optional[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     )				->     List[str]:
        """simple docstring"""





        _UpperCamelCase  :Any       =self.prepare_config_and_inputs()
        _UpperCamelCase       ,      _UpperCamelCase       ,      _UpperCamelCase  :Any       =config_and_inputs
        _UpperCamelCase  :Union[str, Any]       ={"""input_ids""": input_ids, """attention_mask""": attention_mask}
        return config, inputs_dict



   def  _UpperCamelCase		(    self     )				->     List[str]:
        """simple docstring"""





        _UpperCamelCase       ,      _UpperCamelCase       ,      _UpperCamelCase  :List[str]       =self.prepare_config_and_inputs()

        _UpperCamelCase  :int       =floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
        _UpperCamelCase  :Union[str, 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  _UpperCamelCase		(    self					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__     )				->     Union[str, Any]:
        """simple docstring"""





        _UpperCamelCase  :List[str]       =20
        _UpperCamelCase  :Union[str, Any]       =model_class_name(lowerCAmelCase__     )

        _UpperCamelCase  :List[str]       =model.init_cache(input_ids.shape[0]					,      lowerCAmelCase__     )
        _UpperCamelCase  :Any       =jnp.ones((input_ids.shape[0], max_decoder_length)					,      dtype="""i4"""     )

        _UpperCamelCase  :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=lowerCAmelCase__					,      past_key_values=lowerCAmelCase__					,      position_ids=lowerCAmelCase__					,      )

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

        _UpperCamelCase  :int       =model(lowerCAmelCase__     )

        _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}'''     )




   def  _UpperCamelCase		(    self					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__     )				->     Any:
        """simple docstring"""





        _UpperCamelCase  :Dict       =20
        _UpperCamelCase  :List[str]       =model_class_name(lowerCAmelCase__     )

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

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

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

        _UpperCamelCase  :Tuple       =model(lowerCAmelCase__					,      attention_mask=lowerCAmelCase__     )

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



   def  _UpperCamelCase		(    self     )				->     Union[str, Any]:
        """simple docstring"""





        _UpperCamelCase  :int       =FlaxGPTaModelTester(self     )



   def  _UpperCamelCase		(    self     )				->     List[str]:
        """simple docstring"""





        for model_class_name in self.all_model_classes:
             _UpperCamelCase       ,      _UpperCamelCase       ,      _UpperCamelCase  :Dict       =self.model_tester.prepare_config_and_inputs()
             self.model_tester.check_use_cache_forward(lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__     )



   def  _UpperCamelCase		(    self     )				->     Any:
        """simple docstring"""





        for model_class_name in self.all_model_classes:
             _UpperCamelCase       ,      _UpperCamelCase       ,      _UpperCamelCase  :Dict       =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     )				->     List[str]:
        """simple docstring"""





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

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

        _UpperCamelCase  :Any       =jax.jit(model.generate     )

        _UpperCamelCase  :Dict       =jit_generate(inputs["""input_ids"""]					,      attention_mask=inputs["""attention_mask"""]     ).sequences

        _UpperCamelCase  :Optional[int]       =tokenizer.batch_decode(lowerCAmelCase__					,      skip_special_tokens=lowerCAmelCase__     )

        _UpperCamelCase  :Optional[int]       =[
            """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     )				->     Any:
        """simple docstring"""





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

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

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

                  _UpperCamelCase  :Any       =convert_pytorch_state_dict_to_flax(pt_model.state_dict()					,      lowerCAmelCase__     )
                  _UpperCamelCase  :List[str]       =fx_state

                  with torch.no_grad():
                       _UpperCamelCase  :Optional[int]       =pt_model(**lowerCAmelCase__     ).to_tuple()

                  _UpperCamelCase  :List[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:
                       pt_model.save_pretrained(lowerCAmelCase__     )
                       _UpperCamelCase  :Any       =model_class.from_pretrained(lowerCAmelCase__					,      from_pt=lowerCAmelCase__     )

                  _UpperCamelCase  :Optional[int]       =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     )				->     Any:
        """simple docstring"""





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

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

                  _UpperCamelCase  :Union[str, Any]       =pt_model_class(lowerCAmelCase__     ).eval()
                  _UpperCamelCase  :Dict       =model_class(lowerCAmelCase__					,      dtype=jnp.floataa     )

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

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

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

                  _UpperCamelCase  :Tuple       =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__     )
                       _UpperCamelCase  :Dict       =pt_model_class.from_pretrained(lowerCAmelCase__					,      from_flax=lowerCAmelCase__     )

                  with torch.no_grad():
                       _UpperCamelCase  :str       =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     )				->     List[Any]:
        """simple docstring"""





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



"""simple docstring"""
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


SCREAMING_SNAKE_CASE					:       Tuple							   =							{
    """configuration_realm""": ["""REALM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """RealmConfig"""],
    """tokenization_realm""": ["""RealmTokenizer"""],
}

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			SCREAMING_SNAKE_CASE					:       Optional[Any]							   =							["""RealmTokenizerFast"""]

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			SCREAMING_SNAKE_CASE					:       Dict							   =							[
			    """REALM_PRETRAINED_MODEL_ARCHIVE_LIST""",
			    """RealmEmbedder""",
			    """RealmForOpenQA""",
			    """RealmKnowledgeAugEncoder""",
			    """RealmPreTrainedModel""",
			    """RealmReader""",
			    """RealmScorer""",
			    """load_tf_weights_in_realm""",
			]
			SCREAMING_SNAKE_CASE					:       Union[str, Any]							   =							["""RealmRetriever"""]


if TYPE_CHECKING:
			from .configuration_realm import REALM_PRETRAINED_CONFIG_ARCHIVE_MAP, RealmConfig
			from .tokenization_realm import RealmTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_realm import RealmTokenizerFast

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_realm import (
						    REALM_PRETRAINED_MODEL_ARCHIVE_LIST,
						    RealmEmbedder,
						    RealmForOpenQA,
						    RealmKnowledgeAugEncoder,
						    RealmPreTrainedModel,
						    RealmReader,
						    RealmScorer,
						    load_tf_weights_in_realm,
						)
						from .retrieval_realm import RealmRetriever


else:
			import sys

			SCREAMING_SNAKE_CASE					:       str							   =							_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





