


"""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 _lowerCAmelCase							:







   def __init__(      self		,      UpperCamelCase__		,      )      ->						Union[str, Any]:

       '''simple docstring'''




       snake_case       : Optional[Any]    =       parent
       snake_case       : int    =       13
       snake_case       : Optional[int]    =       7
       snake_case       : List[Any]    =       True
       snake_case       : Tuple    =       True
       snake_case       : str    =       False
       snake_case       : Optional[int]    =       True
       snake_case       : List[Any]    =       99
       snake_case       : int    =       32
       snake_case       : Any    =       2
       snake_case       : Any    =       4
       snake_case       : Tuple    =       37
       snake_case       : List[Any]    =       "gelu"
       snake_case       : List[str]    =       0.1
       snake_case       : Dict    =       0.1
       snake_case       : str    =       512
       snake_case       : Optional[int]    =       16
       snake_case       : str    =       2
       snake_case       : Any    =       0.02
       snake_case       : Optional[int]    =       3
       snake_case       : Tuple    =       4
       snake_case       : Union[str, Any]    =       None







   def 					lowerCamelCase						(      self		)      ->						List[Any]:

       '''simple docstring'''




       snake_case       : Tuple    =       ids_tensor([self.batch_size, self.seq_length]		,      self.vocab_size		)

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

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

       snake_case       : Dict    =       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 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)      ->						Union[str, Any]:

       '''simple docstring'''




       snake_case       : Dict    =       TFDistilBertModel(config=UpperCamelCase__		)
       snake_case       : Union[str, Any]    =       {"input_ids": input_ids, "attention_mask": input_mask}

       snake_case       : str    =       model(UpperCamelCase__		)

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

       snake_case       : int    =       model(UpperCamelCase__		)

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







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)      ->						List[Any]:

       '''simple docstring'''




       snake_case       : str    =       TFDistilBertForMaskedLM(config=UpperCamelCase__		)
       snake_case       : int    =       {"input_ids": input_ids, "attention_mask": input_mask}
       snake_case       : Optional[Any]    =       model(UpperCamelCase__		)
       self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.vocab_size)		)







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)      ->						int:

       '''simple docstring'''




       snake_case       : List[Any]    =       TFDistilBertForQuestionAnswering(config=UpperCamelCase__		)
       snake_case       : int    =       {
           "input_ids": input_ids,
           "attention_mask": input_mask,
       }
       snake_case       : Optional[int]    =       model(UpperCamelCase__		)
       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 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)      ->						int:

       '''simple docstring'''




       snake_case       : Tuple    =       self.num_labels
       snake_case       : List[str]    =       TFDistilBertForSequenceClassification(UpperCamelCase__		)
       snake_case       : str    =       {"input_ids": input_ids, "attention_mask": input_mask}
       snake_case       : List[Any]    =       model(UpperCamelCase__		)
       self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_labels)		)







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)      ->						Dict:

       '''simple docstring'''




       snake_case       : int    =       self.num_choices
       snake_case       : List[str]    =       TFDistilBertForMultipleChoice(UpperCamelCase__		)
       snake_case       : Optional[int]    =       tf.tile(tf.expand_dims(UpperCamelCase__		,      1		)		,      (1, self.num_choices, 1)		)
       snake_case       : Any    =       tf.tile(tf.expand_dims(UpperCamelCase__		,      1		)		,      (1, self.num_choices, 1)		)
       snake_case       : Any    =       {
           "input_ids": multiple_choice_inputs_ids,
           "attention_mask": multiple_choice_input_mask,
       }
       snake_case       : Optional[int]    =       model(UpperCamelCase__		)
       self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_choices)		)







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)      ->						Optional[int]:

       '''simple docstring'''




       snake_case       : str    =       self.num_labels
       snake_case       : str    =       TFDistilBertForTokenClassification(UpperCamelCase__		)
       snake_case       : Any    =       {"input_ids": input_ids, "attention_mask": input_mask}
       snake_case       : List[Any]    =       model(UpperCamelCase__		)
       self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.num_labels)		)







   def 					lowerCamelCase						(      self		)      ->						Tuple:

       '''simple docstring'''




       snake_case       : Any    =       self.prepare_config_and_inputs()
       ((snake_case)   ,(snake_case)   ,(snake_case)   ,(snake_case)   ,(snake_case)   ,(snake_case))       : Dict    =       config_and_inputs
       snake_case       : Optional[Any]    =       {"input_ids": input_ids, "attention_mask": input_mask}
       return config, inputs_dict







@require_tf
class _lowerCAmelCase							( snake_case_       ,				snake_case_       ,				unittest.TestCase      ):
   __UpperCAmelCase					:			Tuple				    =    (
       (
           TFDistilBertModel,
           TFDistilBertForMaskedLM,
           TFDistilBertForQuestionAnswering,
           TFDistilBertForSequenceClassification,
           TFDistilBertForTokenClassification,
           TFDistilBertForMultipleChoice,
       )
       if is_tf_available()
       else None
   )
   __UpperCAmelCase					:			Optional[Any]				    =    (
       {
           '''feature-extraction''': TFDistilBertModel,
           '''fill-mask''': TFDistilBertForMaskedLM,
           '''question-answering''': TFDistilBertForQuestionAnswering,
           '''text-classification''': TFDistilBertForSequenceClassification,
           '''token-classification''': TFDistilBertForTokenClassification,
           '''zero-shot''': TFDistilBertForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   __UpperCAmelCase					:			str				    =    False
   __UpperCAmelCase					:			int				    =    False







   def 					lowerCamelCase						(      self		)      ->						Any:

       '''simple docstring'''




       snake_case       : List[Any]    =       TFDistilBertModelTester(self		)
       snake_case       : Optional[int]    =       ConfigTester(self		,      config_class=UpperCamelCase__		,      dim=37		)







   def 					lowerCamelCase						(      self		)      ->						Union[str, Any]:

       '''simple docstring'''




       self.config_tester.run_common_tests()







   def 					lowerCamelCase						(      self		)      ->						List[Any]:

       '''simple docstring'''




       snake_case       : Any    =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_model(*UpperCamelCase__		)







   def 					lowerCamelCase						(      self		)      ->						Optional[Any]:

       '''simple docstring'''




       snake_case       : Union[str, Any]    =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_masked_lm(*UpperCamelCase__		)







   def 					lowerCamelCase						(      self		)      ->						Any:

       '''simple docstring'''




       snake_case       : Tuple    =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_question_answering(*UpperCamelCase__		)







   def 					lowerCamelCase						(      self		)      ->						Optional[int]:

       '''simple docstring'''




       snake_case       : List[str]    =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_sequence_classification(*UpperCamelCase__		)







   def 					lowerCamelCase						(      self		)      ->						str:

       '''simple docstring'''




       snake_case       : int    =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_multiple_choice(*UpperCamelCase__		)







   def 					lowerCamelCase						(      self		)      ->						Optional[Any]:

       '''simple docstring'''




       snake_case       : Tuple    =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_token_classification(*UpperCamelCase__		)







   @slow
   def 					lowerCamelCase						(      self		)      ->						Dict:

       '''simple docstring'''




       for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]		):
           snake_case       : List[Any]    =       TFDistilBertModel.from_pretrained(UpperCamelCase__		)
           self.assertIsNotNone(UpperCamelCase__		)







@require_tf
class _lowerCAmelCase							( unittest.TestCase      ):







   @slow
   def 					lowerCamelCase						(      self		)      ->						Union[str, Any]:

       '''simple docstring'''




       snake_case       : Union[str, Any]    =       TFDistilBertModel.from_pretrained("distilbert-base-uncased"		)
       snake_case       : Tuple    =       tf.constant([[0, 1, 2, 3, 4, 5]]		)
       snake_case       : int    =       model(UpperCamelCase__		)[0]

       snake_case       : Tuple    =       [1, 6, 768]
       self.assertEqual(output.shape		,      UpperCamelCase__		)

       snake_case       : List[Any]    =       tf.constant(
           [
               [
                   [0.19261885, -0.13732955, 0.4119799],
                   [0.22150156, -0.07422661, 0.39037204],
                   [0.22756018, -0.0896414, 0.3701467],
               ]
           ]		)
       tf.debugging.assert_near(output[:, :3, :3]		,      UpperCamelCase__		,      atol=1e-4		)









"""simple docstring"""



from collections import OrderedDict
from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import TensorType, logging


if TYPE_CHECKING:
  from ...onnx.config import PatchingSpec
  from ...tokenization_utils_base import PreTrainedTokenizerBase


__snake_case       =       logging.get_logger(__name__)

__snake_case       =       {
    """allenai/longformer-base-4096""": """https://huggingface.co/allenai/longformer-base-4096/resolve/main/config.json""",
    """allenai/longformer-large-4096""": """https://huggingface.co/allenai/longformer-large-4096/resolve/main/config.json""",
    """allenai/longformer-large-4096-finetuned-triviaqa""": (
        """https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/config.json"""
    ),
    """allenai/longformer-base-4096-extra.pos.embd.only""": (
        """https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/config.json"""
    ),
    """allenai/longformer-large-4096-extra.pos.embd.only""": (
        """https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/config.json"""
    ),
}







