





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	:



      """simple docstring"""






      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			,  )    ->				str:






        '''simple docstring'''






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







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






        '''simple docstring'''






        __snake_case				:			Optional[Any]            =    ids_tensor([self.batch_size, self.seq_length]			,  self.vocab_size				)

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

        __snake_case				:			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=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				)    ->				Any:






        '''simple docstring'''






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







      def   UpperCAmelCase						(	self				)    ->				int:






        '''simple docstring'''






        __snake_case      ,				__snake_case      ,				__snake_case				:			Dict            =    self.prepare_config_and_inputs()

        __snake_case				:			Dict            =    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   UpperCAmelCase						(	self			,  UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase				)    ->				List[str]:






        '''simple docstring'''






        __snake_case				:			str            =    20
        __snake_case				:			Any            =    model_class_name(UpperCAmelCase				)

        __snake_case				:			Optional[int]            =    model.init_cache(input_ids.shape[0]			,  UpperCAmelCase				)
        __snake_case				:			Optional[Any]            =    jnp.ones((input_ids.shape[0], max_decoder_length)			,  dtype="i4"				)

        __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				:			Any            =    model(
            input_ids[:, :-1]			,  attention_mask=UpperCAmelCase			,  past_key_values=UpperCAmelCase			,  position_ids=UpperCAmelCase			,  )

        __snake_case				:			List[Any]            =    jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]			,  dtype="i4"				)
        __snake_case				:			Optional[int]            =    model(
            input_ids[:, -1:]			,  attention_mask=UpperCAmelCase			,  past_key_values=outputs_cache.past_key_values			,  position_ids=UpperCAmelCase			,  )

        __snake_case				:			List[Any]            =    model(UpperCAmelCase				)

        __snake_case				:			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			,  UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase				)    ->				Optional[Any]:






        '''simple docstring'''






        __snake_case				:			Optional[Any]            =    20
        __snake_case				:			Optional[int]            =    model_class_name(UpperCAmelCase				)

        __snake_case				:			str            =    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]			,  UpperCAmelCase				)
        __snake_case				:			str            =    jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1				)[None, :]			,  (input_ids.shape[0], input_ids.shape[-1] - 1)				)

        __snake_case				:			int            =    model(
            input_ids[:, :-1]			,  attention_mask=UpperCAmelCase			,  past_key_values=UpperCAmelCase			,  position_ids=UpperCAmelCase			,  )
        __snake_case				:			Optional[int]            =    jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]			,  dtype="i4"				)
        __snake_case				:			str            =    model(
            input_ids[:, -1:]			,  past_key_values=outputs_cache.past_key_values			,  attention_mask=UpperCAmelCase			,  position_ids=UpperCAmelCase			,  )

        __snake_case				:			List[Any]            =    model(UpperCAmelCase			,  attention_mask=UpperCAmelCase				)

        __snake_case				:			Optional[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      ):



      """simple docstring"""






      UpperCAmelCase_  : List[str]				      =(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
      UpperCAmelCase_  : Tuple				      =(FlaxGPTaLMHeadModel,) if is_flax_available() else ()







      def   UpperCAmelCase						(	self				)    ->				Optional[Any]:






        '''simple docstring'''






        __snake_case				:			List[Any]            =    FlaxGPTaModelTester(self				)







      def   UpperCAmelCase						(	self				)    ->				List[str]:






        '''simple docstring'''






        for model_class_name in self.all_model_classes:
          __snake_case      ,				__snake_case      ,				__snake_case				:			List[Any]            =    self.model_tester.prepare_config_and_inputs()
          self.model_tester.check_use_cache_forward(UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase				)







      def   UpperCAmelCase						(	self				)    ->				str:






        '''simple docstring'''






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






        '''simple docstring'''






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

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

        __snake_case				:			int            =    jax.jit(model.generate				)

        __snake_case				:			Union[str, Any]            =    jit_generate(inputs["input_ids"]			,  attention_mask=inputs["attention_mask"]				).sequences

        __snake_case				:			Optional[int]            =    tokenizer.batch_decode(UpperCAmelCase			,  skip_special_tokens=UpperCAmelCase				)

        __snake_case				:			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(UpperCAmelCase			,  UpperCAmelCase				)







      @is_pt_flax_cross_test
      def   UpperCAmelCase						(	self				)    ->				Union[str, Any]:






        '''simple docstring'''






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

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

            __snake_case      ,				__snake_case				:			Optional[Any]            =    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(UpperCAmelCase				):
              __snake_case				:			Tuple            =    0
              __snake_case				:			Any            =    1
              __snake_case				:			int            =    0
              __snake_case				:			int            =    1
            __snake_case				:			Optional[int]            =    pt_model_class(UpperCAmelCase				).eval()
            __snake_case				:			str            =    model_class(UpperCAmelCase			,  dtype=jnp.floataa				)

            __snake_case				:			Optional[Any]            =    convert_pytorch_state_dict_to_flax(pt_model.state_dict()			,  UpperCAmelCase				)
            __snake_case				:			Dict            =    fx_state

            with torch.no_grad():
              __snake_case				:			List[Any]            =    pt_model(**UpperCAmelCase				).to_tuple()

            __snake_case				:			Optional[int]            =    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				)
              __snake_case				:			Union[str, Any]            =    model_class.from_pretrained(UpperCAmelCase			,  from_pt=UpperCAmelCase				)

            __snake_case				:			Tuple            =    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]:






        '''simple docstring'''






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

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

            __snake_case				:			Any            =    pt_model_class(UpperCAmelCase				).eval()
            __snake_case				:			Optional[int]            =    model_class(UpperCAmelCase			,  dtype=jnp.floataa				)

            __snake_case				:			Optional[int]            =    load_flax_weights_in_pytorch_model(UpperCAmelCase			,  fx_model.params				)
            __snake_case      ,				__snake_case				:			Dict            =    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(UpperCAmelCase				):
              __snake_case				:			Optional[int]            =    0
              __snake_case				:			List[str]            =    1
              __snake_case				:			Tuple            =    0
              __snake_case				:			Optional[Any]            =    1

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

            with torch.no_grad():
              __snake_case				:			Any            =    pt_model(**UpperCAmelCase				).to_tuple()

            __snake_case				:			Tuple            =    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				)
              __snake_case				:			int            =    pt_model_class.from_pretrained(UpperCAmelCase			,  from_flax=UpperCAmelCase				)

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






        '''simple docstring'''






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








