






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  UpperCamelCase    :


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


 def _lowercase			(  self	:		Tuple							)      -> Optional[Any]:
   _a						: str    =			ids_tensor([self.batch_size, self.seq_length]							,		self.vocab_size							)

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

   _a						: Any    =			None
   _a						: int    =			None
   _a						: int    =			None
   if self.use_labels:
     _a						: Tuple    =			ids_tensor([self.batch_size]							,		self.type_sequence_label_size							)
     _a						: int    =			ids_tensor([self.batch_size, self.seq_length]							,		self.num_labels							)
     _a						: Optional[Any]    =			ids_tensor([self.batch_size]							,		self.num_choices							)

   _a						: str    =			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 _lowercase			(  self	:		Optional[Any]							,		UpperCAmelCase__	:		Optional[Any]							,		UpperCAmelCase__	:		Tuple							,		UpperCAmelCase__	:		int							,		UpperCAmelCase__	:		Optional[Any]							,		UpperCAmelCase__	:		List[Any]							,		UpperCAmelCase__	:		Tuple							)      -> Dict:
   _a						: Dict    =			TFMPNetModel(config=UpperCAmelCase__							)
   _a						: Optional[int]    =			{"""input_ids""": input_ids, """attention_mask""": input_mask}
   _a						: Tuple    =			model(UpperCAmelCase__							)
   _a						: int    =			[input_ids, input_mask]
   _a						: Optional[int]    =			model(UpperCAmelCase__							)
   self.parent.assertEqual(result.last_hidden_state.shape							,		(self.batch_size, self.seq_length, self.hidden_size)							)


 def _lowercase			(  self	:		Dict							,		UpperCAmelCase__	:		Tuple							,		UpperCAmelCase__	:		Optional[int]							,		UpperCAmelCase__	:		Dict							,		UpperCAmelCase__	:		Dict							,		UpperCAmelCase__	:		str							,		UpperCAmelCase__	:		Union[str, Any]							)      -> str:
   _a						: Tuple    =			TFMPNetForMaskedLM(config=UpperCAmelCase__							)
   _a						: List[str]    =			{"""input_ids""": input_ids, """attention_mask""": input_mask}
   _a						: List[Any]    =			model(UpperCAmelCase__							)
   self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.seq_length, self.vocab_size)							)


 def _lowercase			(  self	:		Optional[int]							,		UpperCAmelCase__	:		Dict							,		UpperCAmelCase__	:		Optional[Any]							,		UpperCAmelCase__	:		Union[str, Any]							,		UpperCAmelCase__	:		Optional[int]							,		UpperCAmelCase__	:		Tuple							,		UpperCAmelCase__	:		Dict							)      -> int:
   _a						: List[Any]    =			TFMPNetForQuestionAnswering(config=UpperCAmelCase__							)
   _a						: str    =			{
       """input_ids""": input_ids,
       """attention_mask""": input_mask,
   }
   _a						: Optional[int]    =			model(UpperCAmelCase__							)
   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 _lowercase			(  self	:		str							,		UpperCAmelCase__	:		List[Any]							,		UpperCAmelCase__	:		List[Any]							,		UpperCAmelCase__	:		Tuple							,		UpperCAmelCase__	:		Optional[Any]							,		UpperCAmelCase__	:		Tuple							,		UpperCAmelCase__	:		Tuple							)      -> Optional[Any]:
   _a						: Optional[Any]    =			self.num_labels
   _a						: Dict    =			TFMPNetForSequenceClassification(UpperCAmelCase__							)
   _a						: int    =			{"""input_ids""": input_ids, """attention_mask""": input_mask}
   _a						: Optional[Any]    =			model(UpperCAmelCase__							)
   self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.num_labels)							)


 def _lowercase			(  self	:		Union[str, Any]							,		UpperCAmelCase__	:		Optional[int]							,		UpperCAmelCase__	:		str							,		UpperCAmelCase__	:		Union[str, Any]							,		UpperCAmelCase__	:		List[Any]							,		UpperCAmelCase__	:		Optional[int]							,		UpperCAmelCase__	:		List[Any]							)      -> str:
   _a						: List[Any]    =			self.num_choices
   _a						: int    =			TFMPNetForMultipleChoice(UpperCAmelCase__							)
   _a						: Optional[Any]    =			tf.tile(tf.expand_dims(UpperCAmelCase__							,		1							)							,		(1, self.num_choices, 1)							)
   _a						: int    =			tf.tile(tf.expand_dims(UpperCAmelCase__							,		1							)							,		(1, self.num_choices, 1)							)
   _a						: Union[str, Any]    =			{
       """input_ids""": multiple_choice_inputs_ids,
       """attention_mask""": multiple_choice_input_mask,
   }
   _a						: Optional[Any]    =			model(UpperCAmelCase__							)
   self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.num_choices)							)


 def _lowercase			(  self	:		Tuple							,		UpperCAmelCase__	:		List[str]							,		UpperCAmelCase__	:		int							,		UpperCAmelCase__	:		List[str]							,		UpperCAmelCase__	:		int							,		UpperCAmelCase__	:		Optional[int]							,		UpperCAmelCase__	:		Optional[int]							)      -> Tuple:
   _a						: Any    =			self.num_labels
   _a						: Dict    =			TFMPNetForTokenClassification(UpperCAmelCase__							)
   _a						: Union[str, Any]    =			{"""input_ids""": input_ids, """attention_mask""": input_mask}
   _a						: str    =			model(UpperCAmelCase__							)
   self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.seq_length, self.num_labels)							)







 def _lowercase			(  self	:		str							)      -> List[str]:
   _a						: Dict    =			self.prepare_config_and_inputs()
   (_a)						: Any    =			config_and_inputs
   _a						: Optional[Any]    =			{"""input_ids""": input_ids, """attention_mask""": input_mask}
   return config, inputs_dict







