





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




					def __init__(						self			,  lowerCAmelCase_			,  lowerCAmelCase_=13			,  lowerCAmelCase_=7			,  lowerCAmelCase_=True			,  lowerCAmelCase_=True			,  lowerCAmelCase_=False			,  lowerCAmelCase_=True			,  lowerCAmelCase_=99			,  lowerCAmelCase_=32			,  lowerCAmelCase_=5			,  lowerCAmelCase_=4			,  lowerCAmelCase_=37			,  lowerCAmelCase_="gelu"			,  lowerCAmelCase_=0.1			,  lowerCAmelCase_=0.1			,  lowerCAmelCase_=512			,  lowerCAmelCase_=16			,  lowerCAmelCase_=2			,  lowerCAmelCase_=0.02			,  lowerCAmelCase_=3			,  lowerCAmelCase_=4			,  lowerCAmelCase_=None			,  )						->							str:
										_snake_case													=			parent
										_snake_case													=			batch_size
										_snake_case													=			seq_length
										_snake_case													=			is_training
										_snake_case													=			use_input_mask
										_snake_case													=			use_token_type_ids
										_snake_case													=			use_labels
										_snake_case													=			vocab_size
										_snake_case													=			hidden_size
										_snake_case													=			num_hidden_layers
										_snake_case													=			num_attention_heads
										_snake_case													=			intermediate_size
										_snake_case													=			hidden_act
										_snake_case													=			hidden_dropout_prob
										_snake_case													=			attention_probs_dropout_prob
										_snake_case													=			max_position_embeddings
										_snake_case													=			type_vocab_size
										_snake_case													=			type_sequence_label_size
										_snake_case													=			initializer_range
										_snake_case													=			num_labels
										_snake_case													=			num_choices
										_snake_case													=			scope




					def 		lowerCAmelCase     (						self					)						->							Optional[Any]:
										_snake_case													=			ids_tensor([self.batch_size, self.seq_length]			,  self.vocab_size					)

										_snake_case													=			None
										if self.use_input_mask:
															_snake_case													=			random_attention_mask([self.batch_size, self.seq_length]					)

										_snake_case													=			None
										_snake_case													=			None
										_snake_case													=			None
										if self.use_labels:
															_snake_case													=			ids_tensor([self.batch_size]			,  self.type_sequence_label_size					)
															_snake_case													=			ids_tensor([self.batch_size, self.seq_length]			,  self.num_labels					)
															_snake_case													=			ids_tensor([self.batch_size]			,  self.num_choices					)

										_snake_case													=			self.get_config()

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




					def 		lowerCAmelCase     (						self					)						->							List[str]:
										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 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)						->							List[str]:
										_snake_case													=			DistilBertModel(config=lowerCAmelCase_					)
										model.to(lowerCAmelCase_					)
										model.eval()
										_snake_case													=			model(lowerCAmelCase_			,  lowerCAmelCase_					)
										_snake_case													=			model(lowerCAmelCase_					)
										self.parent.assertEqual(result.last_hidden_state.shape			,  (self.batch_size, self.seq_length, self.hidden_size)					)




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)						->							Dict:
										_snake_case													=			DistilBertForMaskedLM(config=lowerCAmelCase_					)
										model.to(lowerCAmelCase_					)
										model.eval()
										_snake_case													=			model(lowerCAmelCase_			,  attention_mask=lowerCAmelCase_			,  labels=lowerCAmelCase_					)
										self.parent.assertEqual(result.logits.shape			,  (self.batch_size, self.seq_length, self.vocab_size)					)




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)						->							str:
										_snake_case													=			DistilBertForQuestionAnswering(config=lowerCAmelCase_					)
										model.to(lowerCAmelCase_					)
										model.eval()
										_snake_case													=			model(
										    lowerCAmelCase_			,  attention_mask=lowerCAmelCase_			,  start_positions=lowerCAmelCase_			,  end_positions=lowerCAmelCase_					)
										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 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)						->							int:
										_snake_case													=			self.num_labels
										_snake_case													=			DistilBertForSequenceClassification(lowerCAmelCase_					)
										model.to(lowerCAmelCase_					)
										model.eval()
										_snake_case													=			model(lowerCAmelCase_			,  attention_mask=lowerCAmelCase_			,  labels=lowerCAmelCase_					)
										self.parent.assertEqual(result.logits.shape			,  (self.batch_size, self.num_labels)					)




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)						->							Union[str, Any]:
										_snake_case													=			self.num_labels
										_snake_case													=			DistilBertForTokenClassification(config=lowerCAmelCase_					)
										model.to(lowerCAmelCase_					)
										model.eval()

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




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)						->							List[Any]:
										_snake_case													=			self.num_choices
										_snake_case													=			DistilBertForMultipleChoice(config=lowerCAmelCase_					)
										model.to(lowerCAmelCase_					)
										model.eval()
										_snake_case													=			input_ids.unsqueeze(1					).expand(-1			,  self.num_choices			,  -1					).contiguous()
										_snake_case													=			input_mask.unsqueeze(1					).expand(-1			,  self.num_choices			,  -1					).contiguous()
										_snake_case													=			model(
										    lowerCAmelCase_			,  attention_mask=lowerCAmelCase_			,  labels=lowerCAmelCase_			,  )
										self.parent.assertEqual(result.logits.shape			,  (self.batch_size, self.num_choices)					)






					def 		lowerCAmelCase     (						self					)						->							str:
										_snake_case													=			self.prepare_config_and_inputs()
										((_snake_case)							,       (_snake_case)							,       (_snake_case)							,       (_snake_case)							,       (_snake_case)							,       (_snake_case))													=			config_and_inputs
										_snake_case													=			{'input_ids': input_ids, 'attention_mask': input_mask}
										return config, inputs_dict





