
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     __UpperCamelCase (							lowerCAmelCase__       ):






     """simple docstring"""







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







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		parent
       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		batch_size
       __SCREAMING_SNAKE_CASE      :	str					=		seq_length
       __SCREAMING_SNAKE_CASE      :	int					=		is_training
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		use_input_mask
       __SCREAMING_SNAKE_CASE      :	str					=		use_token_type_ids
       __SCREAMING_SNAKE_CASE      :	Any					=		use_labels
       __SCREAMING_SNAKE_CASE      :	Any					=		vocab_size
       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		hidden_size
       __SCREAMING_SNAKE_CASE      :	Any					=		num_hidden_layers
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		num_attention_heads
       __SCREAMING_SNAKE_CASE      :	List[str]					=		intermediate_size
       __SCREAMING_SNAKE_CASE      :	List[str]					=		hidden_act
       __SCREAMING_SNAKE_CASE      :	int					=		hidden_dropout_prob
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		attention_probs_dropout_prob
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		max_position_embeddings
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		type_vocab_size
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		type_sequence_label_size
       __SCREAMING_SNAKE_CASE      :	int					=		initializer_range
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		num_labels
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		num_choices
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		scope



     def  UpperCAmelCase__							(				self							:							List[str] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	List[Any]					=		ids_tensor([self.batch_size, self.seq_length]     ,						self.vocab_size )

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

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

       __SCREAMING_SNAKE_CASE      :	Dict					=		self.get_config()

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



     def  UpperCAmelCase__							(				self							:							Optional[int] ):







       """simple docstring"""




       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  UpperCAmelCase__							(				self							:							Optional[int]     ,						_A							:							int     ,						_A							:							Union[str, Any]     ,						_A							:							List[str]     ,						_A							:							Dict     ,						_A							:							Dict     ,						_A							:							Optional[int] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	List[Any]					=		DistilBertModel(config=_A )
       model.to(_A )
       model.eval()
       __SCREAMING_SNAKE_CASE      :	Dict					=		model(_A     ,						_A )
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		model(_A )
       self.parent.assertEqual(result.last_hidden_state.shape     ,						(self.batch_size, self.seq_length, self.hidden_size) )



     def  UpperCAmelCase__							(				self							:							Tuple     ,						_A							:							Dict     ,						_A							:							Tuple     ,						_A							:							str     ,						_A							:							Optional[int]     ,						_A							:							List[str]     ,						_A							:							Tuple ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		DistilBertForMaskedLM(config=_A )
       model.to(_A )
       model.eval()
       __SCREAMING_SNAKE_CASE      :	Tuple					=		model(_A     ,						attention_mask=_A     ,						labels=_A )
       self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.seq_length, self.vocab_size) )



     def  UpperCAmelCase__							(				self							:							Dict     ,						_A							:							Optional[Any]     ,						_A							:							Optional[Any]     ,						_A							:							Union[str, Any]     ,						_A							:							Optional[Any]     ,						_A							:							str     ,						_A							:							Tuple ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		DistilBertForQuestionAnswering(config=_A )
       model.to(_A )
       model.eval()
       __SCREAMING_SNAKE_CASE      :	int					=		model(
           _A     ,						attention_mask=_A     ,						start_positions=_A     ,						end_positions=_A )
       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							:							Dict     ,						_A							:							List[str]     ,						_A							:							Tuple     ,						_A							:							str     ,						_A							:							Tuple     ,						_A							:							Optional[int]     ,						_A							:							Any ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		self.num_labels
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		DistilBertForSequenceClassification(_A )
       model.to(_A )
       model.eval()
       __SCREAMING_SNAKE_CASE      :	Dict					=		model(_A     ,						attention_mask=_A     ,						labels=_A )
       self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.num_labels) )



     def  UpperCAmelCase__							(				self							:							List[str]     ,						_A							:							int     ,						_A							:							List[Any]     ,						_A							:							Any     ,						_A							:							Any     ,						_A							:							str     ,						_A							:							List[Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		self.num_labels
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		DistilBertForTokenClassification(config=_A )
       model.to(_A )
       model.eval()

       __SCREAMING_SNAKE_CASE      :	Dict					=		model(_A     ,						attention_mask=_A     ,						labels=_A )
       self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.seq_length, self.num_labels) )



     def  UpperCAmelCase__							(				self							:							Dict     ,						_A							:							Optional[int]     ,						_A							:							int     ,						_A							:							Optional[int]     ,						_A							:							List[Any]     ,						_A							:							int     ,						_A							:							Dict ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Any					=		self.num_choices
       __SCREAMING_SNAKE_CASE      :	int					=		DistilBertForMultipleChoice(config=_A )
       model.to(_A )
       model.eval()
       __SCREAMING_SNAKE_CASE      :	List[str]					=		input_ids.unsqueeze(1 ).expand(-1     ,						self.num_choices     ,						-1 ).contiguous()
       __SCREAMING_SNAKE_CASE      :	Any					=		input_mask.unsqueeze(1 ).expand(-1     ,						self.num_choices     ,						-1 ).contiguous()
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		model(
           _A     ,						attention_mask=_A     ,						labels=_A     ,						)
       self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.num_choices) )




     def  UpperCAmelCase__							(				self							:							int ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Dict					=		self.prepare_config_and_inputs()
       ((__SCREAMING_SNAKE_CASE),		(__SCREAMING_SNAKE_CASE),		(__SCREAMING_SNAKE_CASE),		(__SCREAMING_SNAKE_CASE),		(__SCREAMING_SNAKE_CASE),		(__SCREAMING_SNAKE_CASE))      :	List[Any]					=		config_and_inputs
       __SCREAMING_SNAKE_CASE      :	Dict					=		{'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict






@require_torch
class     __UpperCamelCase (							lowerCAmelCase__ ,						lowerCAmelCase__ ,						unittest.TestCase       ):






     """simple docstring"""







     lowerCAmelCase_							   =			(
         (
             DistilBertModel,
             DistilBertForMaskedLM,
             DistilBertForMultipleChoice,
             DistilBertForQuestionAnswering,
             DistilBertForSequenceClassification,
             DistilBertForTokenClassification,
         )
         if is_torch_available()
         else None
     )
     lowerCAmelCase_							   =			(
         {
             '''feature-extraction''': DistilBertModel,
             '''fill-mask''': DistilBertForMaskedLM,
             '''question-answering''': DistilBertForQuestionAnswering,
             '''text-classification''': DistilBertForSequenceClassification,
             '''token-classification''': DistilBertForTokenClassification,
             '''zero-shot''': DistilBertForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     lowerCAmelCase_							   =			True
     lowerCAmelCase_							   =			True
     lowerCAmelCase_							   =			True
     lowerCAmelCase_							   =			True



     def  UpperCAmelCase__							(				self							:							Any ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Dict					=		DistilBertModelTester(self )
       __SCREAMING_SNAKE_CASE      :	List[str]					=		ConfigTester(self     ,						config_class=_A     ,						dim=37 )



     def  UpperCAmelCase__							(				self							:							List[Any] ):







       """simple docstring"""




       self.config_tester.run_common_tests()



     def  UpperCAmelCase__							(				self							:							Any ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Tuple					=		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_model(*_A )



     def  UpperCAmelCase__							(				self							:							Tuple ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_masked_lm(*_A )



     def  UpperCAmelCase__							(				self							:							Union[str, Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	List[Any]					=		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_question_answering(*_A )



     def  UpperCAmelCase__							(				self							:							str ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Dict					=		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_sequence_classification(*_A )



     def  UpperCAmelCase__							(				self							:							List[Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_token_classification(*_A )



     def  UpperCAmelCase__							(				self							:							List[Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	int					=		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_multiple_choice(*_A )



     @slow
     def  UpperCAmelCase__							(				self							:							Union[str, Any] ):







       """simple docstring"""




       for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         __SCREAMING_SNAKE_CASE      :	List[Any]					=		DistilBertModel.from_pretrained(_A )
         self.assertIsNotNone(_A )




     @slow
     @require_torch_gpu
     def  UpperCAmelCase__							(				self							:							str ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE      :	List[Any]					=		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

         __SCREAMING_SNAKE_CASE      :	Dict					=		True
         __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		model_class(config=_A )

         __SCREAMING_SNAKE_CASE      :	int					=		self._prepare_for_class(_A     ,						_A )
         __SCREAMING_SNAKE_CASE      :	List[Any]					=		torch.jit.trace(
             _A     ,						(inputs_dict['''input_ids'''].to('''cpu''' ), inputs_dict['''attention_mask'''].to('''cpu''' )) )

         with tempfile.TemporaryDirectory() as tmp:
           torch.jit.save(_A     ,						os.path.join(_A     ,						'''traced_model.pt''' ) )
           __SCREAMING_SNAKE_CASE      :	Optional[int]					=		torch.jit.load(os.path.join(_A     ,						'''traced_model.pt''' )     ,						map_location=_A )
           loaded(inputs_dict['''input_ids'''].to(_A )     ,						inputs_dict['''attention_mask'''].to(_A ) )






@require_torch
class     __UpperCamelCase (							unittest.TestCase       ):






     """simple docstring"""







     @slow
     def  UpperCAmelCase__							(				self							:							Dict ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	List[Any]					=		DistilBertModel.from_pretrained('''distilbert-base-uncased''' )
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]] )
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] )
       with torch.no_grad():
         __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		model(_A     ,						attention_mask=_A )[0]
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		torch.Size((1, 11, 768) )
       self.assertEqual(output.shape     ,						_A )
       __SCREAMING_SNAKE_CASE      :	Any					=		torch.tensor(
           [[[-0.16_39, 0.32_99, 0.16_48], [-0.17_46, 0.32_89, 0.17_10], [-0.18_84, 0.33_57, 0.18_10]]] )

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


from ..utils import is_flax_available, is_torch_available


if is_torch_available():
							from .autoencoder_kl import AutoencoderKL
							from .controlnet import ControlNetModel
							from .dual_transformer_ad import DualTransformeraDModel
							from .modeling_utils import ModelMixin
							from .prior_transformer import PriorTransformer
							from .ta_film_transformer import TaFilmDecoder
							from .transformer_ad import TransformeraDModel
							from .unet_ad import UNetaDModel
							from .unet_ad import UNetaDModel
							from .unet_ad_condition import UNetaDConditionModel
							from .unet_ad_condition import UNetaDConditionModel
							from .vq_model import VQModel

if is_flax_available():
							from .controlnet_flax import FlaxControlNetModel
							from .unet_ad_condition_flax import FlaxUNetaDConditionModel
							from .vae_flax import FlaxAutoencoderKL






