





'''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 			SCREAMING_SNAKE_CASE   :


     def __init__(						self						,     _UpperCAmelCase						,     ):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Optional[Any]								=	ids_tensor([self.batch_size, self.seq_length]						,     self.vocab_size)

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

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

          __A						:			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 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			int								=	TFDistilBertModel(config=_UpperCAmelCase)
          __A						:			int								=	{'input_ids': input_ids, 'attention_mask': input_mask}

          __A						:			Optional[Any]								=	model(_UpperCAmelCase)

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

          __A						:			Tuple								=	model(_UpperCAmelCase)

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


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			List[Any]								=	TFDistilBertForMaskedLM(config=_UpperCAmelCase)
          __A						:			List[Any]								=	{'input_ids': input_ids, 'attention_mask': input_mask}
          __A						:			List[Any]								=	model(_UpperCAmelCase)
          self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.seq_length, self.vocab_size))


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			Any								=	TFDistilBertForQuestionAnswering(config=_UpperCAmelCase)
          __A						:			List[str]								=	{
              'input_ids': input_ids,
              'attention_mask': input_mask,
          }
          __A						:			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 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			Tuple								=	self.num_labels
          __A						:			Union[str, Any]								=	TFDistilBertForSequenceClassification(_UpperCAmelCase)
          __A						:			Tuple								=	{'input_ids': input_ids, 'attention_mask': input_mask}
          __A						:			Any								=	model(_UpperCAmelCase)
          self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.num_labels))


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Dict								=	self.prepare_config_and_inputs()
          ((__A)			,(__A)			,(__A)			,(__A)			,(__A)			,(__A))						:			str								=	config_and_inputs
          __A						:			Tuple								=	{'input_ids': input_ids, 'attention_mask': input_mask}
          return config, inputs_dict




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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Optional[Any]								=	TFDistilBertModelTester(self)
          __A						:			str								=	ConfigTester(self						,     config_class=_UpperCAmelCase						,     dim=37)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          self.config_tester.run_common_tests()


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Union[str, Any]								=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_model(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			List[Any]								=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_masked_lm(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			str								=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_question_answering(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Tuple								=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_sequence_classification(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Union[str, Any]								=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_multiple_choice(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			List[Any]								=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_token_classification(*_UpperCAmelCase)


     @slow
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]):
               __A						:			List[str]								=	TFDistilBertModel.from_pretrained(_UpperCAmelCase)
               self.assertIsNotNone(_UpperCAmelCase)




@require_tf
class 			SCREAMING_SNAKE_CASE   (unittest.TestCase							):


     @slow
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Optional[int]								=	TFDistilBertModel.from_pretrained('distilbert-base-uncased')
          __A						:			int								=	tf.constant([[0, 1, 2, 3, 4, 5]])
          __A						:			Any								=	model(_UpperCAmelCase)[0]

          __A						:			List[Any]								=	[1, 6, 768]
          self.assertEqual(output.shape						,     _UpperCAmelCase)

          __A						:			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 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 			SCREAMING_SNAKE_CASE   (a__							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Optional[int]								=	SMALL_MODEL_IDENTIFIER
          __A						:			Any								=	'pt'
          __A						:			str								=	'tf'


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			Any								=	AutoModel.from_pretrained(self.test_model)
          model_pt.save_pretrained(_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			Union[str, Any]								=	TFAutoModel.from_pretrained(self.test_model						,     from_pt=_UpperCAmelCase)
          model_tf.save_pretrained(_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Union[str, Any]								=	'mock_framework'

          # Framework provided - return whatever the user provides
          __A						:			Union[str, Any]								=	FeaturesManager.determine_framework(self.test_model						,     _UpperCAmelCase)
          self.assertEqual(_UpperCAmelCase						,     _UpperCAmelCase)

          # Local checkpoint and framework provided - return provided framework
          # PyTorch checkpoint
          with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_UpperCAmelCase)
               __A						:			List[str]								=	FeaturesManager.determine_framework(_UpperCAmelCase						,     _UpperCAmelCase)
               self.assertEqual(_UpperCAmelCase						,     _UpperCAmelCase)

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_UpperCAmelCase)
               __A						:			Optional[int]								=	FeaturesManager.determine_framework(_UpperCAmelCase)
               self.assertEqual(_UpperCAmelCase						,     self.framework_pt)

          # TensorFlow checkpoint
          with TemporaryDirectory() as local_tf_ckpt:
               self._setup_tf_ckpt(_UpperCAmelCase)
               __A						:			List[str]								=	FeaturesManager.determine_framework(_UpperCAmelCase)
               self.assertEqual(_UpperCAmelCase						,     self.framework_tf)

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Tuple								=	MagicMock(return_value=_UpperCAmelCase)
          with patch('transformers.onnx.features.is_tf_available'						,     _UpperCAmelCase):
               __A						:			Union[str, Any]								=	FeaturesManager.determine_framework(self.test_model)
               self.assertEqual(_UpperCAmelCase						,     self.framework_pt)

          # PyTorch not in environment -> use TensorFlow
          __A						:			List[str]								=	MagicMock(return_value=_UpperCAmelCase)
          with patch('transformers.onnx.features.is_torch_available'						,     _UpperCAmelCase):
               __A						:			List[Any]								=	FeaturesManager.determine_framework(self.test_model)
               self.assertEqual(_UpperCAmelCase						,     self.framework_tf)

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

          # Both not in environment -> raise error
          __A						:			List[str]								=	MagicMock(return_value=_UpperCAmelCase)
          __A						:			Tuple								=	MagicMock(return_value=_UpperCAmelCase)
          with patch('transformers.onnx.features.is_tf_available'						,     _UpperCAmelCase), patch(
              'transformers.onnx.features.is_torch_available'						,     _UpperCAmelCase):
               with self.assertRaises(_UpperCAmelCase):
                    __A						:			int								=	FeaturesManager.determine_framework(self.test_model)