

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



       '''simple docstring'''







       def __init__(							self ,   lowercase__ ,   ):







             '''simple docstring'''



             __A    =parent
             __A    =1_3
             __A    =7
             __A    =True
             __A    =True
             __A    =False
             __A    =True
             __A    =9_9
             __A    =3_2
             __A    =2
             __A    =4
             __A    =3_7
             __A    ="""gelu"""
             __A    =0.1
             __A    =0.1
             __A    =5_1_2
             __A    =1_6
             __A    =2
             __A    =0.02
             __A    =3
             __A    =4
             __A    =None

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =ids_tensor([self.batch_size, self.seq_length] ,   self.vocab_size     )

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

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

             __A    =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 ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =TFDistilBertModel(config=lowerCAmelCase_     )
             __A    ={"""input_ids""": input_ids, """attention_mask""": input_mask}

             __A    =model(lowerCAmelCase_     )

             __A    =[input_ids, input_mask]

             __A    =model(lowerCAmelCase_     )

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

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =TFDistilBertForMaskedLM(config=lowerCAmelCase_     )
             __A    ={"""input_ids""": input_ids, """attention_mask""": input_mask}
             __A    =model(lowerCAmelCase_     )
             self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.seq_length, self.vocab_size)     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =TFDistilBertForQuestionAnswering(config=lowerCAmelCase_     )
             __A    ={
                 """input_ids""": input_ids,
                 """attention_mask""": input_mask,
             }
             __A    =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 ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =self.num_labels
             __A    =TFDistilBertForSequenceClassification(lowerCAmelCase_     )
             __A    ={"""input_ids""": input_ids, """attention_mask""": input_mask}
             __A    =model(lowerCAmelCase_     )
             self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.num_labels)     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =self.num_choices
             __A    =TFDistilBertForMultipleChoice(lowerCAmelCase_     )
             __A    =tf.tile(tf.expand_dims(lowerCAmelCase_ ,   1     ) ,   (1, self.num_choices, 1)     )
             __A    =tf.tile(tf.expand_dims(lowerCAmelCase_ ,   1     ) ,   (1, self.num_choices, 1)     )
             __A    ={
                 """input_ids""": multiple_choice_inputs_ids,
                 """attention_mask""": multiple_choice_input_mask,
             }
             __A    =model(lowerCAmelCase_     )
             self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.num_choices)     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =self.num_labels
             __A    =TFDistilBertForTokenClassification(lowerCAmelCase_     )
             __A    ={"""input_ids""": input_ids, """attention_mask""": input_mask}
             __A    =model(lowerCAmelCase_     )
             self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.seq_length, self.num_labels)     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.prepare_config_and_inputs()
             (__A)    =config_and_inputs
             __A    ={"""input_ids""": input_ids, """attention_mask""": input_mask}
             return config, inputs_dict





@require_tf
class        lowerCAmelCase__  (					_SCREAMING_SNAKE_CASE				,							_SCREAMING_SNAKE_CASE				,							unittest.TestCase			):



       '''simple docstring'''







       lowercase_						=    (
           (
               TFDistilBertModel,
               TFDistilBertForMaskedLM,
               TFDistilBertForQuestionAnswering,
               TFDistilBertForSequenceClassification,
               TFDistilBertForTokenClassification,
               TFDistilBertForMultipleChoice,
           )
           if is_tf_available()
           else None
       )
       lowercase_						=    (
           {
               """feature-extraction""": TFDistilBertModel,
               """fill-mask""": TFDistilBertForMaskedLM,
               """question-answering""": TFDistilBertForQuestionAnswering,
               """text-classification""": TFDistilBertForSequenceClassification,
               """token-classification""": TFDistilBertForTokenClassification,
               """zero-shot""": TFDistilBertForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )
       lowercase_						=    False
       lowercase_						=    False

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =TFDistilBertModelTester(self     )
             __A    =ConfigTester(self ,   config_class=lowerCAmelCase_ ,   dim=3_7     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             self.config_tester.run_common_tests()

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_model(*lowerCAmelCase_     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_for_masked_lm(*lowerCAmelCase_     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_for_question_answering(*lowerCAmelCase_     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowerCAmelCase_     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowerCAmelCase_     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_for_token_classification(*lowerCAmelCase_     )


       @slow
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]     ):
                   __A    =TFDistilBertModel.from_pretrained(lowerCAmelCase_     )
                   self.assertIsNotNone(lowerCAmelCase_     )





@require_tf
class        lowerCAmelCase__  (					unittest.TestCase			):



       '''simple docstring'''







       @slow
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =TFDistilBertModel.from_pretrained('''distilbert-base-uncased'''     )
             __A    =tf.constant([[0, 1, 2, 3, 4, 5]]     )
             __A    =model(lowerCAmelCase_     )[0]

             __A    =[1, 6, 7_6_8]
             self.assertEqual(output.shape ,   lowerCAmelCase_     )

             __A    =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     )










"""simple docstring"""
from typing import Dict

from .base import GenericTensor, Pipeline



class   SCREAMING_SNAKE_CASE	(					_SCREAMING_SNAKE_CASE    ):





	"""simple docstring"""




	def  lowerCamelCase(self  ,  lowerCAmelCase_=None  ,  lowerCAmelCase_=None  ,  lowerCAmelCase_=None  ,  **lowerCAmelCase_				):
		if tokenize_kwargs is None:
			A_		:    Dict						     = {}

		if truncation is not None:
			if "truncation" in tokenize_kwargs:
				raise ValueError(
				    """truncation parameter defined twice (given as keyword argument as well as in tokenize_kwargs)"""				)
			A_		:    Optional[Any]						     = truncation

		A_		:    Optional[int]						     = tokenize_kwargs

		A_		:    Optional[Any]						     = {}
		if return_tensors is not None:
			A_		:    Optional[Any]						     = return_tensors

		return preprocess_params, {}, postprocess_params





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  **lowerCAmelCase_				):
		A_		:    List[Any]						     = self.framework
		A_		:    Optional[int]						     = self.tokenizer(lowerCAmelCase_  ,  return_tensors=lowerCAmelCase_  ,  **lowerCAmelCase_				)
		return model_inputs





	def  lowerCamelCase(self  ,  lowerCAmelCase_				):
		A_		:    List[Any]						     = self.model(**lowerCAmelCase_				)
		return model_outputs





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_=False				):
		# [0] is the first available tensor, logits or last_hidden_state.
		if return_tensors:
			return model_outputs[0]
		if self.framework == "pt":
			return model_outputs[0].tolist()
		elif self.framework == "tf":
			return model_outputs[0].numpy().tolist()





	def __call__(self  ,  *lowerCAmelCase_  ,  **lowerCAmelCase_				):
		return super().__call__(*lowerCAmelCase_  ,  **lowerCAmelCase_				)

