



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







						"""simple docstring"""







						def __init__(self   ,		lowercase__   ,		):
											snake_case_					:		Tuple 					=					parent
											snake_case_					:		int 					=					13
											snake_case_					:		Dict 					=					7
											snake_case_					:		int 					=					True
											snake_case_					:		Union[str, Any] 					=					True
											snake_case_					:		int 					=					False
											snake_case_					:		Any 					=					True
											snake_case_					:		Any 					=					99
											snake_case_					:		Any 					=					32
											snake_case_					:		Any 					=					2
											snake_case_					:		Any 					=					4
											snake_case_					:		str 					=					37
											snake_case_					:		Optional[Any] 					=					"""gelu"""
											snake_case_					:		Any 					=					0.1
											snake_case_					:		Union[str, Any] 					=					0.1
											snake_case_					:		Optional[int] 					=					5_12
											snake_case_					:		Any 					=					16
											snake_case_					:		Optional[int] 					=					2
											snake_case_					:		Optional[int] 					=					0.02
											snake_case_					:		List[str] 					=					3
											snake_case_					:		List[str] 					=					4
											snake_case_					:		List[str] 					=					None







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Any 					=					ids_tensor([self.batch_size, self.seq_length]   ,		self.vocab_size       )

											snake_case_					:		Any 					=					None
											if self.use_input_mask:
																snake_case_					:		Tuple 					=					random_attention_mask([self.batch_size, self.seq_length]       )

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

											snake_case_					:		int 					=					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 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		List[str] 					=					TFDistilBertModel(config=lowercase__       )
											snake_case_					:		Any 					=					{"""input_ids""": input_ids, """attention_mask""": input_mask}

											snake_case_					:		List[Any] 					=					model(lowercase__       )

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

											snake_case_					:		Dict 					=					model(lowercase__       )

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







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







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







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







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







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Optional[Any] 					=					self.prepare_config_and_inputs()
											(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    __lowercase     (					_UpperCAmelCase					,    _UpperCAmelCase					,    unittest.TestCase):







						"""simple docstring"""







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







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Any 					=					TFDistilBertModelTester(self       )
											snake_case_					:		Tuple 					=					ConfigTester(self   ,		config_class=lowercase__   ,		dim=37       )







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







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		str 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_model(*lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Union[str, Any] 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_for_masked_lm(*lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Dict 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_for_question_answering(*lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Optional[int] 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Optional[int] 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Tuple 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_for_token_classification(*lowercase__       )







						@slow
						def 						__UpperCamelCase							(self       ):
											for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]       ):
																snake_case_					:		List[Any] 					=					TFDistilBertModel.from_pretrained(lowercase__       )
																self.assertIsNotNone(lowercase__       )







@require_tf
class    __lowercase     (					unittest.TestCase):







						"""simple docstring"""







						@slow
						def 						__UpperCamelCase							(self       ):
											snake_case_					:		str 					=					TFDistilBertModel.from_pretrained("""distilbert-base-uncased"""       )
											snake_case_					:		Optional[int] 					=					tf.constant([[0, 1, 2, 3, 4, 5]]       )
											snake_case_					:		Union[str, Any] 					=					model(lowercase__       )[0]

											snake_case_					:		int 					=					[1, 6, 7_68]
											self.assertEqual(output.shape   ,		lowercase__       )

											snake_case_					:		Optional[int] 					=					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]   ,		lowercase__   ,		atol=1e-4       )









"""simple docstring"""

import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.text import TextDatasetReader

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  List[Any]				,    SCREAMING_SNAKE_CASE__ :  Dict			):



					"""simple docstring"""





					assert isinstance(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__			)
					assert dataset.num_rows == 4
					assert dataset.num_columns == 1
					assert dataset.column_names == ["text"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize("""keep_in_memory"""				,    [False, True]			)
def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Union[str, Any]				,    SCREAMING_SNAKE_CASE__ :  Optional[Any]				,    SCREAMING_SNAKE_CASE__ :  Optional[Any]			):



					"""simple docstring"""





					snake_case_					:		Optional[Any] 					=					tmp_path / """cache"""
					snake_case_					:		Optional[int] 					=					{"""text""": """string"""}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										snake_case_					:		Tuple 					=					TextDatasetReader(SCREAMING_SNAKE_CASE__				,    cache_dir=SCREAMING_SNAKE_CASE__				,    keep_in_memory=SCREAMING_SNAKE_CASE__			).read()
					_check_text_dataset(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__			)







@pytest.mark.parametrize(
    """features"""				,    [
        None,
        {"""text""": """string"""},
        {"""text""": """int32"""},
        {"""text""": """float32"""},
    ]				,    )
def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Dict				,    SCREAMING_SNAKE_CASE__ :  List[str]				,    SCREAMING_SNAKE_CASE__ :  Union[str, Any]			):



					"""simple docstring"""





					snake_case_					:		List[Any] 					=					tmp_path / """cache"""
					snake_case_					:		int 					=					{"""text""": """string"""}
					snake_case_					:		Any 					=					features.copy() if features else default_expected_features
					snake_case_					:		List[Any] 					=					(
					    Features({feature: Value(SCREAMING_SNAKE_CASE__			) for feature, dtype in features.items()}			) if features is not None else None
					)
					snake_case_					:		Dict 					=					TextDatasetReader(SCREAMING_SNAKE_CASE__				,    features=SCREAMING_SNAKE_CASE__				,    cache_dir=SCREAMING_SNAKE_CASE__			).read()
					_check_text_dataset(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__			)







