



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





    def __init__(	self       ,     _lowerCAmelCase       ,     ):
          _lowerCAmelCase 			=  parent
          _lowerCAmelCase 			=  13
          _lowerCAmelCase 			=  7
          _lowerCAmelCase 			=  True
          _lowerCAmelCase 			=  True
          _lowerCAmelCase 			=  False
          _lowerCAmelCase 			=  True
          _lowerCAmelCase 			=  99
          _lowerCAmelCase 			=  32
          _lowerCAmelCase 			=  2
          _lowerCAmelCase 			=  4
          _lowerCAmelCase 			=  37
          _lowerCAmelCase 			=  """gelu"""
          _lowerCAmelCase 			=  0.1
          _lowerCAmelCase 			=  0.1
          _lowerCAmelCase 			=  512
          _lowerCAmelCase 			=  16
          _lowerCAmelCase 			=  2
          _lowerCAmelCase 			=  0.02
          _lowerCAmelCase 			=  3
          _lowerCAmelCase 			=  4
          _lowerCAmelCase 			=  None





    def 	__lowerCAmelCase							(	self			):
          _lowerCAmelCase 			=  ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size			)

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

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

          _lowerCAmelCase 			=  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       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase			):
          _lowerCAmelCase 			=  TFDistilBertModel(config=_lowercase			)
          _lowerCAmelCase 			=  {"""input_ids""": input_ids, """attention_mask""": input_mask}

          _lowerCAmelCase 			=  model(_lowercase			)

          _lowerCAmelCase 			=  [input_ids, input_mask]

          _lowerCAmelCase 			=  model(_lowercase			)

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





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





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





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





    def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase			):
          _lowerCAmelCase 			=  self.num_labels
          _lowerCAmelCase 			=  TFDistilBertForTokenClassification(_lowercase			)
          _lowerCAmelCase 			=  {"""input_ids""": input_ids, """attention_mask""": input_mask}
          _lowerCAmelCase 			=  model(_lowercase			)
          self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.num_labels)			)





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






@require_tf
class 		UpperCAmelCase						(						__snake_case       ,__snake_case       ,unittest.TestCase				):
    SCREAMING_SNAKE_CASE__							=     (
        (
            TFDistilBertModel,
            TFDistilBertForMaskedLM,
            TFDistilBertForQuestionAnswering,
            TFDistilBertForSequenceClassification,
            TFDistilBertForTokenClassification,
            TFDistilBertForMultipleChoice,
        )
        if is_tf_available()
        else None
    )
    SCREAMING_SNAKE_CASE__							=     (
        {
            '''feature-extraction''': TFDistilBertModel,
            '''fill-mask''': TFDistilBertForMaskedLM,
            '''question-answering''': TFDistilBertForQuestionAnswering,
            '''text-classification''': TFDistilBertForSequenceClassification,
            '''token-classification''': TFDistilBertForTokenClassification,
            '''zero-shot''': TFDistilBertForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    SCREAMING_SNAKE_CASE__							=     False
    SCREAMING_SNAKE_CASE__							=     False





    def 	__lowerCAmelCase							(	self			):
          _lowerCAmelCase 			=  TFDistilBertModelTester(self			)
          _lowerCAmelCase 			=  ConfigTester(self       ,     config_class=_lowercase       ,     dim=37			)





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





    def 	__lowerCAmelCase							(	self			):
          _lowerCAmelCase 			=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_model(*_lowercase			)





    def 	__lowerCAmelCase							(	self			):
          _lowerCAmelCase 			=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_masked_lm(*_lowercase			)





    def 	__lowerCAmelCase							(	self			):
          _lowerCAmelCase 			=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_question_answering(*_lowercase			)





    def 	__lowerCAmelCase							(	self			):
          _lowerCAmelCase 			=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_sequence_classification(*_lowercase			)





    def 	__lowerCAmelCase							(	self			):
          _lowerCAmelCase 			=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_multiple_choice(*_lowercase			)





    def 	__lowerCAmelCase							(	self			):
          _lowerCAmelCase 			=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_token_classification(*_lowercase			)





    @slow
    def 	__lowerCAmelCase							(	self			):

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






@require_tf
class 		UpperCAmelCase						(						unittest.TestCase				):





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

          _lowerCAmelCase 			=  [1, 6, 768]
          self.assertEqual(output.shape       ,     _lowercase			)

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




from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch





@require_torch
@require_tf
class 		UpperCAmelCase						(						snake_case_				):





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  SMALL_MODEL_IDENTIFIER
         _lowerCAmelCase 			=  '''pt'''
         _lowerCAmelCase 			=  '''tf'''





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase			):
         _lowerCAmelCase 			=  AutoModel.from_pretrained(self.test_model			)
         model_pt.save_pretrained(_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase			):
         _lowerCAmelCase 			=  TFAutoModel.from_pretrained(self.test_model       ,     from_pt=_lowerCAmelCase			)
         model_tf.save_pretrained(_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  '''mock_framework'''

         # Framework provided - return whatever the user provides
         _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model       ,     _lowerCAmelCase			)
         self.assertEqual(_lowerCAmelCase       ,     _lowerCAmelCase			)

         # Local checkpoint and framework provided - return provided framework
         # PyTorch checkpoint
         with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase       ,     _lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     _lowerCAmelCase			)

         # TensorFlow checkpoint
         with TemporaryDirectory() as local_tf_ckpt:
               self._setup_tf_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase       ,     _lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     _lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         # PyTorch checkpoint
         with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_pt			)

         # TensorFlow checkpoint
         with TemporaryDirectory() as local_tf_ckpt:
               self._setup_tf_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_tf			)

         # Invalid local checkpoint
         with TemporaryDirectory() as local_invalid_ckpt:
               with self.assertRaises(_lowerCAmelCase			):
                     _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_tf_available'''       ,     _lowerCAmelCase			):
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_pt			)

         # PyTorch not in environment -> use TensorFlow
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_torch_available'''       ,     _lowerCAmelCase			):
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_tf			)

         # Both in environment -> use PyTorch
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_tf_available'''       ,     _lowerCAmelCase			), patch(
             '''transformers.onnx.features.is_torch_available'''       ,     _lowerCAmelCase			):
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_pt			)

         # Both not in environment -> raise error
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_tf_available'''       ,     _lowerCAmelCase			), patch(
             '''transformers.onnx.features.is_torch_available'''       ,     _lowerCAmelCase			):
               with self.assertRaises(_lowerCAmelCase			):
                     _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)