



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_	):


						'''simple docstring'''





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






							'''simple docstring'''



							__lowercase		       =			parent
							__lowercase		       =			batch_size
							__lowercase		       =			seq_length
							__lowercase		       =			is_training
							__lowercase		       =			use_input_mask
							__lowercase		       =			use_token_type_ids
							__lowercase		       =			use_labels
							__lowercase		       =			vocab_size
							__lowercase		       =			hidden_size
							__lowercase		       =			num_hidden_layers
							__lowercase		       =			num_attention_heads
							__lowercase		       =			intermediate_size
							__lowercase		       =			hidden_act
							__lowercase		       =			hidden_dropout_prob
							__lowercase		       =			attention_probs_dropout_prob
							__lowercase		       =			max_position_embeddings
							__lowercase		       =			type_vocab_size
							__lowercase		       =			type_sequence_label_size
							__lowercase		       =			initializer_range
							__lowercase		       =			num_labels
							__lowercase		       =			num_choices
							__lowercase		       =			scope



						def     _snake_case	(	self	:							Dict	):






							'''simple docstring'''



							__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		       =			self.get_config()

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



						def     _snake_case	(	self	:							Any	):






							'''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     _snake_case	(	self	:							Dict							,      lowerCamelCase	:							Optional[int]							,      lowerCamelCase	:							Dict							,      lowerCamelCase	:							Optional[int]							,      lowerCamelCase	:							Union[str, Any]							,      lowerCamelCase	:							Union[str, Any]							,      lowerCamelCase	:							Optional[int]	):






							'''simple docstring'''



							__lowercase		       =			DistilBertModel(config=_lowercase	)
							model.to(_lowercase	)
							model.eval()
							__lowercase		       =			model(_lowercase							,      _lowercase	)
							__lowercase		       =			model(_lowercase	)
							self.parent.assertEqual(result.last_hidden_state.shape							,      (self.batch_size, self.seq_length, self.hidden_size)	)



						def     _snake_case	(	self	:							Tuple							,      lowerCamelCase	:							int							,      lowerCamelCase	:							Tuple							,      lowerCamelCase	:							Optional[Any]							,      lowerCamelCase	:							str							,      lowerCamelCase	:							List[str]							,      lowerCamelCase	:							Optional[int]	):






							'''simple docstring'''



							__lowercase		       =			DistilBertForMaskedLM(config=_lowercase	)
							model.to(_lowercase	)
							model.eval()
							__lowercase		       =			model(_lowercase							,      attention_mask=_lowercase							,      labels=_lowercase	)
							self.parent.assertEqual(result.logits.shape							,      (self.batch_size, self.seq_length, self.vocab_size)	)



						def     _snake_case	(	self	:							Any							,      lowerCamelCase	:							str							,      lowerCamelCase	:							Dict							,      lowerCamelCase	:							Optional[Any]							,      lowerCamelCase	:							Dict							,      lowerCamelCase	:							List[str]							,      lowerCamelCase	:							Optional[Any]	):






							'''simple docstring'''



							__lowercase		       =			DistilBertForQuestionAnswering(config=_lowercase	)
							model.to(_lowercase	)
							model.eval()
							__lowercase		       =			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     _snake_case	(	self	:							Any							,      lowerCamelCase	:							List[str]							,      lowerCamelCase	:							Optional[int]							,      lowerCamelCase	:							List[Any]							,      lowerCamelCase	:							Optional[Any]							,      lowerCamelCase	:							int							,      lowerCamelCase	:							Dict	):






							'''simple docstring'''



							__lowercase		       =			self.num_labels
							__lowercase		       =			DistilBertForSequenceClassification(_lowercase	)
							model.to(_lowercase	)
							model.eval()
							__lowercase		       =			model(_lowercase							,      attention_mask=_lowercase							,      labels=_lowercase	)
							self.parent.assertEqual(result.logits.shape							,      (self.batch_size, self.num_labels)	)



						def     _snake_case	(	self	:							Optional[int]							,      lowerCamelCase	:							List[Any]							,      lowerCamelCase	:							Dict							,      lowerCamelCase	:							Optional[Any]							,      lowerCamelCase	:							int							,      lowerCamelCase	:							Dict							,      lowerCamelCase	:							List[str]	):






							'''simple docstring'''



							__lowercase		       =			self.num_labels
							__lowercase		       =			DistilBertForTokenClassification(config=_lowercase	)
							model.to(_lowercase	)
							model.eval()

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



						def     _snake_case	(	self	:							List[Any]							,      lowerCamelCase	:							List[Any]							,      lowerCamelCase	:							Union[str, Any]							,      lowerCamelCase	:							Dict							,      lowerCamelCase	:							Tuple							,      lowerCamelCase	:							Any							,      lowerCamelCase	:							Union[str, Any]	):






							'''simple docstring'''



							__lowercase		       =			self.num_choices
							__lowercase		       =			DistilBertForMultipleChoice(config=_lowercase	)
							model.to(_lowercase	)
							model.eval()
							__lowercase		       =			input_ids.unsqueeze(1	).expand(-1							,      self.num_choices							,      -1	).contiguous()
							__lowercase		       =			input_mask.unsqueeze(1	).expand(-1							,      self.num_choices							,      -1	).contiguous()
							__lowercase		       =			model(
							    _lowercase							,      attention_mask=_lowercase							,      labels=_lowercase							,      )
							self.parent.assertEqual(result.logits.shape							,      (self.batch_size, self.num_choices)	)



						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							__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_torch
class   _A   ( UpperCAmelCase_  ,						UpperCAmelCase_  ,						unittest.TestCase	):


						'''simple docstring'''





						_snake_case			:     Optional[int]      							=      (
						    (
						        DistilBertModel,
						        DistilBertForMaskedLM,
						        DistilBertForMultipleChoice,
						        DistilBertForQuestionAnswering,
						        DistilBertForSequenceClassification,
						        DistilBertForTokenClassification,
						    )
						    if is_torch_available()
						    else None
						)
						_snake_case			:     List[Any]      							=      (
						    {
						        """feature-extraction""": DistilBertModel,
						        """fill-mask""": DistilBertForMaskedLM,
						        """question-answering""": DistilBertForQuestionAnswering,
						        """text-classification""": DistilBertForSequenceClassification,
						        """token-classification""": DistilBertForTokenClassification,
						        """zero-shot""": DistilBertForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						_snake_case			:     Optional[Any]      							=      True
						_snake_case			:     List[str]      							=      True
						_snake_case			:     Tuple      							=      True
						_snake_case			:     Any      							=      True



						def     _snake_case	(	self	:							Dict	):






							'''simple docstring'''



							__lowercase		       =			DistilBertModelTester(self	)
							__lowercase		       =			ConfigTester(self							,      config_class=_lowercase							,      dim=37	)



						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							self.config_tester.run_common_tests()



						def     _snake_case	(	self	:							List[Any]	):






							'''simple docstring'''



							__lowercase		       =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_model(*_lowercase	)



						def     _snake_case	(	self	:							str	):






							'''simple docstring'''



							__lowercase		       =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_masked_lm(*_lowercase	)



						def     _snake_case	(	self	:							Any	):






							'''simple docstring'''



							__lowercase		       =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_question_answering(*_lowercase	)



						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							__lowercase		       =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_sequence_classification(*_lowercase	)



						def     _snake_case	(	self	:							str	):






							'''simple docstring'''



							__lowercase		       =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_token_classification(*_lowercase	)



						def     _snake_case	(	self	:							str	):






							'''simple docstring'''



							__lowercase		       =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_multiple_choice(*_lowercase	)



						@slow
						def     _snake_case	(	self	:							Tuple	):






							'''simple docstring'''



							for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								__lowercase		       =			DistilBertModel.from_pretrained(_lowercase	)
								self.assertIsNotNone(_lowercase	)



						@slow
						@require_torch_gpu
						def     _snake_case	(	self	:							Optional[int]	):






							'''simple docstring'''



							__lowercase       ,    __lowercase		       =			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

								__lowercase		       =			True
								__lowercase		       =			model_class(config=_lowercase	)

								__lowercase		       =			self._prepare_for_class(_lowercase							,      _lowercase	)
								__lowercase		       =			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"	)	)
									__lowercase		       =			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	):


						'''simple docstring'''





						@slow
						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							__lowercase		       =			DistilBertModel.from_pretrained("distilbert-base-uncased"	)
							__lowercase		       =			torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]	)
							__lowercase		       =			torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]	)
							with torch.no_grad():
								__lowercase		       =			model(_lowercase							,      attention_mask=_lowercase	)[0]
							__lowercase		       =			torch.Size((1, 11, 768)	)
							self.assertEqual(output.shape							,      _lowercase	)
							__lowercase		       =			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 ....utils import logging


snake_case__					:     List[Any]    		=   logging.get_logger(__name__)

class   _A   ( _lowercase	):


						'''simple docstring'''





						def __init__(	self	:							List[str]							,      lowerCamelCase	:							Any							,      lowerCamelCase	:							Dict=None							,      lowerCamelCase	:							Dict=2_048	):






							'''simple docstring'''



							__lowercase		       =			config.__dict__
							__lowercase		       =			modal_hidden_size
							if num_labels:
								__lowercase		       =			num_labels