class _lowerCAmelCase							( snake_case_      ):
   __UpperCAmelCase					:			Tuple				    =    '''longformer'''







   def __init__(      self		,      UpperCamelCase__ = 512		,      UpperCamelCase__ = 2		,      UpperCamelCase__ = 1		,      UpperCamelCase__ = 0		,      UpperCamelCase__ = 2		,      UpperCamelCase__ = 3_0522		,      UpperCamelCase__ = 768		,      UpperCamelCase__ = 12		,      UpperCamelCase__ = 12		,      UpperCamelCase__ = 3072		,      UpperCamelCase__ = "gelu"		,      UpperCamelCase__ = 0.1		,      UpperCamelCase__ = 0.1		,      UpperCamelCase__ = 512		,      UpperCamelCase__ = 2		,      UpperCamelCase__ = 0.02		,      UpperCamelCase__ = 1e-12		,      UpperCamelCase__ = False		,      **UpperCamelCase__		,      )      ->						str:

       '''simple docstring'''




       super().__init__(pad_token_id=UpperCamelCase__		,      **UpperCamelCase__		)

       snake_case       : int    =       attention_window
       snake_case       : Any    =       sep_token_id
       snake_case       : Dict    =       bos_token_id
       snake_case       : int    =       eos_token_id
       snake_case       : List[str]    =       vocab_size
       snake_case       : Dict    =       hidden_size
       snake_case       : Optional[int]    =       num_hidden_layers
       snake_case       : List[str]    =       num_attention_heads
       snake_case       : Dict    =       hidden_act
       snake_case       : Union[str, Any]    =       intermediate_size
       snake_case       : Optional[int]    =       hidden_dropout_prob
       snake_case       : List[Any]    =       attention_probs_dropout_prob
       snake_case       : Any    =       max_position_embeddings
       snake_case       : Union[str, Any]    =       type_vocab_size
       snake_case       : Union[str, Any]    =       initializer_range
       snake_case       : int    =       layer_norm_eps
       snake_case       : str    =       onnx_export







