





'''simple docstring'''






from __future__ import annotations

import unittest

from transformers import DistilBertConfig, 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.distilbert.modeling_tf_distilbert import (
          TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFDistilBertForMaskedLM,
          TFDistilBertForMultipleChoice,
          TFDistilBertForQuestionAnswering,
          TFDistilBertForSequenceClassification,
          TFDistilBertForTokenClassification,
          TFDistilBertModel,
      )




class 			UpperCAmelCase					:




      def __init__(      self      :			Dict      ,      __snake_case      :			Union[str, Any]      ,      )    ->				str:
            _lowerCAmelCase												=							parent
            _lowerCAmelCase												=							13
            _lowerCAmelCase												=							7
            _lowerCAmelCase												=							True
            _lowerCAmelCase												=							True
            _lowerCAmelCase												=							False
            _lowerCAmelCase												=							True
            _lowerCAmelCase												=							99
            _lowerCAmelCase												=							32
            _lowerCAmelCase												=							2
            _lowerCAmelCase												=							4
            _lowerCAmelCase												=							37
            _lowerCAmelCase												=							"""gelu"""
            _lowerCAmelCase												=							0.1
            _lowerCAmelCase												=							0.1
            _lowerCAmelCase												=							5_12
            _lowerCAmelCase												=							16
            _lowerCAmelCase												=							2
            _lowerCAmelCase												=							0.02
            _lowerCAmelCase												=							3
            _lowerCAmelCase												=							4
            _lowerCAmelCase												=							None




      def     lowercase__   (      self      :			int  )    ->				List[Any]:
            _lowerCAmelCase												=							ids_tensor([self.batch_size, self.seq_length]      ,      self.vocab_size  )

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

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

            _lowerCAmelCase												=							DistilBertConfig(
                vocab_size=self.vocab_size      ,      dim=self.hidden_size      ,      n_layers=self.num_hidden_layers      ,      n_heads=self.num_attention_heads      ,      hidden_dim=self.intermediate_size      ,      hidden_act=self.hidden_act      ,      dropout=self.hidden_dropout_prob      ,      attention_dropout=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      :			Any      ,      __snake_case      :			Optional[Any]      ,      __snake_case      :			Tuple      ,      __snake_case      :			List[Any]      ,      __snake_case      :			Tuple      ,      __snake_case      :			Optional[int]      ,      __snake_case      :			Dict  )    ->				Union[str, Any]:
            _lowerCAmelCase												=							TFDistilBertModel(config=__snake_case  )
            _lowerCAmelCase												=							{"""input_ids""": input_ids, """attention_mask""": input_mask}

            _lowerCAmelCase												=							model(__snake_case  )

            _lowerCAmelCase												=							[input_ids, input_mask]

            _lowerCAmelCase												=							model(__snake_case  )

            self.parent.assertEqual(result.last_hidden_state.shape      ,      (self.batch_size, self.seq_length, self.hidden_size)  )




      def     lowercase__   (      self      :			Optional[Any]      ,      __snake_case      :			Union[str, Any]      ,      __snake_case      :			Any      ,      __snake_case      :			Dict      ,      __snake_case      :			Tuple      ,      __snake_case      :			Any      ,      __snake_case      :			str  )    ->				Optional[Any]:
            _lowerCAmelCase												=							TFDistilBertForMaskedLM(config=__snake_case  )
            _lowerCAmelCase												=							{"""input_ids""": input_ids, """attention_mask""": input_mask}
            _lowerCAmelCase												=							model(__snake_case  )
            self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)  )




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




      def     lowercase__   (      self      :			Optional[int]      ,      __snake_case      :			Union[str, Any]      ,      __snake_case      :			Any      ,      __snake_case      :			Dict      ,      __snake_case      :			Any      ,      __snake_case      :			str      ,      __snake_case      :			Tuple  )    ->				Optional[Any]:
            _lowerCAmelCase												=							self.num_choices
            _lowerCAmelCase												=							TFDistilBertForMultipleChoice(__snake_case  )
            _lowerCAmelCase												=							tf.tile(tf.expand_dims(__snake_case      ,      1  )      ,      (1, self.num_choices, 1)  )
            _lowerCAmelCase												=							tf.tile(tf.expand_dims(__snake_case      ,      1  )      ,      (1, self.num_choices, 1)  )
            _lowerCAmelCase												=							{
                """input_ids""": multiple_choice_inputs_ids,
                """attention_mask""": multiple_choice_input_mask,
            }
            _lowerCAmelCase												=							model(__snake_case  )
            self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.num_choices)  )




      def     lowercase__   (      self      :			Optional[Any]      ,      __snake_case      :			List[str]      ,      __snake_case      :			int      ,      __snake_case      :			List[Any]      ,      __snake_case      :			Any      ,      __snake_case      :			List[str]      ,      __snake_case      :			int  )    ->				List[str]:
            _lowerCAmelCase												=							self.num_labels
            _lowerCAmelCase												=							TFDistilBertForTokenClassification(__snake_case  )
            _lowerCAmelCase												=							{"""input_ids""": input_ids, """attention_mask""": input_mask}
            _lowerCAmelCase												=							model(__snake_case  )
            self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.num_labels)  )




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




