
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						,       snake_case						,       ):
		lowercase	    =       parent
		lowercase	    =       13
		lowercase	    =       7
		lowercase	    =       True
		lowercase	    =       True
		lowercase	    =       False
		lowercase	    =       True
		lowercase	    =       99
		lowercase	    =       32
		lowercase	    =       2
		lowercase	    =       4
		lowercase	    =       37
		lowercase	    =       'gelu'
		lowercase	    =       0.1
		lowercase	    =       0.1
		lowercase	    =       512
		lowercase	    =       16
		lowercase	    =       2
		lowercase	    =       0.02
		lowercase	    =       3
		lowercase	    =       4
		lowercase	    =       None





	def       SCREAMING_SNAKE_CASE__		(			self			):
		lowercase	    =       ids_tensor([self.batch_size, self.seq_length]						,       self.vocab_size			)

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

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

		lowercase	    =       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       SCREAMING_SNAKE_CASE__		(			self						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case			):
		lowercase	    =       TFDistilBertModel(config=snake_case			)
		lowercase	    =       {'input_ids': input_ids, 'attention_mask': input_mask}

		lowercase	    =       model(snake_case			)

		lowercase	    =       [input_ids, input_mask]

		lowercase	    =       model(snake_case			)

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





	def       SCREAMING_SNAKE_CASE__		(			self						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case			):
		lowercase	    =       TFDistilBertForMaskedLM(config=snake_case			)
		lowercase	    =       {'input_ids': input_ids, 'attention_mask': input_mask}
		lowercase	    =       model(snake_case			)
		self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.seq_length, self.vocab_size)			)





	def       SCREAMING_SNAKE_CASE__		(			self						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case			):
		lowercase	    =       TFDistilBertForQuestionAnswering(config=snake_case			)
		lowercase	    =       {
		    'input_ids': input_ids,
		    'attention_mask': input_mask,
		}
		lowercase	    =       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       SCREAMING_SNAKE_CASE__		(			self						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case			):
		lowercase	    =       self.num_labels
		lowercase	    =       TFDistilBertForSequenceClassification(snake_case			)
		lowercase	    =       {'input_ids': input_ids, 'attention_mask': input_mask}
		lowercase	    =       model(snake_case			)
		self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.num_labels)			)





	def       SCREAMING_SNAKE_CASE__		(			self						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case			):
		lowercase	    =       self.num_choices
		lowercase	    =       TFDistilBertForMultipleChoice(snake_case			)
		lowercase	    =       tf.tile(tf.expand_dims(snake_case						,       1			)						,       (1, self.num_choices, 1)			)
		lowercase	    =       tf.tile(tf.expand_dims(snake_case						,       1			)						,       (1, self.num_choices, 1)			)
		lowercase	    =       {
		    'input_ids': multiple_choice_inputs_ids,
		    'attention_mask': multiple_choice_input_mask,
		}
		lowercase	    =       model(snake_case			)
		self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.num_choices)			)





	def       SCREAMING_SNAKE_CASE__		(			self						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case			):
		lowercase	    =       self.num_labels
		lowercase	    =       TFDistilBertForTokenClassification(snake_case			)
		lowercase	    =       {'input_ids': input_ids, 'attention_mask': input_mask}
		lowercase	    =       model(snake_case			)
		self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.seq_length, self.num_labels)			)





	def       SCREAMING_SNAKE_CASE__		(			self			):
		lowercase	    =       self.prepare_config_and_inputs()
		((lowercase)    ,   (lowercase)    ,   (lowercase)    ,   (lowercase)    ,   (lowercase)    ,   (lowercase))	    =       config_and_inputs
		lowercase	    =       {'input_ids': input_ids, 'attention_mask': input_mask}
		return config, inputs_dict







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




	'''simple docstring'''


	_UpperCamelCase	:	List[Any]									=			(
	    (
	        TFDistilBertModel,
	        TFDistilBertForMaskedLM,
	        TFDistilBertForQuestionAnswering,
	        TFDistilBertForSequenceClassification,
	        TFDistilBertForTokenClassification,
	        TFDistilBertForMultipleChoice,
	    )
	    if is_tf_available()
	    else None
	)
	_UpperCamelCase	:	int									=			(
	    {
	        """feature-extraction""": TFDistilBertModel,
	        """fill-mask""": TFDistilBertForMaskedLM,
	        """question-answering""": TFDistilBertForQuestionAnswering,
	        """text-classification""": TFDistilBertForSequenceClassification,
	        """token-classification""": TFDistilBertForTokenClassification,
	        """zero-shot""": TFDistilBertForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	_UpperCamelCase	:	List[str]									=			False
	_UpperCamelCase	:	Dict									=			False





	def       SCREAMING_SNAKE_CASE__		(			self			):
		lowercase	    =       TFDistilBertModelTester(self			)
		lowercase	    =       ConfigTester(self						,       config_class=snake_case						,       dim=37			)





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





	def       SCREAMING_SNAKE_CASE__		(			self			):
		lowercase	    =       self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_model(*snake_case			)





	def       SCREAMING_SNAKE_CASE__		(			self			):
		lowercase	    =       self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_for_masked_lm(*snake_case			)





	def       SCREAMING_SNAKE_CASE__		(			self			):
		lowercase	    =       self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_for_question_answering(*snake_case			)





	def       SCREAMING_SNAKE_CASE__		(			self			):
		lowercase	    =       self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_for_sequence_classification(*snake_case			)





	def       SCREAMING_SNAKE_CASE__		(			self			):
		lowercase	    =       self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_for_multiple_choice(*snake_case			)





	def       SCREAMING_SNAKE_CASE__		(			self			):
		lowercase	    =       self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_for_token_classification(*snake_case			)





	@slow
	def       SCREAMING_SNAKE_CASE__		(			self			):
		for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]			):
			lowercase	    =       TFDistilBertModel.from_pretrained(snake_case			)
			self.assertIsNotNone(snake_case			)







