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

import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / """utils"""))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


_UpperCamelCase		     =    get_tests_dir("""fixtures""")
_UpperCamelCase		     =    get_tests_dir("""fixtures/dummy_feature_extractor_config.json""")
_UpperCamelCase		     =    get_tests_dir("""fixtures/dummy-config.json""")






class 			__a					(						unittest.TestCase       ):





       """simple docstring"""


       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				0






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base-960h"   )
        self.assertIsInstance(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Tuple 			=				AutoFeatureExtractor.from_pretrained(snake_case   )
        self.assertIsInstance(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        with tempfile.TemporaryDirectory() as tmpdirname:
         lowerCAmelCase__   :				int 			=				WavaVecaConfig()

         # remove feature_extractor_type to make sure config.json alone is enough to load feature processor locally
         lowerCAmelCase__   :				Union[str, Any] 			=				AutoFeatureExtractor.from_pretrained(snake_case   ).to_dict()

         config_dict.pop("feature_extractor_type"   )
         lowerCAmelCase__   :				Any 			=				WavaVecaFeatureExtractor(**snake_case   )

         # save in new folder
         model_config.save_pretrained(snake_case   )
         config.save_pretrained(snake_case   )

         lowerCAmelCase__   :				List[str] 			=				AutoFeatureExtractor.from_pretrained(snake_case   )

         # make sure private variable is not incorrectly saved
         lowerCAmelCase__   :				Optional[int] 			=				json.loads(config.to_json_string()   )
         self.assertTrue("_processor_class" not in dict_as_saved   )

        self.assertIsInstance(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[str] 			=				AutoFeatureExtractor.from_pretrained(snake_case   )
        self.assertIsInstance(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        with self.assertRaisesRegex(
            snake_case			,					"bert-base is not a local folder and is not a valid model identifier"   ):
         lowerCAmelCase__   :				Optional[int] 			=				AutoFeatureExtractor.from_pretrained("bert-base"   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        with self.assertRaisesRegex(
            snake_case			,					R"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"   ):
         lowerCAmelCase__   :				Optional[Any] 			=				AutoFeatureExtractor.from_pretrained(snake_case			,					revision="aaaaaa"   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        with self.assertRaisesRegex(
            snake_case			,					"hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json."			,					):
         lowerCAmelCase__   :				Dict 			=				AutoFeatureExtractor.from_pretrained("hf-internal-testing/config-no-model"   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        with self.assertRaises(snake_case   ):
         lowerCAmelCase__   :				Tuple 			=				AutoFeatureExtractor.from_pretrained(
             "hf-internal-testing/test_dynamic_feature_extractor"   )
        # If remote code is disabled, we can't load this config.
        with self.assertRaises(snake_case   ):
         lowerCAmelCase__   :				str 			=				AutoFeatureExtractor.from_pretrained(
             "hf-internal-testing/test_dynamic_feature_extractor"			,					trust_remote_code=snake_case   )

        lowerCAmelCase__   :				Optional[int] 			=				AutoFeatureExtractor.from_pretrained(
            "hf-internal-testing/test_dynamic_feature_extractor"			,					trust_remote_code=snake_case   )
        self.assertEqual(feature_extractor.__class__.__name__			,					"NewFeatureExtractor"   )

        # Test feature extractor can be reloaded.
        with tempfile.TemporaryDirectory() as tmp_dir:
         feature_extractor.save_pretrained(snake_case   )
         lowerCAmelCase__   :				Optional[Any] 			=				AutoFeatureExtractor.from_pretrained(snake_case			,					trust_remote_code=snake_case   )
        self.assertEqual(reloaded_feature_extractor.__class__.__name__			,					"NewFeatureExtractor"   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        try:
         AutoConfig.register("custom"			,					snake_case   )
         AutoFeatureExtractor.register(snake_case			,					snake_case   )
         # Trying to register something existing in the Transformers library will raise an error
         with self.assertRaises(snake_case   ):
          AutoFeatureExtractor.register(snake_case			,					snake_case   )

         # Now that the config is registered, it can be used as any other config with the auto-API
         lowerCAmelCase__   :				List[Any] 			=				CustomFeatureExtractor.from_pretrained(snake_case   )
         with tempfile.TemporaryDirectory() as tmp_dir:
          feature_extractor.save_pretrained(snake_case   )
          lowerCAmelCase__   :				Any 			=				AutoFeatureExtractor.from_pretrained(snake_case   )
          self.assertIsInstance(snake_case			,					snake_case   )

        finally:
         if "custom" in CONFIG_MAPPING._extra_content:
          del CONFIG_MAPPING._extra_content["custom"]
         if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
          del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        class 			__a					(						__magic_name__       ):





               """simple docstring"""


               __UpperCamelCase					:      Optional[int]    =						True

        try:
         AutoConfig.register("custom"			,					snake_case   )
         AutoFeatureExtractor.register(snake_case			,					snake_case   )
         # If remote code is not set, the default is to use local
         lowerCAmelCase__   :				List[Any] 			=				AutoFeatureExtractor.from_pretrained(
             "hf-internal-testing/test_dynamic_feature_extractor"   )
         self.assertEqual(feature_extractor.__class__.__name__			,					"NewFeatureExtractor"   )
         self.assertTrue(feature_extractor.is_local   )

         # If remote code is disabled, we load the local one.
         lowerCAmelCase__   :				Dict 			=				AutoFeatureExtractor.from_pretrained(
             "hf-internal-testing/test_dynamic_feature_extractor"			,					trust_remote_code=snake_case   )
         self.assertEqual(feature_extractor.__class__.__name__			,					"NewFeatureExtractor"   )
         self.assertTrue(feature_extractor.is_local   )

         # If remote is enabled, we load from the Hub
         lowerCAmelCase__   :				Tuple 			=				AutoFeatureExtractor.from_pretrained(
             "hf-internal-testing/test_dynamic_feature_extractor"			,					trust_remote_code=snake_case   )
         self.assertEqual(feature_extractor.__class__.__name__			,					"NewFeatureExtractor"   )
         self.assertTrue(not hasattr(snake_case			,					"is_local"   )   )

        finally:
         if "custom" in CONFIG_MAPPING._extra_content:
          del CONFIG_MAPPING._extra_content["custom"]
         if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
          del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]