@require_torch
class      UpperCamelCase_						(					_lowerCamelCase							,					_lowerCamelCase							,					unittest.TestCase   ):
					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 		lowerCAmelCase     (						self					)						->							Tuple:
										_snake_case													=			DistilBertModelTester(self					)
										_snake_case													=			ConfigTester(self			,  config_class=lowerCAmelCase_			,  dim=37					)




					def 		lowerCAmelCase     (						self					)						->							Dict:
										self.config_tester.run_common_tests()




					def 		lowerCAmelCase     (						self					)						->							str:
										_snake_case													=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_model(*lowerCAmelCase_					)




					def 		lowerCAmelCase     (						self					)						->							Tuple:
										_snake_case													=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_for_masked_lm(*lowerCAmelCase_					)




					def 		lowerCAmelCase     (						self					)						->							List[Any]:
										_snake_case													=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_for_question_answering(*lowerCAmelCase_					)




					def 		lowerCAmelCase     (						self					)						->							Dict:
										_snake_case													=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowerCAmelCase_					)




					def 		lowerCAmelCase     (						self					)						->							List[str]:
										_snake_case													=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_for_token_classification(*lowerCAmelCase_					)




					def 		lowerCAmelCase     (						self					)						->							Any:
										_snake_case													=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowerCAmelCase_					)




					@slow
					def 		lowerCAmelCase     (						self					)						->							Any:
										for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_snake_case													=			DistilBertModel.from_pretrained(lowerCAmelCase_					)
															self.assertIsNotNone(lowerCAmelCase_					)






					@slow
					@require_torch_gpu
					def 		lowerCAmelCase     (						self					)						->							Any:
										_snake_case							,       _snake_case													=			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

															_snake_case													=			True
															_snake_case													=			model_class(config=lowerCAmelCase_					)

															_snake_case													=			self._prepare_for_class(lowerCAmelCase_			,  lowerCAmelCase_					)
															_snake_case													=			torch.jit.trace(
															    lowerCAmelCase_			,  (inputs_dict['input_ids'].to('cpu'					), inputs_dict['attention_mask'].to('cpu'					))					)

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





@require_torch
class      UpperCamelCase_						(					unittest.TestCase   ):




					@slow
					def 		lowerCAmelCase     (						self					)						->							Dict:
										_snake_case													=			DistilBertModel.from_pretrained('distilbert-base-uncased'					)
										_snake_case													=			torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]					)
										_snake_case													=			torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]					)
										with torch.no_grad():
															_snake_case													=			model(lowerCAmelCase_			,  attention_mask=lowerCAmelCase_					)[0]
										_snake_case													=			torch.Size((1, 11, 768)					)
										self.assertEqual(output.shape			,  lowerCAmelCase_					)
										_snake_case													=			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]			,  lowerCAmelCase_			,  atol=1E-4					)					)










import random





def   lowerCamelCase__       (       UpperCamelCase__     :    int    ,      UpperCamelCase__     :    float    ,      UpperCamelCase__     :    bool = False  )     ->	dict:



					'''simple docstring'''




					_snake_case													=			{i: [] for i in range(UpperCamelCase__  )}

					# if probability is greater or equal than 1, then generate a complete graph
					if probability >= 1:
										return complete_graph(UpperCamelCase__  )
					# if probability is lower or equal than 0, then return a graph without edges
					if probability <= 0:
										return graph

					# for each couple of nodes, add an edge from u to v
					# if the number randomly generated is greater than probability probability
					for i in range(UpperCamelCase__  ):
										for j in range(i + 1    ,      UpperCamelCase__  ):
															if random.random() < probability:
																				graph[i].append(UpperCamelCase__  )
																				if not directed:
																									# if the graph is undirected, add an edge in from j to i, either
																									graph[j].append(UpperCamelCase__  )
					return graph





def   lowerCamelCase__       (       UpperCamelCase__     :    int  )     ->	dict:



					'''simple docstring'''




					return {
					    i: [j for j in range(UpperCamelCase__  ) if i != j] for i in range(UpperCamelCase__  )
					}


if __name__ == "__main__":
					import doctest

					doctest.testmod()