from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_UpperCamelCase    =	{
    '''configuration_dpr''': ['''DPR_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''DPRConfig'''],
    '''tokenization_dpr''': [
        '''DPRContextEncoderTokenizer''',
        '''DPRQuestionEncoderTokenizer''',
        '''DPRReaderOutput''',
        '''DPRReaderTokenizer''',
    ],
}


try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCamelCase    =	[
         '''DPRContextEncoderTokenizerFast''',
         '''DPRQuestionEncoderTokenizerFast''',
         '''DPRReaderTokenizerFast''',
     ]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCamelCase    =	[
         '''DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''DPRContextEncoder''',
         '''DPRPretrainedContextEncoder''',
         '''DPRPreTrainedModel''',
         '''DPRPretrainedQuestionEncoder''',
         '''DPRPretrainedReader''',
         '''DPRQuestionEncoder''',
         '''DPRReader''',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCamelCase    =	[
         '''TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TFDPRContextEncoder''',
         '''TFDPRPretrainedContextEncoder''',
         '''TFDPRPretrainedQuestionEncoder''',
         '''TFDPRPretrainedReader''',
         '''TFDPRQuestionEncoder''',
         '''TFDPRReader''',
     ]


if TYPE_CHECKING:
     from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
     from .tokenization_dpr import (
         DPRContextEncoderTokenizer,
         DPRQuestionEncoderTokenizer,
         DPRReaderOutput,
         DPRReaderTokenizer,
     )

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_dpr_fast import (
              DPRContextEncoderTokenizerFast,
              DPRQuestionEncoderTokenizerFast,
              DPRReaderTokenizerFast,
          )

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_dpr import (
              DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPRContextEncoder,
              DPRPretrainedContextEncoder,
              DPRPreTrainedModel,
              DPRPretrainedQuestionEncoder,
              DPRPretrainedReader,
              DPRQuestionEncoder,
              DPRReader,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_dpr import (
              TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFDPRContextEncoder,
              TFDPRPretrainedContextEncoder,
              TFDPRPretrainedQuestionEncoder,
              TFDPRPretrainedReader,
              TFDPRQuestionEncoder,
              TFDPRReader,
          )

else:
     import sys

     _UpperCamelCase    =	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

