

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





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





    def 				snake_case_    (				self):
         lowercase__   :       List[Any]				=	ids_tensor([self.batch_size, self.seq_length]						,       self.vocab_size)

         lowercase__   :       int				=	None
         if self.use_input_mask:
              lowercase__   :       Union[str, Any]				=	random_attention_mask([self.batch_size, self.seq_length])

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

         lowercase__   :       List[Any]				=	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 				snake_case_    (				self						,       a						,       a						,       a						,       a						,       a						,       a):
         lowercase__   :       Optional[int]				=	TFDistilBertModel(config=a)
         lowercase__   :       Dict				=	{'input_ids': input_ids, 'attention_mask': input_mask}

         lowercase__   :       Optional[Any]				=	model(a)

         lowercase__   :       str				=	[input_ids, input_mask]

         lowercase__   :       Optional[Any]				=	model(a)

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





    def 				snake_case_    (				self						,       a						,       a						,       a						,       a						,       a						,       a):
         lowercase__   :       int				=	TFDistilBertForMaskedLM(config=a)
         lowercase__   :       Optional[int]				=	{'input_ids': input_ids, 'attention_mask': input_mask}
         lowercase__   :       List[str]				=	model(a)
         self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.seq_length, self.vocab_size))





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





    def 				snake_case_    (				self						,       a						,       a						,       a						,       a						,       a						,       a):
         lowercase__   :       Any				=	self.num_choices
         lowercase__   :       str				=	TFDistilBertForMultipleChoice(a)
         lowercase__   :       Union[str, Any]				=	tf.tile(tf.expand_dims(a						,       1)						,       (1, self.num_choices, 1))
         lowercase__   :       Tuple				=	tf.tile(tf.expand_dims(a						,       1)						,       (1, self.num_choices, 1))
         lowercase__   :       int				=	{
             'input_ids': multiple_choice_inputs_ids,
             'attention_mask': multiple_choice_input_mask,
         }
         lowercase__   :       Tuple				=	model(a)
         self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.num_choices))





    def 				snake_case_    (				self						,       a						,       a						,       a						,       a						,       a						,       a):
         lowercase__   :       Any				=	self.num_labels
         lowercase__   :       Any				=	TFDistilBertForTokenClassification(a)
         lowercase__   :       Tuple				=	{'input_ids': input_ids, 'attention_mask': input_mask}
         lowercase__   :       List[Any]				=	model(a)
         self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.seq_length, self.num_labels))





    def 				snake_case_    (				self):
         lowercase__   :       Tuple				=	self.prepare_config_and_inputs()
         (lowercase__)   :       Optional[Any]				=	config_and_inputs
         lowercase__   :       str				=	{'input_ids': input_ids, 'attention_mask': input_mask}
         return config, inputs_dict



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





    def 				snake_case_    (				self):
         lowercase__   :       List[str]				=	TFDistilBertModelTester(self)
         lowercase__   :       int				=	ConfigTester(self						,       config_class=a						,       dim=37)





    def 				snake_case_    (				self):
         self.config_tester.run_common_tests()





    def 				snake_case_    (				self):
         lowercase__   :       Union[str, Any]				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_model(*a)





    def 				snake_case_    (				self):
         lowercase__   :       str				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_masked_lm(*a)





    def 				snake_case_    (				self):
         lowercase__   :       List[Any]				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_question_answering(*a)





    def 				snake_case_    (				self):
         lowercase__   :       Any				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_sequence_classification(*a)





    def 				snake_case_    (				self):
         lowercase__   :       Union[str, Any]				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_multiple_choice(*a)





    def 				snake_case_    (				self):
         lowercase__   :       Dict				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_token_classification(*a)





    @slow
    def 				snake_case_    (				self):
         for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]):
              lowercase__   :       List[Any]				=	TFDistilBertModel.from_pretrained(a)
              self.assertIsNotNone(a)



@require_tf
class 				SCREAMING_SNAKE_CASE__ (unittest.TestCase					):





    @slow
    def 				snake_case_    (				self):
         lowercase__   :       List[str]				=	TFDistilBertModel.from_pretrained('distilbert-base-uncased')
         lowercase__   :       Union[str, Any]				=	tf.constant([[0, 1, 2, 3, 4, 5]])
         lowercase__   :       Union[str, Any]				=	model(a)[0]

         lowercase__   :       Any				=	[1, 6, 768]
         self.assertEqual(output.shape						,       a)

         lowercase__   :       Union[str, Any]				=	tf.constant(
             [
                 [
                     [0.19_261_885, -0.13_732_955, 0.4_119_799],
                     [0.22_150_156, -0.07_422_661, 0.39_037_204],
                     [0.22_756_018, -0.0_896_414, 0.3_701_467],
                 ]
             ])
         tf.debugging.assert_near(output[:, :3, :3]						,       a						,       atol=1e-4)




import numpy as np

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   str     ,					SCREAMING_SNAKE_CASE_		:   Optional[int]     ,					SCREAMING_SNAKE_CASE_		:   Optional[int]     ,					SCREAMING_SNAKE_CASE_		:   Union[str, Any]     ,					SCREAMING_SNAKE_CASE_		:   List[Any]      ):



     '''simple docstring'''





     lowercase__   :       List[Any]				=	int(np.ceil((x_end - xa) / h      )      )
     lowercase__   :       Dict				=	np.zeros((n + 1,)      )
     lowercase__   :       Dict				=	ya
     lowercase__   :       Union[str, Any]				=	xa

     for k in range(SCREAMING_SNAKE_CASE_      ):
          lowercase__   :       Optional[int]				=	f(SCREAMING_SNAKE_CASE_     ,					y[k]      )
          lowercase__   :       Dict				=	f(x + 0.5 * h     ,					y[k] + 0.5 * h * ka      )
          lowercase__   :       Any				=	f(x + 0.5 * h     ,					y[k] + 0.5 * h * ka      )
          lowercase__   :       List[Any]				=	f(x + h     ,					y[k] + h * ka      )
          lowercase__   :       Tuple				=	y[k] + (1 / 6) * h * (ka + 2 * ka + 2 * ka + ka)
          x += h

     return y


if __name__ == "__main__":
       import doctest

       doctest.testmod()

