

import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
 import os

 # The slow tests are often failing with OOM error on GPU
 # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
 # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
 _lowerCamelCase								="platform"

 import jax
 import jax.numpy as jnp

 from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer







def       snake_case__    (       lowerCAmelCase_,   lowerCAmelCase_,   lowerCAmelCase_=None,   lowerCAmelCase_=None      ):





     """simple docstring"""







     if attention_mask is None:
          SCREAMING_SNAKE_CASE										=np.where(input_ids != config.pad_token_id,   1,   0      )
     return {
         "input_ids": input_ids,
         "attention_mask": attention_mask,
     }



@require_flax
class 		a_	:



      """simple docstring"""







      def __init__(				self					:							int      ,snake_case					:							str      ,snake_case					:							int=13      ,snake_case					:							List[str]=7      ,snake_case					:							Any=True      ,snake_case					:							Optional[int]=False      ,snake_case					:							Any=99      ,snake_case					:							str=16      ,snake_case					:							List[Any]=2      ,snake_case					:							List[Any]=4      ,snake_case					:							Optional[int]=4      ,snake_case					:							Optional[int]="gelu"      ,snake_case					:							str=0.1      ,snake_case					:							int=0.1      ,snake_case					:							int=20      ,snake_case					:							Union[str, Any]=2      ,snake_case					:							List[str]=1      ,snake_case					:							Tuple=0      ,snake_case					:							Tuple=16      ,snake_case					:							Tuple=16      ,snake_case					:							Tuple=0.02      ,):
           SCREAMING_SNAKE_CASE										=parent
           SCREAMING_SNAKE_CASE										=batch_size
           SCREAMING_SNAKE_CASE										=seq_length
           SCREAMING_SNAKE_CASE										=is_training
           SCREAMING_SNAKE_CASE										=use_labels
           SCREAMING_SNAKE_CASE										=vocab_size
           SCREAMING_SNAKE_CASE										=hidden_size
           SCREAMING_SNAKE_CASE										=num_hidden_layers
           SCREAMING_SNAKE_CASE										=num_attention_heads
           SCREAMING_SNAKE_CASE										=intermediate_size
           SCREAMING_SNAKE_CASE										=hidden_act
           SCREAMING_SNAKE_CASE										=hidden_dropout_prob
           SCREAMING_SNAKE_CASE										=attention_probs_dropout_prob
           SCREAMING_SNAKE_CASE										=max_position_embeddings
           SCREAMING_SNAKE_CASE										=eos_token_id
           SCREAMING_SNAKE_CASE										=pad_token_id
           SCREAMING_SNAKE_CASE										=bos_token_id
           SCREAMING_SNAKE_CASE										=embed_dim
           SCREAMING_SNAKE_CASE										=word_embed_proj_dim
           SCREAMING_SNAKE_CASE										=initializer_range
           SCREAMING_SNAKE_CASE										=False
      def 	_lowerCAmelCase					(				self					:							Tuple    ):
           SCREAMING_SNAKE_CASE										=np.clip(ids_tensor([self.batch_size, self.seq_length - 1]      ,self.vocab_size    )      ,3      ,self.vocab_size    )
           SCREAMING_SNAKE_CASE										=np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)      ,dtype=np.intaa    ))      ,-1    )

           SCREAMING_SNAKE_CASE										=OPTConfig(
               vocab_size=self.vocab_size      ,hidden_size=self.hidden_size      ,num_hidden_layers=self.num_hidden_layers      ,num_attention_heads=self.num_attention_heads      ,ffn_dim=self.intermediate_size      ,dropout=self.hidden_dropout_prob      ,attention_dropout=self.attention_probs_dropout_prob      ,max_position_embeddings=self.max_position_embeddings      ,eos_token_id=self.eos_token_id      ,bos_token_id=self.bos_token_id      ,pad_token_id=self.pad_token_id      ,embed_dim=self.embed_dim      ,is_encoder_decoder=snake_case      ,word_embed_proj_dim=self.word_embed_proj_dim      ,initializer_range=self.initializer_range      ,use_cache=snake_case      ,)
           SCREAMING_SNAKE_CASE										=prepare_opt_inputs_dict(snake_case      ,snake_case    )
           return config, inputs_dict
      def 	_lowerCAmelCase					(				self					:							List[Any]    ):
           SCREAMING_SNAKE_CASE						,  SCREAMING_SNAKE_CASE										=self.prepare_config_and_inputs()
           return config, inputs_dict
      def 	_lowerCAmelCase					(				self					:							Dict      ,snake_case					:							Any      ,snake_case					:							Tuple      ,snake_case					:							Dict    ):
           SCREAMING_SNAKE_CASE										=20
           SCREAMING_SNAKE_CASE										=model_class_name(snake_case    )

           SCREAMING_SNAKE_CASE										=inputs_dict['input_ids']
           SCREAMING_SNAKE_CASE										=inputs_dict['attention_mask']

           SCREAMING_SNAKE_CASE										=model.init_cache(input_ids.shape[0]      ,snake_case    )
           SCREAMING_SNAKE_CASE										=jnp.ones((input_ids.shape[0], max_length)      ,dtype='i4'    )

           SCREAMING_SNAKE_CASE										=jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1    )[None, :]      ,(input_ids.shape[0], input_ids.shape[-1] - 1)      ,)
           SCREAMING_SNAKE_CASE										=model(
               input_ids[:, :-1]      ,attention_mask=snake_case      ,past_key_values=snake_case      ,position_ids=snake_case      ,)

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

           SCREAMING_SNAKE_CASE										=model(snake_case    )

           SCREAMING_SNAKE_CASE										=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 	_lowerCAmelCase					(				self					:							Optional[int]      ,snake_case					:							Optional[int]      ,snake_case					:							Any      ,snake_case					:							int    ):
           SCREAMING_SNAKE_CASE										=20
           SCREAMING_SNAKE_CASE										=model_class_name(snake_case    )

           SCREAMING_SNAKE_CASE						,  SCREAMING_SNAKE_CASE										=(
               inputs_dict['input_ids'],
               inputs_dict['attention_mask'],
           )

           SCREAMING_SNAKE_CASE										=jnp.concatenate(
               [
                   attention_mask,
                   jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])    ),
               ]      ,axis=-1      ,)

           SCREAMING_SNAKE_CASE										=model.init_cache(input_ids.shape[0]      ,snake_case    )
           SCREAMING_SNAKE_CASE										=jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1    )[None, :]      ,(input_ids.shape[0], input_ids.shape[-1] - 1)      ,)

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

           SCREAMING_SNAKE_CASE										=model(snake_case      ,attention_mask=snake_case    )

           SCREAMING_SNAKE_CASE										=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 		a_	(							lowerCamelCase_			,      unittest.TestCase			,      lowerCamelCase_					):



      """simple docstring"""







      __UpperCAmelCase      		=				(FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
      __UpperCAmelCase      		=				() if is_flax_available() else ()
      def 	_lowerCAmelCase					(				self					:							Tuple    ):
           SCREAMING_SNAKE_CASE										=FlaxOPTModelTester(self    )
      def 	_lowerCAmelCase					(				self					:							Optional[int]    ):
           SCREAMING_SNAKE_CASE						,  SCREAMING_SNAKE_CASE										=self.model_tester.prepare_config_and_inputs()
           for model_class in self.all_model_classes:
                self.model_tester.check_use_cache_forward(snake_case      ,snake_case      ,snake_case    )
      def 	_lowerCAmelCase					(				self					:							Optional[int]    ):
           SCREAMING_SNAKE_CASE						,  SCREAMING_SNAKE_CASE										=self.model_tester.prepare_config_and_inputs()
           for model_class in self.all_model_classes:
                self.model_tester.check_use_cache_forward_with_attn_mask(snake_case      ,snake_case      ,snake_case    )


      @slow
      def 	_lowerCAmelCase					(				self					:							Any    ):
           for model_class_name in self.all_model_classes:
                SCREAMING_SNAKE_CASE										=model_class_name.from_pretrained('facebook/opt-125m'    )
                SCREAMING_SNAKE_CASE										=np.ones((1, 1)    ) * model.config.eos_token_id
                SCREAMING_SNAKE_CASE										=model(snake_case    )
                self.assertIsNotNone(snake_case    )







@require_sentencepiece
@require_flax
class 		a_	(							unittest.TestCase					):



      """simple docstring"""







      @slow
      def 	_lowerCAmelCase					(				self					:							Dict    ):
           SCREAMING_SNAKE_CASE										=FlaxOPTModel.from_pretrained('facebook/opt-350m'    )
           SCREAMING_SNAKE_CASE										=jnp.array([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]    )
           SCREAMING_SNAKE_CASE										=model(input_ids=snake_case    ).last_hidden_state
           SCREAMING_SNAKE_CASE										=(1, 11, 512)
           self.assertEqual(output.shape      ,snake_case    )
           SCREAMING_SNAKE_CASE										=jnp.array(
               [[-0.2_867, -1.9_256, -0.3_062], [-1.2_711, -0.1_337, -0.1_897], [0.4_109, 0.1_187, -1.3_142]]    )
           self.assertTrue(jnp.allclose(output[:, :3, :3]      ,snake_case      ,atol=4e-2    )    )







@require_flax
@slow
class 		a_	(							unittest.TestCase					):



      """simple docstring"""







      def 	_lowerCAmelCase					(				self					:							List[Any]    ):
           super().setUp()
           SCREAMING_SNAKE_CASE										='facebook/opt-350m'


      def 	_lowerCAmelCase					(				self					:							Dict    ):
           SCREAMING_SNAKE_CASE										=FlaxOPTForCausalLM.from_pretrained(self.path_model    )
           SCREAMING_SNAKE_CASE										=GPTaTokenizer.from_pretrained(self.path_model    )

           SCREAMING_SNAKE_CASE										=[
               'Today is a beautiful day and I want to',
               'In the city of',
               'Paris is the capital of France and',
               'Computers and mobile phones have taken',
           ]
           # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
           SCREAMING_SNAKE_CASE										=tokenizer(snake_case      ,return_tensors='jax'      ,padding=snake_case      ,add_special_tokens=snake_case    )
           SCREAMING_SNAKE_CASE										=model(inputs.input_ids      ,attention_mask=inputs.attention_mask    )[0].mean(axis=-1    )
           SCREAMING_SNAKE_CASE										=jnp.array(
               [
                   [1.3_851, -13.8_923, -10.5_229, -10.7_533, -0.2_309, -10.2_384, -0.5_365, -9.0_947, -5.1_670],
                   [-4.7_073, -10.6_276, -3.9_415, -21.5_242, -0.2_822, -0.2_822, -0.2_822, -0.2_822, -0.2_822],
                   [0.6_247, -3.4_229, -8.9_179, -1.4_297, -14.1_650, 1.4_146, -9.0_218, -0.2_703, -0.2_703],
                   [6.4_783, -1.9_913, -10.7_926, -2.3_336, 1.5_092, -0.9_974, -6.8_213, 1.3_477, 1.3_477],
               ]    )
           self.assertTrue(jnp.allclose(snake_case      ,snake_case      ,atol=4e-2    )    )

           SCREAMING_SNAKE_CASE										=jax.jit(snake_case    )
           SCREAMING_SNAKE_CASE										=model(inputs.input_ids      ,attention_mask=inputs.attention_mask    )[0].mean(axis=-1    )
           self.assertTrue(jnp.allclose(snake_case      ,snake_case      ,atol=4e-2    )    )







@require_flax
@slow
class 		a_	(							unittest.TestCase					):



      """simple docstring"""







      @property
      def 	_lowerCAmelCase					(				self					:							int    ):
           return [
               "Today is a beautiful day and I want",
               "In the city of",
               "Paris is the capital of France and",
               "Computers and mobile phones have taken",
           ]
      def 	_lowerCAmelCase					(				self					:							str    ):
           SCREAMING_SNAKE_CASE										='facebook/opt-125m'

           SCREAMING_SNAKE_CASE										=[
               'Today is a beautiful day and I want to',
               'In the city of New York, the city',
               'Paris is the capital of France and the capital',
               'Computers and mobile phones have taken over the',
           ]

           SCREAMING_SNAKE_CASE										=[]

           SCREAMING_SNAKE_CASE										=FlaxOPTForCausalLM.from_pretrained(snake_case    )
           SCREAMING_SNAKE_CASE										=GPTaTokenizer.from_pretrained(snake_case    )

           for prompt in self.prompts:
                SCREAMING_SNAKE_CASE										=tokenizer(snake_case      ,return_tensors='jax'    ).input_ids

                SCREAMING_SNAKE_CASE										=model.generate(snake_case      ,max_length=10    )
                SCREAMING_SNAKE_CASE										=generated_ids[0]

                SCREAMING_SNAKE_CASE										=tokenizer.batch_decode(snake_case      ,skip_special_tokens=snake_case    )
                predicted_outputs += generated_string

           self.assertListEqual(snake_case      ,snake_case    )
      def 	_lowerCAmelCase					(				self					:							List[str]    ):
           SCREAMING_SNAKE_CASE										='facebook/opt-350m'

           SCREAMING_SNAKE_CASE										=[
               'Today is a beautiful day and I want to',
               'In the city of San Francisco, the city',
               'Paris is the capital of France and the capital',
               'Computers and mobile phones have taken over the',
           ]

           SCREAMING_SNAKE_CASE										=[]
           SCREAMING_SNAKE_CASE										=FlaxOPTForCausalLM.from_pretrained(snake_case    )
           SCREAMING_SNAKE_CASE										=GPTaTokenizer.from_pretrained(snake_case    )

           for prompt in self.prompts:
                SCREAMING_SNAKE_CASE										=tokenizer(snake_case      ,return_tensors='jax'    ).input_ids

                SCREAMING_SNAKE_CASE										=model.generate(snake_case      ,max_length=10    )
                SCREAMING_SNAKE_CASE										=generated_ids[0]

                SCREAMING_SNAKE_CASE										=tokenizer.batch_decode(snake_case      ,skip_special_tokens=snake_case    )
                predicted_outputs += generated_string

           self.assertListEqual(snake_case      ,snake_case    )
      def 	_lowerCAmelCase					(				self					:							List[str]    ):
           SCREAMING_SNAKE_CASE										='facebook/opt-125m'
           SCREAMING_SNAKE_CASE										=[
               'Today is a beautiful day and I want to thank',
               'In the city of Rome Canaver Canaver Canaver Canaver',
           ]
           SCREAMING_SNAKE_CASE										=FlaxOPTForCausalLM.from_pretrained(snake_case    )
           SCREAMING_SNAKE_CASE										=GPTaTokenizer.from_pretrained(snake_case    )
           SCREAMING_SNAKE_CASE										=tokenizer(
               [
                   'Today is a beautiful day and I want to',
                   'In the city of',
               ]      ,return_tensors='jax'      ,padding=snake_case      ,)

           SCREAMING_SNAKE_CASE										=jax.jit(model.generate    )

           SCREAMING_SNAKE_CASE										=jit_generate(inputs['input_ids']      ,attention_mask=inputs['attention_mask']    ).sequences

           SCREAMING_SNAKE_CASE										=tokenizer.batch_decode(snake_case      ,skip_special_tokens=snake_case    )

           self.assertIsNotNone(snake_case      ,snake_case    )


      def 	_lowerCAmelCase					(				self					:							Union[str, Any]    ):
           SCREAMING_SNAKE_CASE										='facebook/opt-350m'

           SCREAMING_SNAKE_CASE										=GPTaTokenizer.from_pretrained(snake_case    )
           SCREAMING_SNAKE_CASE										=FlaxOPTForCausalLM.from_pretrained(snake_case    )

           SCREAMING_SNAKE_CASE										='left'

           # use different length sentences to test batching
           SCREAMING_SNAKE_CASE										=[
               'Hello, my dog is a little',
               'Today, I',
           ]

           SCREAMING_SNAKE_CASE										=tokenizer(snake_case      ,return_tensors='jax'      ,padding=snake_case    )
           SCREAMING_SNAKE_CASE										=inputs['input_ids']

           SCREAMING_SNAKE_CASE										=model.generate(input_ids=snake_case      ,attention_mask=inputs['attention_mask']      ,trace=snake_case    )

           SCREAMING_SNAKE_CASE										=tokenizer(sentences[0]      ,return_tensors='jax'    ).input_ids
           SCREAMING_SNAKE_CASE										=model.generate(input_ids=snake_case    )

           SCREAMING_SNAKE_CASE										=inputs_non_padded.shape[-1] - inputs['attention_mask'][-1].sum()
           SCREAMING_SNAKE_CASE										=tokenizer(sentences[1]      ,return_tensors='jax'    ).input_ids
           SCREAMING_SNAKE_CASE										=model.generate(input_ids=snake_case      ,max_length=model.config.max_length - num_paddings    )

           SCREAMING_SNAKE_CASE										=tokenizer.batch_decode(outputs[0]      ,skip_special_tokens=snake_case    )
           SCREAMING_SNAKE_CASE										=tokenizer.decode(output_non_padded[0][0]      ,skip_special_tokens=snake_case    )
           SCREAMING_SNAKE_CASE										=tokenizer.decode(output_padded[0][0]      ,skip_special_tokens=snake_case    )

           SCREAMING_SNAKE_CASE										=[
               'Hello, my dog is a little bit of a dork.\nI\'m a little bit',
               'Today, I was in the middle of a conversation with a friend about the',
           ]
           self.assertListEqual(snake_case      ,snake_case    )
           self.assertListEqual(snake_case      ,[non_padded_sentence, padded_sentence]    )




from __future__ import annotations

import unittest

from transformers import MPNetConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
 import tensorflow as tf

 from transformers.models.mpnet.modeling_tf_mpnet import (
     TFMPNetForMaskedLM,
     TFMPNetForMultipleChoice,
     TFMPNetForQuestionAnswering,
     TFMPNetForSequenceClassification,
     TFMPNetForTokenClassification,
     TFMPNetModel,
 )



class 		a_	:



      """simple docstring"""







      def __init__(				self					:							Any      ,snake_case					:							List[Any]      ,snake_case					:							Optional[Any]=13      ,snake_case					:							str=7      ,snake_case					:							List[Any]=True      ,snake_case					:							Optional[Any]=True      ,snake_case					:							int=False      ,snake_case					:							List[str]=True      ,snake_case					:							Optional[int]=99      ,snake_case					:							Any=64      ,snake_case					:							int=2      ,snake_case					:							str=4      ,snake_case					:							str=64      ,snake_case					:							List[Any]="gelu"      ,snake_case					:							Dict=0.1      ,snake_case					:							int=0.1      ,snake_case					:							List[Any]=512      ,snake_case					:							int=16      ,snake_case					:							List[Any]=2      ,snake_case					:							Union[str, Any]=0.02      ,snake_case					:							List[str]=3      ,snake_case					:							Any=4      ,snake_case					:							Tuple=None      ,):
           SCREAMING_SNAKE_CASE										=parent
           SCREAMING_SNAKE_CASE										=batch_size
           SCREAMING_SNAKE_CASE										=seq_length
           SCREAMING_SNAKE_CASE										=is_training
           SCREAMING_SNAKE_CASE										=use_input_mask
           SCREAMING_SNAKE_CASE										=use_token_type_ids
           SCREAMING_SNAKE_CASE										=use_labels
           SCREAMING_SNAKE_CASE										=vocab_size
           SCREAMING_SNAKE_CASE										=hidden_size
           SCREAMING_SNAKE_CASE										=num_hidden_layers
           SCREAMING_SNAKE_CASE										=num_attention_heads
           SCREAMING_SNAKE_CASE										=intermediate_size
           SCREAMING_SNAKE_CASE										=hidden_act
           SCREAMING_SNAKE_CASE										=hidden_dropout_prob
           SCREAMING_SNAKE_CASE										=attention_probs_dropout_prob
           SCREAMING_SNAKE_CASE										=max_position_embeddings
           SCREAMING_SNAKE_CASE										=type_vocab_size
           SCREAMING_SNAKE_CASE										=type_sequence_label_size
           SCREAMING_SNAKE_CASE										=initializer_range
           SCREAMING_SNAKE_CASE										=num_labels
           SCREAMING_SNAKE_CASE										=num_choices
           SCREAMING_SNAKE_CASE										=scope
      def 	_lowerCAmelCase					(				self					:							List[Any]    ):
           SCREAMING_SNAKE_CASE										=ids_tensor([self.batch_size, self.seq_length]      ,self.vocab_size    )

           SCREAMING_SNAKE_CASE										=None
           if self.use_input_mask:
                SCREAMING_SNAKE_CASE										=random_attention_mask([self.batch_size, self.seq_length]    )

           SCREAMING_SNAKE_CASE										=None
           SCREAMING_SNAKE_CASE										=None
           SCREAMING_SNAKE_CASE										=None
           if self.use_labels:
                SCREAMING_SNAKE_CASE										=ids_tensor([self.batch_size]      ,self.type_sequence_label_size    )
                SCREAMING_SNAKE_CASE										=ids_tensor([self.batch_size, self.seq_length]      ,self.num_labels    )
                SCREAMING_SNAKE_CASE										=ids_tensor([self.batch_size]      ,self.num_choices    )

           SCREAMING_SNAKE_CASE										=MPNetConfig(
               vocab_size=self.vocab_size      ,hidden_size=self.hidden_size      ,num_hidden_layers=self.num_hidden_layers      ,num_attention_heads=self.num_attention_heads      ,intermediate_size=self.intermediate_size      ,hidden_act=self.hidden_act      ,hidden_dropout_prob=self.hidden_dropout_prob      ,attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,max_position_embeddings=self.max_position_embeddings      ,initializer_range=self.initializer_range      ,)
           return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
      def 	_lowerCAmelCase					(				self					:							Dict      ,snake_case					:							Any      ,snake_case					:							Optional[int]      ,snake_case					:							str      ,snake_case					:							Union[str, Any]      ,snake_case					:							Optional[Any]      ,snake_case					:							Any    ):
           SCREAMING_SNAKE_CASE										=TFMPNetModel(config=snake_case    )
           SCREAMING_SNAKE_CASE										={'input_ids': input_ids, 'attention_mask': input_mask}
           SCREAMING_SNAKE_CASE										=model(snake_case    )
           SCREAMING_SNAKE_CASE										=[input_ids, input_mask]
           SCREAMING_SNAKE_CASE										=model(snake_case    )
           self.parent.assertEqual(result.last_hidden_state.shape      ,(self.batch_size, self.seq_length, self.hidden_size)    )
      def 	_lowerCAmelCase					(				self					:							List[Any]      ,snake_case					:							List[str]      ,snake_case					:							Tuple      ,snake_case					:							Dict      ,snake_case					:							Optional[Any]      ,snake_case					:							Optional[int]      ,snake_case					:							Dict    ):
           SCREAMING_SNAKE_CASE										=TFMPNetForMaskedLM(config=snake_case    )
           SCREAMING_SNAKE_CASE										={'input_ids': input_ids, 'attention_mask': input_mask}
           SCREAMING_SNAKE_CASE										=model(snake_case    )
           self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.vocab_size)    )
      def 	_lowerCAmelCase					(				self					:							int      ,snake_case					:							List[str]      ,snake_case					:							Dict      ,snake_case					:							List[Any]      ,snake_case					:							Optional[int]      ,snake_case					:							Optional[int]      ,snake_case					:							Tuple    ):
           SCREAMING_SNAKE_CASE										=TFMPNetForQuestionAnswering(config=snake_case    )
           SCREAMING_SNAKE_CASE										={
               'input_ids': input_ids,
               'attention_mask': input_mask,
           }
           SCREAMING_SNAKE_CASE										=model(snake_case    )
           self.parent.assertEqual(result.start_logits.shape      ,(self.batch_size, self.seq_length)    )
           self.parent.assertEqual(result.end_logits.shape      ,(self.batch_size, self.seq_length)    )
      def 	_lowerCAmelCase					(				self					:							List[Any]      ,snake_case					:							Union[str, Any]      ,snake_case					:							str      ,snake_case					:							str      ,snake_case					:							Any      ,snake_case					:							Tuple      ,snake_case					:							Optional[Any]    ):
           SCREAMING_SNAKE_CASE										=self.num_labels
           SCREAMING_SNAKE_CASE										=TFMPNetForSequenceClassification(snake_case    )
           SCREAMING_SNAKE_CASE										={'input_ids': input_ids, 'attention_mask': input_mask}
           SCREAMING_SNAKE_CASE										=model(snake_case    )
           self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.num_labels)    )
      def 	_lowerCAmelCase					(				self					:							int      ,snake_case					:							Any      ,snake_case					:							Union[str, Any]      ,snake_case					:							int      ,snake_case					:							List[str]      ,snake_case					:							Tuple      ,snake_case					:							Union[str, Any]    ):
           SCREAMING_SNAKE_CASE										=self.num_choices
           SCREAMING_SNAKE_CASE										=TFMPNetForMultipleChoice(snake_case    )
           SCREAMING_SNAKE_CASE										=tf.tile(tf.expand_dims(snake_case      ,1    )      ,(1, self.num_choices, 1)    )
           SCREAMING_SNAKE_CASE										=tf.tile(tf.expand_dims(snake_case      ,1    )      ,(1, self.num_choices, 1)    )
           SCREAMING_SNAKE_CASE										={
               'input_ids': multiple_choice_inputs_ids,
               'attention_mask': multiple_choice_input_mask,
           }
           SCREAMING_SNAKE_CASE										=model(snake_case    )
           self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.num_choices)    )
      def 	_lowerCAmelCase					(				self					:							Any      ,snake_case					:							Tuple      ,snake_case					:							Dict      ,snake_case					:							Tuple      ,snake_case					:							Union[str, Any]      ,snake_case					:							Optional[int]      ,snake_case					:							Dict    ):
           SCREAMING_SNAKE_CASE										=self.num_labels
           SCREAMING_SNAKE_CASE										=TFMPNetForTokenClassification(snake_case    )
           SCREAMING_SNAKE_CASE										={'input_ids': input_ids, 'attention_mask': input_mask}
           SCREAMING_SNAKE_CASE										=model(snake_case    )
           self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.num_labels)    )


      def 	_lowerCAmelCase					(				self					:							Optional[int]    ):
           SCREAMING_SNAKE_CASE										=self.prepare_config_and_inputs()
           ((SCREAMING_SNAKE_CASE)						,  (SCREAMING_SNAKE_CASE)						,  (SCREAMING_SNAKE_CASE)						,  (SCREAMING_SNAKE_CASE)						,  (SCREAMING_SNAKE_CASE)						,  (SCREAMING_SNAKE_CASE))										=config_and_inputs
           SCREAMING_SNAKE_CASE										={'input_ids': input_ids, 'attention_mask': input_mask}
           return config, inputs_dict