@require_tf
class 			A_  (     unittest.TestCase  ):




	'''simple docstring'''


	@slow
	def       SCREAMING_SNAKE_CASE__		(			self			):
		lowercase	    =       TFDistilBertModel.from_pretrained('distilbert-base-uncased'			)
		lowercase	    =       tf.constant([[0, 1, 2, 3, 4, 5]]			)
		lowercase	    =       model(snake_case			)[0]

		lowercase	    =       [1, 6, 768]
		self.assertEqual(output.shape						,       snake_case			)

		lowercase	    =       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]						,       snake_case						,       atol=1E-4			)








"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE__   =  {
    "configuration_roformer": ["ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "RoFormerConfig", "RoFormerOnnxConfig"],
    "tokenization_roformer": ["RoFormerTokenizer"],
}

try:
       if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  ["RoFormerTokenizerFast"]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "RoFormerForCausalLM",
           "RoFormerForMaskedLM",
           "RoFormerForMultipleChoice",
           "RoFormerForQuestionAnswering",
           "RoFormerForSequenceClassification",
           "RoFormerForTokenClassification",
           "RoFormerLayer",
           "RoFormerModel",
           "RoFormerPreTrainedModel",
           "load_tf_weights_in_roformer",
       ]


try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TFRoFormerForCausalLM",
           "TFRoFormerForMaskedLM",
           "TFRoFormerForMultipleChoice",
           "TFRoFormerForQuestionAnswering",
           "TFRoFormerForSequenceClassification",
           "TFRoFormerForTokenClassification",
           "TFRoFormerLayer",
           "TFRoFormerModel",
           "TFRoFormerPreTrainedModel",
       ]


try:
       if not is_flax_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "FLAX_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "FlaxRoFormerForMaskedLM",
           "FlaxRoFormerForMultipleChoice",
           "FlaxRoFormerForQuestionAnswering",
           "FlaxRoFormerForSequenceClassification",
           "FlaxRoFormerForTokenClassification",
           "FlaxRoFormerModel",
           "FlaxRoFormerPreTrainedModel",
       ]


if TYPE_CHECKING:
       from .configuration_roformer import ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, RoFormerConfig, RoFormerOnnxConfig
       from .tokenization_roformer import RoFormerTokenizer

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_roformer_fast import RoFormerTokenizerFast

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_roformer import (
                  ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  RoFormerForCausalLM,
                  RoFormerForMaskedLM,
                  RoFormerForMultipleChoice,
                  RoFormerForQuestionAnswering,
                  RoFormerForSequenceClassification,
                  RoFormerForTokenClassification,
                  RoFormerLayer,
                  RoFormerModel,
                  RoFormerPreTrainedModel,
                  load_tf_weights_in_roformer,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_roformer import (
                  TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFRoFormerForCausalLM,
                  TFRoFormerForMaskedLM,
                  TFRoFormerForMultipleChoice,
                  TFRoFormerForQuestionAnswering,
                  TFRoFormerForSequenceClassification,
                  TFRoFormerForTokenClassification,
                  TFRoFormerLayer,
                  TFRoFormerModel,
                  TFRoFormerPreTrainedModel,
              )

       try:
              if not is_flax_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_flax_roformer import (
                  FLAX_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  FlaxRoFormerForMaskedLM,
                  FlaxRoFormerForMultipleChoice,
                  FlaxRoFormerForQuestionAnswering,
                  FlaxRoFormerForSequenceClassification,
                  FlaxRoFormerForTokenClassification,
                  FlaxRoFormerModel,
                  FlaxRoFormerPreTrainedModel,
              )


else:
       import sys

       SCREAMING_SNAKE_CASE__   =  _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






