

'''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        A_ :






     '''simple docstring'''



     def __init__(		self :     Optional[Any]     ,			lowercase_ :     int     ,			)						->	List[Any]:
      UpperCAmelCase       :	Optional[Any]											= parent
      UpperCAmelCase       :	Tuple											= 13
      UpperCAmelCase       :	int											= 7
      UpperCAmelCase       :	List[Any]											= True
      UpperCAmelCase       :	List[Any]											= True
      UpperCAmelCase       :	Union[str, Any]											= False
      UpperCAmelCase       :	Tuple											= True
      UpperCAmelCase       :	Optional[Any]											= 99
      UpperCAmelCase       :	Tuple											= 32
      UpperCAmelCase       :	Any											= 2
      UpperCAmelCase       :	List[Any]											= 4
      UpperCAmelCase       :	Union[str, Any]											= 37
      UpperCAmelCase       :	Dict											= 'gelu'
      UpperCAmelCase       :	List[str]											= 0.1
      UpperCAmelCase       :	Optional[int]											= 0.1
      UpperCAmelCase       :	Tuple											= 512
      UpperCAmelCase       :	List[str]											= 16
      UpperCAmelCase       :	int											= 2
      UpperCAmelCase       :	List[Any]											= 0.02
      UpperCAmelCase       :	Dict											= 3
      UpperCAmelCase       :	List[str]											= 4
      UpperCAmelCase       :	Optional[Any]											= None


     def    UpperCAmelCase_       (		self :     int       )						->	Optional[int]:
      UpperCAmelCase       :	Union[str, Any]											= ids_tensor([self.batch_size, self.seq_length]     ,			self.vocab_size       )

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

      UpperCAmelCase       :	Tuple											= None
      UpperCAmelCase       :	Dict											= None
      UpperCAmelCase       :	List[Any]											= None
      if self.use_labels:
       UpperCAmelCase       :	List[Any]											= ids_tensor([self.batch_size]     ,			self.type_sequence_label_size       )
       UpperCAmelCase       :	List[Any]											= ids_tensor([self.batch_size, self.seq_length]     ,			self.num_labels       )
       UpperCAmelCase       :	Any											= ids_tensor([self.batch_size]     ,			self.num_choices       )

      UpperCAmelCase       :	Tuple											= 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    UpperCAmelCase_       (		self :     List[Any]     ,			lowercase_ :     List[str]     ,			lowercase_ :     Optional[int]     ,			lowercase_ :     Optional[Any]     ,			lowercase_ :     int     ,			lowercase_ :     Union[str, Any]     ,			lowercase_ :     Tuple       )						->	Union[str, Any]:
      UpperCAmelCase       :	Any											= TFDistilBertModel(config=lowerCAmelCase_       )
      UpperCAmelCase       :	Optional[Any]											= {'input_ids': input_ids, 'attention_mask': input_mask}

      UpperCAmelCase       :	List[str]											= model(lowerCAmelCase_       )

      UpperCAmelCase       :	Union[str, Any]											= [input_ids, input_mask]

      UpperCAmelCase       :	str											= model(lowerCAmelCase_       )

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


     def    UpperCAmelCase_       (		self :     Any     ,			lowercase_ :     int     ,			lowercase_ :     int     ,			lowercase_ :     Optional[int]     ,			lowercase_ :     List[Any]     ,			lowercase_ :     Union[str, Any]     ,			lowercase_ :     Optional[int]       )						->	List[str]:
      UpperCAmelCase       :	Tuple											= TFDistilBertForMaskedLM(config=lowerCAmelCase_       )
      UpperCAmelCase       :	int											= {'input_ids': input_ids, 'attention_mask': input_mask}
      UpperCAmelCase       :	int											= model(lowerCAmelCase_       )
      self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.vocab_size)       )


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


     def    UpperCAmelCase_       (		self :     List[str]     ,			lowercase_ :     Union[str, Any]     ,			lowercase_ :     Any     ,			lowercase_ :     Any     ,			lowercase_ :     List[str]     ,			lowercase_ :     Dict     ,			lowercase_ :     int       )						->	Dict:
      UpperCAmelCase       :	List[str]											= self.num_choices
      UpperCAmelCase       :	Tuple											= TFDistilBertForMultipleChoice(lowerCAmelCase_       )
      UpperCAmelCase       :	Dict											= tf.tile(tf.expand_dims(lowerCAmelCase_     ,			1       )     ,			(1, self.num_choices, 1)       )
      UpperCAmelCase       :	List[str]											= tf.tile(tf.expand_dims(lowerCAmelCase_     ,			1       )     ,			(1, self.num_choices, 1)       )
      UpperCAmelCase       :	List[Any]											= {
          'input_ids': multiple_choice_inputs_ids,
          'attention_mask': multiple_choice_input_mask,
      }
      UpperCAmelCase       :	List[Any]											= model(lowerCAmelCase_       )
      self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.num_choices)       )


     def    UpperCAmelCase_       (		self :     Union[str, Any]     ,			lowercase_ :     Optional[int]     ,			lowercase_ :     Tuple     ,			lowercase_ :     Union[str, Any]     ,			lowercase_ :     Tuple     ,			lowercase_ :     Any     ,			lowercase_ :     Optional[Any]       )						->	Optional[int]:
      UpperCAmelCase       :	Tuple											= self.num_labels
      UpperCAmelCase       :	int											= TFDistilBertForTokenClassification(lowerCAmelCase_       )
      UpperCAmelCase       :	str											= {'input_ids': input_ids, 'attention_mask': input_mask}
      UpperCAmelCase       :	str											= model(lowerCAmelCase_       )
      self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.num_labels)       )




     def    UpperCAmelCase_       (		self :     Optional[int]       )						->	str:
      UpperCAmelCase       :	Tuple											= self.prepare_config_and_inputs()
      ((UpperCAmelCase)					, (UpperCAmelCase)					, (UpperCAmelCase)					, (UpperCAmelCase)					, (UpperCAmelCase)					, (UpperCAmelCase))       :	str											= config_and_inputs
      UpperCAmelCase       :	Any											= {'input_ids': input_ids, 'attention_mask': input_mask}
      return config, inputs_dict




@require_tf
class        A_ (				_UpperCamelCase       ,   _UpperCamelCase       ,   unittest.TestCase							):






     '''simple docstring'''



     UpperCAmelCase_						:	List[str]  		=  (
         (
             TFDistilBertModel,
             TFDistilBertForMaskedLM,
             TFDistilBertForQuestionAnswering,
             TFDistilBertForSequenceClassification,
             TFDistilBertForTokenClassification,
             TFDistilBertForMultipleChoice,
         )
         if is_tf_available()
         else None
     )
     UpperCAmelCase_						:	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 {}
     )
     UpperCAmelCase_						:	Optional[Any]  		=  False
     UpperCAmelCase_						:	List[Any]  		=  False


     def    UpperCAmelCase_       (		self :     List[str]       )						->	Optional[int]:
      UpperCAmelCase       :	int											= TFDistilBertModelTester(self       )
      UpperCAmelCase       :	Dict											= ConfigTester(self     ,			config_class=lowerCAmelCase_     ,			dim=37       )


     def    UpperCAmelCase_       (		self :     Tuple       )						->	Any:
      self.config_tester.run_common_tests()


     def    UpperCAmelCase_       (		self :     Optional[int]       )						->	str:
      UpperCAmelCase       :	List[Any]											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_model(*lowerCAmelCase_       )


     def    UpperCAmelCase_       (		self :     Any       )						->	Tuple:
      UpperCAmelCase       :	Tuple											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_for_masked_lm(*lowerCAmelCase_       )


     def    UpperCAmelCase_       (		self :     int       )						->	str:
      UpperCAmelCase       :	List[Any]											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_for_question_answering(*lowerCAmelCase_       )


     def    UpperCAmelCase_       (		self :     str       )						->	int:
      UpperCAmelCase       :	Union[str, Any]											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowerCAmelCase_       )


     def    UpperCAmelCase_       (		self :     List[str]       )						->	Union[str, Any]:
      UpperCAmelCase       :	int											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowerCAmelCase_       )


     def    UpperCAmelCase_       (		self :     Dict       )						->	Optional[int]:
      UpperCAmelCase       :	Tuple											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_for_token_classification(*lowerCAmelCase_       )




     @slow
     def    UpperCAmelCase_       (		self :     List[str]       )						->	List[str]:
      for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]       ):
       UpperCAmelCase       :	List[str]											= TFDistilBertModel.from_pretrained(lowerCAmelCase_       )
       self.assertIsNotNone(lowerCAmelCase_       )




@require_tf
class        A_ (				unittest.TestCase							):






     '''simple docstring'''



     @slow
     def    UpperCAmelCase_       (		self :     Optional[int]       )						->	List[Any]:
      UpperCAmelCase       :	int											= TFDistilBertModel.from_pretrained('distilbert-base-uncased'       )
      UpperCAmelCase       :	List[str]											= tf.constant([[0, 1, 2, 3, 4, 5]]       )
      UpperCAmelCase       :	int											= model(lowerCAmelCase_       )[0]

      UpperCAmelCase       :	int											= [1, 6, 768]
      self.assertEqual(output.shape     ,			lowerCAmelCase_       )

      UpperCAmelCase       :	Any											= tf.constant(
          [
              [
                  [0.1926_1885, -0.1373_2955, 0.411_9799],
                  [0.2215_0156, -0.0742_2661, 0.3903_7204],
                  [0.2275_6018, -0.089_6414, 0.370_1467],
              ]
          ]       )
      tf.debugging.assert_near(output[:, :3, :3]     ,			lowerCAmelCase_     ,			atol=1E-4       )











def   a_    ( lowerCAmelCase_      :			str,	lowerCAmelCase_      :			int	):
   return [sentence[i : i + ngram_size] for i in range(len(lowerCAmelCase_	) - ngram_size + 1	)]


if __name__ == "__main__":
     from doctest import testmod

     testmod()