@require_tf
class  UpperCamelCase    (			snake_case_	,  snake_case_	,  unittest.TestCase   ):
 UpperCamelCase      : str    =      (
     (
         TFMPNetForMaskedLM,
         TFMPNetForMultipleChoice,
         TFMPNetForQuestionAnswering,
         TFMPNetForSequenceClassification,
         TFMPNetForTokenClassification,
         TFMPNetModel,
     )
     if is_tf_available()
     else ()
 )
 UpperCamelCase      : int    =      (
     {
         '''feature-extraction''': TFMPNetModel,
         '''fill-mask''': TFMPNetForMaskedLM,
         '''question-answering''': TFMPNetForQuestionAnswering,
         '''text-classification''': TFMPNetForSequenceClassification,
         '''token-classification''': TFMPNetForTokenClassification,
         '''zero-shot''': TFMPNetForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 UpperCamelCase      : Tuple    =      False
 UpperCamelCase      : Union[str, Any]    =      False


 def _lowercase			(  self	:		List[Any]							)      -> Any:
   _a						: str    =			TFMPNetModelTester(self							)
   _a						: str    =			ConfigTester(self							,		config_class=UpperCAmelCase__							,		hidden_size=37							)


 def _lowercase			(  self	:		Tuple							)      -> Dict:
   self.config_tester.run_common_tests()


 def _lowercase			(  self	:		str							)      -> Dict:
   _a						: Optional[int]    =			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mpnet_model(*UpperCAmelCase__							)


 def _lowercase			(  self	:		Any							)      -> Optional[int]:
   _a						: Tuple    =			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mpnet_for_masked_lm(*UpperCAmelCase__							)


 def _lowercase			(  self	:		int							)      -> Optional[Any]:
   _a						: List[str]    =			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mpnet_for_question_answering(*UpperCAmelCase__							)


 def _lowercase			(  self	:		Optional[int]							)      -> Any:
   _a						: Dict    =			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mpnet_for_sequence_classification(*UpperCAmelCase__							)


 def _lowercase			(  self	:		Optional[int]							)      -> List[str]:
   _a						: List[str]    =			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mpnet_for_multiple_choice(*UpperCAmelCase__							)


 def _lowercase			(  self	:		Optional[Any]							)      -> Optional[Any]:
   _a						: List[str]    =			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_mpnet_for_token_classification(*UpperCAmelCase__							)







 @slow
 def _lowercase			(  self	:		Any							)      -> List[Any]:
   for model_name in ["microsoft/mpnet-base"]:
     _a						: str    =			TFMPNetModel.from_pretrained(UpperCAmelCase__							)
     self.assertIsNotNone(UpperCAmelCase__							)







@require_tf
class  UpperCamelCase    (			unittest.TestCase   ):


 @slow
 def _lowercase			(  self	:		Union[str, Any]							)      -> str:
   _a						: List[Any]    =			TFMPNetModel.from_pretrained("""microsoft/mpnet-base"""							)
   _a						: int    =			tf.constant([[0, 1, 2, 3, 4, 5]]							)
   _a						: str    =			model(UpperCAmelCase__							)[0]

   _a						: Optional[int]    =			[1, 6, 768]
   self.assertEqual(output.shape							,		UpperCAmelCase__							)

   _a						: List[str]    =			tf.constant(
       [
           [
               [-0.1_0_6_7_1_7_2, 0.0_8_2_1_6_4_7_3, 0.0_0_2_4_5_4_3],
               [-0.0_3_4_6_5_8_7_9, 0.8_3_5_4_1_1_8, -0.0_3_2_5_2_2_8_8],
               [-0.0_6_5_6_9_4_7_6, -0.1_2_4_2_4_1_1_1, -0.0_4_9_4_4_3_6],
           ]
       ]							)
   tf.debugging.assert_near(output[:, :3, :3]							,		UpperCAmelCase__							,		atol=1E-4							)









"""simple docstring"""




import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


_snake_case					   =     1000_0000_0000
_snake_case					   =     1_0000

_snake_case   ,				_snake_case					   =     os.path.split(__file__)
_snake_case					   =     os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))

