
import os
import tempfile
import unittest

from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    DistilBertForMaskedLM,
		    DistilBertForMultipleChoice,
		    DistilBertForQuestionAnswering,
		    DistilBertForSequenceClassification,
		    DistilBertForTokenClassification,
		    DistilBertModel,
		)





class       __A(				UpperCAmelCase_       ):

					def __init__(	self					:     Any   ,					__UpperCamelCase					:     List[str]   ,					__UpperCamelCase					:     List[str]=1_3   ,					__UpperCamelCase					:     int=7   ,					__UpperCamelCase					:     int=True   ,					__UpperCamelCase					:     List[str]=True   ,					__UpperCamelCase					:     str=False   ,					__UpperCamelCase					:     str=True   ,					__UpperCamelCase					:     Tuple=9_9   ,					__UpperCamelCase					:     List[str]=3_2   ,					__UpperCamelCase					:     Optional[int]=5   ,					__UpperCamelCase					:     Union[str, Any]=4   ,					__UpperCamelCase					:     int=3_7   ,					__UpperCamelCase					:     Union[str, Any]="gelu"   ,					__UpperCamelCase					:     Tuple=0.1   ,					__UpperCamelCase					:     Optional[int]=0.1   ,					__UpperCamelCase					:     Union[str, Any]=5_1_2   ,					__UpperCamelCase					:     Optional[Any]=1_6   ,					__UpperCamelCase					:     List[str]=2   ,					__UpperCamelCase					:     int=0.02   ,					__UpperCamelCase					:     str=3   ,					__UpperCamelCase					:     Optional[int]=4   ,					__UpperCamelCase					:     Optional[int]=None   ,					):
												lowerCamelCase_							  =					parent
												lowerCamelCase_							  =					batch_size
												lowerCamelCase_							  =					seq_length
												lowerCamelCase_							  =					is_training
												lowerCamelCase_							  =					use_input_mask
												lowerCamelCase_							  =					use_token_type_ids
												lowerCamelCase_							  =					use_labels
												lowerCamelCase_							  =					vocab_size
												lowerCamelCase_							  =					hidden_size
												lowerCamelCase_							  =					num_hidden_layers
												lowerCamelCase_							  =					num_attention_heads
												lowerCamelCase_							  =					intermediate_size
												lowerCamelCase_							  =					hidden_act
												lowerCamelCase_							  =					hidden_dropout_prob
												lowerCamelCase_							  =					attention_probs_dropout_prob
												lowerCamelCase_							  =					max_position_embeddings
												lowerCamelCase_							  =					type_vocab_size
												lowerCamelCase_							  =					type_sequence_label_size
												lowerCamelCase_							  =					initializer_range
												lowerCamelCase_							  =					num_labels
												lowerCamelCase_							  =					num_choices
												lowerCamelCase_							  =					scope

					def    lowercase__      (	self					:     Dict   ):
												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_							  =					self.get_config()

												return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels

					def    lowercase__      (	self					:     Any   ):
												return 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   ,					)

					def    lowercase__      (	self					:     Dict   ,					__UpperCamelCase					:     Optional[int]   ,					__UpperCamelCase					:     Dict   ,					__UpperCamelCase					:     Optional[int]   ,					__UpperCamelCase					:     Union[str, Any]   ,					__UpperCamelCase					:     Union[str, Any]   ,					__UpperCamelCase					:     Optional[int]   ):
												lowerCamelCase_							  =					DistilBertModel(config=_lowercase   )
												model.to(_lowercase   )
												model.eval()
												lowerCamelCase_							  =					model(_lowercase   ,					_lowercase   )
												lowerCamelCase_							  =					model(_lowercase   )
												self.parent.assertEqual(result.last_hidden_state.shape   ,					(self.batch_size, self.seq_length, self.hidden_size)   )

					def    lowercase__      (	self					:     Tuple   ,					__UpperCamelCase					:     int   ,					__UpperCamelCase					:     Tuple   ,					__UpperCamelCase					:     Optional[Any]   ,					__UpperCamelCase					:     str   ,					__UpperCamelCase					:     List[str]   ,					__UpperCamelCase					:     Optional[int]   ):
												lowerCamelCase_							  =					DistilBertForMaskedLM(config=_lowercase   )
												model.to(_lowercase   )
												model.eval()
												lowerCamelCase_							  =					model(_lowercase   ,					attention_mask=_lowercase   ,					labels=_lowercase   )
												self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.seq_length, self.vocab_size)   )

					def    lowercase__      (	self					:     Any   ,					__UpperCamelCase					:     str   ,					__UpperCamelCase					:     Dict   ,					__UpperCamelCase					:     Optional[Any]   ,					__UpperCamelCase					:     Dict   ,					__UpperCamelCase					:     List[str]   ,					__UpperCamelCase					:     Optional[Any]   ):
												lowerCamelCase_							  =					DistilBertForQuestionAnswering(config=_lowercase   )
												model.to(_lowercase   )
												model.eval()
												lowerCamelCase_							  =					model(
												    _lowercase   ,					attention_mask=_lowercase   ,					start_positions=_lowercase   ,					end_positions=_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    lowercase__      (	self					:     Any   ,					__UpperCamelCase					:     List[str]   ,					__UpperCamelCase					:     Optional[int]   ,					__UpperCamelCase					:     List[Any]   ,					__UpperCamelCase					:     Optional[Any]   ,					__UpperCamelCase					:     int   ,					__UpperCamelCase					:     Dict   ):
												lowerCamelCase_							  =					self.num_labels
												lowerCamelCase_							  =					DistilBertForSequenceClassification(_lowercase   )
												model.to(_lowercase   )
												model.eval()
												lowerCamelCase_							  =					model(_lowercase   ,					attention_mask=_lowercase   ,					labels=_lowercase   )
												self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.num_labels)   )

					def    lowercase__      (	self					:     Optional[int]   ,					__UpperCamelCase					:     List[Any]   ,					__UpperCamelCase					:     Dict   ,					__UpperCamelCase					:     Optional[Any]   ,					__UpperCamelCase					:     int   ,					__UpperCamelCase					:     Dict   ,					__UpperCamelCase					:     List[str]   ):
												lowerCamelCase_							  =					self.num_labels
												lowerCamelCase_							  =					DistilBertForTokenClassification(config=_lowercase   )
												model.to(_lowercase   )
												model.eval()

												lowerCamelCase_							  =					model(_lowercase   ,					attention_mask=_lowercase   ,					labels=_lowercase   )
												self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.seq_length, self.num_labels)   )

					def    lowercase__      (	self					:     List[Any]   ,					__UpperCamelCase					:     List[Any]   ,					__UpperCamelCase					:     Union[str, Any]   ,					__UpperCamelCase					:     Dict   ,					__UpperCamelCase					:     Tuple   ,					__UpperCamelCase					:     Any   ,					__UpperCamelCase					:     Union[str, Any]   ):
												lowerCamelCase_							  =					self.num_choices
												lowerCamelCase_							  =					DistilBertForMultipleChoice(config=_lowercase   )
												model.to(_lowercase   )
												model.eval()
												lowerCamelCase_							  =					input_ids.unsqueeze(1   ).expand(-1   ,					self.num_choices   ,					-1   ).contiguous()
												lowerCamelCase_							  =					input_mask.unsqueeze(1   ).expand(-1   ,					self.num_choices   ,					-1   ).contiguous()
												lowerCamelCase_							  =					model(
												    _lowercase   ,					attention_mask=_lowercase   ,					labels=_lowercase   ,					)
												self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.num_choices)   )



					def    lowercase__      (	self					:     int   ):
												lowerCamelCase_							  =					self.prepare_config_and_inputs()
												((lowerCamelCase_)  ,   (lowerCamelCase_)  ,   (lowerCamelCase_)  ,   (lowerCamelCase_)  ,   (lowerCamelCase_)  ,   (lowerCamelCase_))							  =					config_and_inputs
												lowerCamelCase_							  =					{"""input_ids""": input_ids, """attention_mask""": input_mask}
												return config, inputs_dict