@require_tf
class 		a_	(							lowerCamelCase_			,      lowerCamelCase_			,      unittest.TestCase					):



      """simple docstring"""







      __UpperCAmelCase      		=				(
          (
              TFMPNetForMaskedLM,
              TFMPNetForMultipleChoice,
              TFMPNetForQuestionAnswering,
              TFMPNetForSequenceClassification,
              TFMPNetForTokenClassification,
              TFMPNetModel,
          )
          if is_tf_available()
          else ()
      )
      __UpperCAmelCase      		=				(
          {
              'feature-extraction': TFMPNetModel,
              'fill-mask': TFMPNetForMaskedLM,
              'question-answering': TFMPNetForQuestionAnswering,
              'text-classification': TFMPNetForSequenceClassification,
              'token-classification': TFMPNetForTokenClassification,
              'zero-shot': TFMPNetForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      __UpperCAmelCase      		=				False
      __UpperCAmelCase      		=				False
      def 	_lowerCAmelCase					(				self					:							Dict    ):
           SCREAMING_SNAKE_CASE										=TFMPNetModelTester(self    )
           SCREAMING_SNAKE_CASE										=ConfigTester(self      ,config_class=snake_case      ,hidden_size=37    )
      def 	_lowerCAmelCase					(				self					:							Dict    ):
           self.config_tester.run_common_tests()
      def 	_lowerCAmelCase					(				self					:							Union[str, Any]    ):
           SCREAMING_SNAKE_CASE										=self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mpnet_model(*snake_case    )
      def 	_lowerCAmelCase					(				self					:							List[Any]    ):
           SCREAMING_SNAKE_CASE										=self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mpnet_for_masked_lm(*snake_case    )
      def 	_lowerCAmelCase					(				self					:							Dict    ):
           SCREAMING_SNAKE_CASE										=self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mpnet_for_question_answering(*snake_case    )
      def 	_lowerCAmelCase					(				self					:							str    ):
           SCREAMING_SNAKE_CASE										=self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mpnet_for_sequence_classification(*snake_case    )
      def 	_lowerCAmelCase					(				self					:							Dict    ):
           SCREAMING_SNAKE_CASE										=self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mpnet_for_multiple_choice(*snake_case    )
      def 	_lowerCAmelCase					(				self					:							List[str]    ):
           SCREAMING_SNAKE_CASE										=self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mpnet_for_token_classification(*snake_case    )


      @slow
      def 	_lowerCAmelCase					(				self					:							Optional[Any]    ):
           for model_name in ["microsoft/mpnet-base"]:
                SCREAMING_SNAKE_CASE										=TFMPNetModel.from_pretrained(snake_case    )
                self.assertIsNotNone(snake_case    )







@require_tf
class 		a_	(							unittest.TestCase					):



      """simple docstring"""







      @slow
      def 	_lowerCAmelCase					(				self					:							Optional[Any]    ):
           SCREAMING_SNAKE_CASE										=TFMPNetModel.from_pretrained('microsoft/mpnet-base'    )
           SCREAMING_SNAKE_CASE										=tf.constant([[0, 1, 2, 3, 4, 5]]    )
           SCREAMING_SNAKE_CASE										=model(snake_case    )[0]

           SCREAMING_SNAKE_CASE										=[1, 6, 768]
           self.assertEqual(output.shape      ,snake_case    )

           SCREAMING_SNAKE_CASE										=tf.constant(
               [
                   [
                       [-0.1_067_172, 0.08_216_473, 0.0_024_543],
                       [-0.03_465_879, 0.8_354_118, -0.03_252_288],
                       [-0.06_569_476, -0.12_424_111, -0.0_494_436],
                   ]
               ]    )
           tf.debugging.assert_near(output[:, :3, :3]      ,snake_case      ,atol=1e-4    )