@pytest.mark.parametrize("""split"""				,    [None, NamedSplit("""train"""			), """train""", """test"""]			)
def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Tuple				,    SCREAMING_SNAKE_CASE__ :  int				,    SCREAMING_SNAKE_CASE__ :  str			):



					"""simple docstring"""





					snake_case_					:		Union[str, Any] 					=					tmp_path / """cache"""
					snake_case_					:		Optional[Any] 					=					{"""text""": """string"""}
					snake_case_					:		Optional[int] 					=					TextDatasetReader(SCREAMING_SNAKE_CASE__				,    cache_dir=SCREAMING_SNAKE_CASE__				,    split=SCREAMING_SNAKE_CASE__			).read()
					_check_text_dataset(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__			)
					assert dataset.split == split if split else "train"







@pytest.mark.parametrize("""path_type"""				,    [str, list]			)
def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Tuple				,    SCREAMING_SNAKE_CASE__ :  List[str]				,    SCREAMING_SNAKE_CASE__ :  Dict			):



					"""simple docstring"""





					if issubclass(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__			):
										snake_case_					:		List[str] 					=					text_path
					elif issubclass(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__			):
										snake_case_					:		str 					=					[text_path]
					snake_case_					:		List[str] 					=					tmp_path / """cache"""
					snake_case_					:		List[str] 					=					{"""text""": """string"""}
					snake_case_					:		Dict 					=					TextDatasetReader(SCREAMING_SNAKE_CASE__				,    cache_dir=SCREAMING_SNAKE_CASE__			).read()
					_check_text_dataset(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__			)







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  str				,    SCREAMING_SNAKE_CASE__ :  Dict				,    SCREAMING_SNAKE_CASE__ :  List[str]=("train",)			):



					"""simple docstring"""





					assert isinstance(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__			)
					for split in splits:
										snake_case_					:		Dict 					=					dataset_dict[split]
										assert dataset.num_rows == 4
										assert dataset.num_columns == 1
										assert dataset.column_names == ["text"]
										for feature, expected_dtype in expected_features.items():
															assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize("""keep_in_memory"""				,    [False, True]			)
def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  List[Any]				,    SCREAMING_SNAKE_CASE__ :  Tuple				,    SCREAMING_SNAKE_CASE__ :  Any			):



					"""simple docstring"""





					snake_case_					:		int 					=					tmp_path / """cache"""
					snake_case_					:		List[str] 					=					{"""text""": """string"""}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										snake_case_					:		Optional[Any] 					=					TextDatasetReader({"""train""": text_path}				,    cache_dir=SCREAMING_SNAKE_CASE__				,    keep_in_memory=SCREAMING_SNAKE_CASE__			).read()
					_check_text_datasetdict(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__			)







@pytest.mark.parametrize(
    """features"""				,    [
        None,
        {"""text""": """string"""},
        {"""text""": """int32"""},
        {"""text""": """float32"""},
    ]				,    )
def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  int				,    SCREAMING_SNAKE_CASE__ :  int				,    SCREAMING_SNAKE_CASE__ :  List[str]			):



					"""simple docstring"""





					snake_case_					:		Tuple 					=					tmp_path / """cache"""
					# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
					snake_case_					:		List[str] 					=					{"""text""": """string"""}
					snake_case_					:		int 					=					features.copy() if features else default_expected_features
					snake_case_					:		Tuple 					=					(
					    Features({feature: Value(SCREAMING_SNAKE_CASE__			) for feature, dtype in features.items()}			) if features is not None else None
					)
					snake_case_					:		str 					=					TextDatasetReader({"""train""": text_path}				,    features=SCREAMING_SNAKE_CASE__				,    cache_dir=SCREAMING_SNAKE_CASE__			).read()
					_check_text_datasetdict(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__			)







@pytest.mark.parametrize("""split"""				,    [None, NamedSplit("""train"""			), """train""", """test"""]			)
def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Dict				,    SCREAMING_SNAKE_CASE__ :  Dict				,    SCREAMING_SNAKE_CASE__ :  Any			):



					"""simple docstring"""





					if split:
										snake_case_					:		Union[str, Any] 					=					{split: text_path}
					else:
										snake_case_					:		Union[str, Any] 					=					"""train"""
										snake_case_					:		int 					=					{"""train""": text_path, """test""": text_path}
					snake_case_					:		List[Any] 					=					tmp_path / """cache"""
					snake_case_					:		Tuple 					=					{"""text""": """string"""}
					snake_case_					:		int 					=					TextDatasetReader(SCREAMING_SNAKE_CASE__				,    cache_dir=SCREAMING_SNAKE_CASE__			).read()
					_check_text_datasetdict(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__				,    splits=list(path.keys()			)			)
					assert all(dataset[split].split == split for split in path.keys()			)