@require_tf
class 			UpperCAmelCase					(      snake_case_							,		snake_case_							,		unittest.TestCase       ):
      _lowercase:      Dict											=       (
          (
              TFDistilBertModel,
              TFDistilBertForMaskedLM,
              TFDistilBertForQuestionAnswering,
              TFDistilBertForSequenceClassification,
              TFDistilBertForTokenClassification,
              TFDistilBertForMultipleChoice,
          )
          if is_tf_available()
          else None
      )
      _lowercase:      List[str]											=       (
          {
              '''feature-extraction''': TFDistilBertModel,
              '''fill-mask''': TFDistilBertForMaskedLM,
              '''question-answering''': TFDistilBertForQuestionAnswering,
              '''text-classification''': TFDistilBertForSequenceClassification,
              '''token-classification''': TFDistilBertForTokenClassification,
              '''zero-shot''': TFDistilBertForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      _lowercase:      Optional[Any]											=       False
      _lowercase:      int											=       False




      def     lowercase__   (      self      :			Dict  )    ->				Optional[Any]:
            _lowerCAmelCase												=							TFDistilBertModelTester(self  )
            _lowerCAmelCase												=							ConfigTester(self      ,      config_class=__snake_case      ,      dim=37  )




      def     lowercase__   (      self      :			Optional[int]  )    ->				str:
            self.config_tester.run_common_tests()




      def     lowercase__   (      self      :			Any  )    ->				List[Any]:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_model(*__snake_case  )




      def     lowercase__   (      self      :			Any  )    ->				Optional[Any]:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_masked_lm(*__snake_case  )




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				Optional[int]:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_question_answering(*__snake_case  )




      def     lowercase__   (      self      :			Any  )    ->				Dict:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_sequence_classification(*__snake_case  )




      def     lowercase__   (      self      :			Tuple  )    ->				Any:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_multiple_choice(*__snake_case  )




      def     lowercase__   (      self      :			int  )    ->				Tuple:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_token_classification(*__snake_case  )




      @slow
      def     lowercase__   (      self      :			Optional[int]  )    ->				Optional[int]:
            for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]  ):
                  _lowerCAmelCase												=							TFDistilBertModel.from_pretrained(__snake_case  )
                  self.assertIsNotNone(__snake_case  )




