
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						:				Tuple ,				_snake_case						:				str ,				):
          __lowercase						:	str					=				parent
          __lowercase						:	Optional[int]					=				13
          __lowercase						:	List[str]					=				7
          __lowercase						:	Union[str, Any]					=				True
          __lowercase						:	Any					=				True
          __lowercase						:	Dict					=				False
          __lowercase						:	int					=				True
          __lowercase						:	Optional[int]					=				99
          __lowercase						:	Any					=				32
          __lowercase						:	Dict					=				2
          __lowercase						:	List[str]					=				4
          __lowercase						:	Optional[int]					=				37
          __lowercase						:	List[str]					=				'''gelu'''
          __lowercase						:	int					=				0.1
          __lowercase						:	Optional[Any]					=				0.1
          __lowercase						:	Any					=				512
          __lowercase						:	Union[str, Any]					=				16
          __lowercase						:	Optional[int]					=				2
          __lowercase						:	List[Any]					=				0.02
          __lowercase						:	Dict					=				3
          __lowercase						:	Any					=				4
          __lowercase						:	Optional[int]					=				None




      def 						snake_case_	(					self						:				Union[str, Any] ):
          __lowercase						:	Tuple					=				ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size )

          __lowercase						:	str					=				None
          if self.use_input_mask:
              __lowercase						:	List[str]					=				random_attention_mask([self.batch_size, self.seq_length] )

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

          __lowercase						:	str					=				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						:				Optional[Any] ,				_snake_case						:				str ,				_snake_case						:				Any ,				_snake_case						:				str ,				_snake_case						:				Any ,				_snake_case						:				Dict ,				_snake_case						:				Optional[Any] ):
          __lowercase						:	Optional[Any]					=				TFDistilBertModel(config=_snake_case )
          __lowercase						:	str					=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}

          __lowercase						:	str					=				model(_snake_case )

          __lowercase						:	Optional[Any]					=				[input_ids, input_mask]

          __lowercase						:	Optional[int]					=				model(_snake_case )

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




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




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




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




      def 						snake_case_	(					self						:				str ,				_snake_case						:				int ,				_snake_case						:				Union[str, Any] ,				_snake_case						:				Dict ,				_snake_case						:				Optional[int] ,				_snake_case						:				Union[str, Any] ,				_snake_case						:				int ):
          __lowercase						:	Union[str, Any]					=				self.num_labels
          __lowercase						:	str					=				TFDistilBertForTokenClassification(_snake_case )
          __lowercase						:	Optional[int]					=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
          __lowercase						:	Any					=				model(_snake_case )
          self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.num_labels) )




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






@require_tf
class   __lowerCAmelCase		(		lowerCAmelCase_       ,       lowerCAmelCase_       ,       unittest.TestCase					):



      """simple docstring"""







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




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




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




      def 						snake_case_	(					self						:				List[str] ):
          __lowercase						:	List[str]					=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_model(*_snake_case )




      def 						snake_case_	(					self						:				int ):
          __lowercase						:	Dict					=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_masked_lm(*_snake_case )




      def 						snake_case_	(					self						:				List[str] ):
          __lowercase						:	List[Any]					=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_question_answering(*_snake_case )




      def 						snake_case_	(					self						:				Optional[Any] ):
          __lowercase						:	int					=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_sequence_classification(*_snake_case )




      def 						snake_case_	(					self						:				Optional[int] ):
          __lowercase						:	Dict					=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_multiple_choice(*_snake_case )




      def 						snake_case_	(					self						:				Any ):
          __lowercase						:	Any					=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_token_classification(*_snake_case )




      @slow
      def 						snake_case_	(					self						:				Union[str, Any] ):
          for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1] ):
              __lowercase						:	int					=				TFDistilBertModel.from_pretrained(_snake_case )
              self.assertIsNotNone(_snake_case )






@require_tf
class   __lowerCAmelCase		(		unittest.TestCase					):



      """simple docstring"""







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

          __lowercase						:	Optional[Any]					=				[1, 6, 768]
          self.assertEqual(output.shape ,				_snake_case )

          __lowercase						:	int					=				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"""
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


_lowerCAmelCase    :str     =					{
    'configuration_squeezebert': [
        'SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'SqueezeBertConfig',
        'SqueezeBertOnnxConfig',
    ],
    'tokenization_squeezebert': ['SqueezeBertTokenizer'],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase    :Optional[int]     =					['SqueezeBertTokenizerFast']

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase    :str     =					[
      'SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
      'SqueezeBertForMaskedLM',
      'SqueezeBertForMultipleChoice',
      'SqueezeBertForQuestionAnswering',
      'SqueezeBertForSequenceClassification',
      'SqueezeBertForTokenClassification',
      'SqueezeBertModel',
      'SqueezeBertModule',
      'SqueezeBertPreTrainedModel',
  ]


if TYPE_CHECKING:
  from .configuration_squeezebert import (
      SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      SqueezeBertConfig,
      SqueezeBertOnnxConfig,
  )
  from .tokenization_squeezebert import SqueezeBertTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_squeezebert_fast import SqueezeBertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_squeezebert import (
        SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        SqueezeBertForMaskedLM,
        SqueezeBertForMultipleChoice,
        SqueezeBertForQuestionAnswering,
        SqueezeBertForSequenceClassification,
        SqueezeBertForTokenClassification,
        SqueezeBertModel,
        SqueezeBertModule,
        SqueezeBertPreTrainedModel,
    )

else:
  import sys

  _lowerCAmelCase    :Optional[Any]     =					_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