def     lowerCAmelCase__  (  UpperCamelCase__  ,     UpperCamelCase__       ):



  '''simple docstring'''
  _a						: Union[str, Any]    =			pa.Table.from_pydict({"""col""": [0] * chunk_size}       )
  _a						: Dict    =			pa.concat_tables([table] * (num_examples // chunk_size)       )
  return datasets.Dataset(UpperCamelCase__  ,     fingerprint="""table_100B"""       )






@dataclass
class  UpperCamelCase    :
 UpperCamelCase      : int
 UpperCamelCase      : int
 UpperCamelCase      : int
 UpperCamelCase      : int


 def _lowercase			(  self	:		Optional[int]							)      -> Dict:
   _a						: Tuple    =			np.random.default_rng(self.seed							)
   _a						: Dict    =			rng.integers(low=self.low							,		high=self.high							,		size=self.size							).tolist()


 def __iter__(  self	:		List[Any]							)      -> Optional[Any]:
   return iter(self._sampled_values							)







 def __len__(  self	:		Dict							)      -> Optional[Any]:
   return self.size







@get_duration
def     lowerCAmelCase__  (  UpperCamelCase__       ):



  '''simple docstring'''
  _a						: Union[str, Any]    =			dataset[0]

@get_duration
def     lowerCAmelCase__  (  UpperCamelCase__       ):



  '''simple docstring'''
  _a						: Optional[Any]    =			dataset[-1]

@get_duration
def     lowerCAmelCase__  (  UpperCamelCase__       ):



  '''simple docstring'''
  _a						: str    =			dataset[range(len(UpperCamelCase__       ) // 2  ,     len(UpperCamelCase__       ) // 2 + 1_0_2_4       )]

@get_duration
def     lowerCAmelCase__  (  UpperCamelCase__       ):



  '''simple docstring'''
  _a						: Tuple    =			dataset[RandIter(0  ,     len(UpperCamelCase__       )  ,     1_0_2_4  ,     seed=4_2       )]

def     lowerCAmelCase__  (  ):



  '''simple docstring'''
  _a						: str    =			{"""num examples""": SPEED_TEST_N_EXAMPLES}
  _a						: Dict    =			(get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
  print("""generating dataset"""       )
  _a						: List[str]    =			generate_100B_dataset(num_examples=UpperCamelCase__  ,     chunk_size=UpperCamelCase__       )
  print("""Functions"""       )
  for func in functions:
    print(func.__name__       )
    _a						: Optional[Any]    =			func(UpperCamelCase__       )

  with open(UpperCamelCase__  ,     """wb"""       ) as f:
    f.write(json.dumps(UpperCamelCase__       ).encode("""utf-8"""       )       )


if __name__ == "__main__":  # useful to run the profiler
 benchmark_table_100B()