class _lowerCAmelCase							( snake_case_      ):







   def __init__(      self		,      UpperCamelCase__		,      UpperCamelCase__ = "default"		,      UpperCamelCase__ = None		)      ->						List[str]:

       '''simple docstring'''




       super().__init__(UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)
       snake_case       : int    =       True







   @property
   def 					lowerCamelCase						(      self		)      ->						Mapping[str, Mapping[int, str]]:

       '''simple docstring'''




       if self.task == "multiple-choice":
           snake_case       : str    =       {0: "batch", 1: "choice", 2: "sequence"}
       else:
           snake_case       : List[str]    =       {0: "batch", 1: "sequence"}
       return OrderedDict(
           [
               ("input_ids", dynamic_axis),
               ("attention_mask", dynamic_axis),
               ("global_attention_mask", dynamic_axis),
           ]		)







   @property
   def 					lowerCamelCase						(      self		)      ->						Mapping[str, Mapping[int, str]]:

       '''simple docstring'''




       snake_case       : str    =       super().outputs
       if self.task == "default":
           snake_case       : str    =       {0: "batch"}
       return outputs







   @property
   def 					lowerCamelCase						(      self		)      ->						float:

       '''simple docstring'''




       return 1e-4







   @property
   def 					lowerCamelCase						(      self		)      ->						int:

       '''simple docstring'''




       return max(super().default_onnx_opset		,      14		)







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__ = -1		,      UpperCamelCase__ = -1		,      UpperCamelCase__ = False		,      UpperCamelCase__ = None		,      )      ->						Mapping[str, Any]:

       '''simple docstring'''




       snake_case       : Optional[int]    =       super().generate_dummy_inputs(
           preprocessor=UpperCamelCase__		,      batch_size=UpperCamelCase__		,      seq_length=UpperCamelCase__		,      is_pair=UpperCamelCase__		,      framework=UpperCamelCase__		)
       import torch

       # for some reason, replacing this code by inputs["global_attention_mask"] = torch.randint(2, inputs["input_ids"].shape, dtype=torch.int64)
       # makes the export fail randomly
       snake_case       : Any    =       torch.zeros_like(inputs["input_ids"]		)
       # make every second token global
       snake_case       : Any    =       1

       return inputs