@require_torch
class       __A(				UpperCAmelCase_     ,		UpperCAmelCase_     ,		unittest.TestCase       ):
					SCREAMING_SNAKE_CASE			    =    (
					    (
					        DistilBertModel,
					        DistilBertForMaskedLM,
					        DistilBertForMultipleChoice,
					        DistilBertForQuestionAnswering,
					        DistilBertForSequenceClassification,
					        DistilBertForTokenClassification,
					    )
					    if is_torch_available()
					    else None
					)
					SCREAMING_SNAKE_CASE			    =    (
					    {
					        '''feature-extraction''': DistilBertModel,
					        '''fill-mask''': DistilBertForMaskedLM,
					        '''question-answering''': DistilBertForQuestionAnswering,
					        '''text-classification''': DistilBertForSequenceClassification,
					        '''token-classification''': DistilBertForTokenClassification,
					        '''zero-shot''': DistilBertForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					SCREAMING_SNAKE_CASE			    =    True
					SCREAMING_SNAKE_CASE			    =    True
					SCREAMING_SNAKE_CASE			    =    True
					SCREAMING_SNAKE_CASE			    =    True

					def    lowercase__      (	self					:     Dict   ):
												lowerCamelCase_							  =					DistilBertModelTester(self   )
												lowerCamelCase_							  =					ConfigTester(self   ,					config_class=_lowercase   ,					dim=3_7   )

					def    lowercase__      (	self					:     int   ):
												self.config_tester.run_common_tests()

					def    lowercase__      (	self					:     List[Any]   ):
												lowerCamelCase_							  =					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_distilbert_model(*_lowercase   )

					def    lowercase__      (	self					:     str   ):
												lowerCamelCase_							  =					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_distilbert_for_masked_lm(*_lowercase   )

					def    lowercase__      (	self					:     Any   ):
												lowerCamelCase_							  =					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_distilbert_for_question_answering(*_lowercase   )

					def    lowercase__      (	self					:     int   ):
												lowerCamelCase_							  =					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_distilbert_for_sequence_classification(*_lowercase   )

					def    lowercase__      (	self					:     str   ):
												lowerCamelCase_							  =					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_distilbert_for_token_classification(*_lowercase   )

					def    lowercase__      (	self					:     str   ):
												lowerCamelCase_							  =					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_distilbert_for_multiple_choice(*_lowercase   )

					@slow
					def    lowercase__      (	self					:     Tuple   ):
												for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			lowerCamelCase_							  =					DistilBertModel.from_pretrained(_lowercase   )
																			self.assertIsNotNone(_lowercase   )



					@slow
					@require_torch_gpu
					def    lowercase__      (	self					:     Optional[int]   ):
												lowerCamelCase_  ,   lowerCamelCase_							  =					self.model_tester.prepare_config_and_inputs_for_common()
												for model_class in self.all_model_classes:
																			# BertForMultipleChoice behaves incorrectly in JIT environments.
																			if model_class == DistilBertForMultipleChoice:
																										return

																			lowerCamelCase_							  =					True
																			lowerCamelCase_							  =					model_class(config=_lowercase   )

																			lowerCamelCase_							  =					self._prepare_for_class(_lowercase   ,					_lowercase   )
																			lowerCamelCase_							  =					torch.jit.trace(
																			    _lowercase   ,					(inputs_dict["""input_ids"""].to("""cpu"""   ), inputs_dict["""attention_mask"""].to("""cpu"""   ))   )

																			with tempfile.TemporaryDirectory() as tmp:
																										torch.jit.save(_lowercase   ,					os.path.join(_lowercase   ,					"""traced_model.pt"""   )   )
																										lowerCamelCase_							  =					torch.jit.load(os.path.join(_lowercase   ,					"""traced_model.pt"""   )   ,					map_location=_lowercase   )
																										loaded(inputs_dict["""input_ids"""].to(_lowercase   )   ,					inputs_dict["""attention_mask"""].to(_lowercase   )   )





@require_torch
class       __A(				unittest.TestCase       ):

					@slow
					def    lowercase__      (	self					:     int   ):
												lowerCamelCase_							  =					DistilBertModel.from_pretrained("""distilbert-base-uncased"""   )
												lowerCamelCase_							  =					torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]   )
												lowerCamelCase_							  =					torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]   )
												with torch.no_grad():
																			lowerCamelCase_							  =					model(_lowercase   ,					attention_mask=_lowercase   )[0]
												lowerCamelCase_							  =					torch.Size((1, 1_1, 7_6_8)   )
												self.assertEqual(output.shape   ,					_lowercase   )
												lowerCamelCase_							  =					torch.tensor(
												    [[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]   )

												self.assertTrue(torch.allclose(output[:, 1:4, 1:4]   ,					_lowercase   ,					atol=1E-4   )   )







from abc import ABC, abstractmethod
from argparse import ArgumentParser





class       __A(				UpperCAmelCase       ):

					@staticmethod
					@abstractmethod
					def    lowercase__      (	__UpperCamelCase					:     ArgumentParser   ):
												raise NotImplementedError()



					@abstractmethod
					def    lowercase__      (	self					:     Any   ):
												raise NotImplementedError()