@require_tf
class 			UpperCAmelCase					(      unittest.TestCase       ):




      @slow
      def     lowercase__   (      self      :			Optional[int]  )    ->				str:
            _lowerCAmelCase												=							TFDistilBertModel.from_pretrained("""distilbert-base-uncased"""  )
            _lowerCAmelCase												=							tf.constant([[0, 1, 2, 3, 4, 5]]  )
            _lowerCAmelCase												=							model(__snake_case  )[0]

            _lowerCAmelCase												=							[1, 6, 7_68]
            self.assertEqual(output.shape      ,      __snake_case  )

            _lowerCAmelCase												=							tf.constant(
                [
                    [
                        [0.19_26_18_85, -0.13_73_29_55, 0.4_11_97_99],
                        [0.22_15_01_56, -0.07_42_26_61, 0.39_03_72_04],
                        [0.22_75_60_18, -0.0_89_64_14, 0.3_70_14_67],
                    ]
                ]  )
            tf.debugging.assert_near(output[:, :3, :3]      ,      __snake_case      ,      atol=1E-4  )













'''simple docstring'''






import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


A__    :       str              =5_00_00
A__    :       Optional[int]              =50_00

A__       ,						A__    :       Optional[int]              =os.path.split(__file__)
A__    :       Tuple              =os.path.join(RESULTS_BASEPATH, '''results''', RESULTS_FILENAME.replace('''.py''', '''.json'''))
@get_duration
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      for i in range(lowerCAmelCase  ):
            _lowerCAmelCase												=							dataset[i]
@get_duration
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      for i in range(0							,     len(lowerCAmelCase  )							,     lowerCAmelCase  ):
            _lowerCAmelCase												=							dataset[i : i + batch_size]
@get_duration
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      with dataset.formatted_as(type=lowerCAmelCase  ):
            for i in range(lowerCAmelCase  ):
                  _lowerCAmelCase												=							dataset[i]
@get_duration
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      with dataset.formatted_as(type=lowerCAmelCase  ):
            for i in range(0							,     lowerCAmelCase							,     lowerCAmelCase  ):
                  _lowerCAmelCase												=							dataset[i : i + batch_size]
def        UpperCamelCase__				(		):





      """simple docstring"""


      _lowerCAmelCase												=							{"""num examples""": SPEED_TEST_N_EXAMPLES}
      _lowerCAmelCase												=							[
          (read, {"""length""": SMALL_TEST}),
          (read, {"""length""": SPEED_TEST_N_EXAMPLES}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 1_00}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10_00}),
          (read_formatted, {"""type""": """numpy""", """length""": SMALL_TEST}),
          (read_formatted, {"""type""": """pandas""", """length""": SMALL_TEST}),
          (read_formatted, {"""type""": """torch""", """length""": SMALL_TEST}),
          (read_formatted, {"""type""": """tensorflow""", """length""": SMALL_TEST}),
          (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10}),
          (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10_00}),
      ]

      _lowerCAmelCase												=							[
          (read, {"""length""": SMALL_TEST}),
          (read, {"""length""": SPEED_TEST_N_EXAMPLES}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 1_00}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10_00}),
          (read_formatted, {"""type""": """numpy""", """length""": SMALL_TEST}),
          (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10}),
          (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10_00}),
      ]
      with tempfile.TemporaryDirectory() as tmp_dir:
            print("""generating dataset"""  )
            _lowerCAmelCase												=							datasets.Features(
                {"""list""": datasets.Sequence(datasets.Value("""float32"""  )  ), """numbers""": datasets.Value("""float32"""  )}  )
            _lowerCAmelCase												=							generate_example_dataset(
                os.path.join(lowerCAmelCase							,     """dataset.arrow"""  )							,     lowerCAmelCase							,     num_examples=lowerCAmelCase							,     seq_shapes={"""list""": (1_00,)}							,     )
            print("""first set of iterations"""  )
            for func, kwargs in functions:
                  print(func.__name__							,     str(lowerCAmelCase  )  )
                  _lowerCAmelCase												=							func(lowerCAmelCase							,     **lowerCAmelCase  )

            print("""shuffling dataset"""  )
            _lowerCAmelCase												=							dataset.shuffle()
            print("""Second set of iterations (after shuffling"""  )
            for func, kwargs in functions_shuffled:
                  print("""shuffled """							,     func.__name__							,     str(lowerCAmelCase  )  )
                  _lowerCAmelCase												=							func(
                      lowerCAmelCase							,     **lowerCAmelCase  )

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


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






