





'''simple docstring'''




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  SCREAMING_SNAKE_CASE    (a__							):


     def __init__(						self						,	_UpperCAmelCase						,	_UpperCAmelCase=13						,	_UpperCAmelCase=7						,	_UpperCAmelCase=True						,	_UpperCAmelCase=True						,	_UpperCAmelCase=False						,	_UpperCAmelCase=True						,	_UpperCAmelCase=99						,	_UpperCAmelCase=32						,	_UpperCAmelCase=5						,	_UpperCAmelCase=4						,	_UpperCAmelCase=37						,	_UpperCAmelCase="gelu"						,	_UpperCAmelCase=0.1						,	_UpperCAmelCase=0.1						,	_UpperCAmelCase=512						,	_UpperCAmelCase=16						,	_UpperCAmelCase=2						,	_UpperCAmelCase=0.02						,	_UpperCAmelCase=3						,	_UpperCAmelCase=4						,	_UpperCAmelCase=None						,	):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       ids_tensor([self.batch_size, self.seq_length]						,	self.vocab_size)

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

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

          __A						: List[str]      =       self.get_config()

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''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 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: Any      =       DistilBertModel(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A						: int      =       model(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: Optional[int]      =       model(_UpperCAmelCase)
          self.parent.assertEqual(result.last_hidden_state.shape						,	(self.batch_size, self.seq_length, self.hidden_size))


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: List[Any]      =       DistilBertForMaskedLM(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A						: Optional[int]      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	labels=_UpperCAmelCase)
          self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.vocab_size))


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: Dict      =       DistilBertForQuestionAnswering(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A						: List[str]      =       model(
              _UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	start_positions=_UpperCAmelCase						,	end_positions=_UpperCAmelCase)
          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 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: Union[str, Any]      =       self.num_labels
          __A						: Union[str, Any]      =       DistilBertForSequenceClassification(_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A						: Optional[Any]      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	labels=_UpperCAmelCase)
          self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.num_labels))


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: Any      =       self.num_labels
          __A						: Union[str, Any]      =       DistilBertForTokenClassification(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()

          __A						: List[Any]      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	labels=_UpperCAmelCase)
          self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.num_labels))


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       self.prepare_config_and_inputs()
          ((__A) ,(__A) ,(__A) ,(__A) ,(__A) ,(__A))						: str      =       config_and_inputs
          __A						: Union[str, Any]      =       {'input_ids': input_ids, 'attention_mask': input_mask}
          return config, inputs_dict




@require_torch
class  SCREAMING_SNAKE_CASE    (a__					,		a__					,		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 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Optional[Any]      =       DistilBertModelTester(self)
          __A						: Optional[int]      =       ConfigTester(self						,	config_class=_UpperCAmelCase						,	dim=37)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          self.config_tester.run_common_tests()


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Optional[int]      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_model(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_masked_lm(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_question_answering(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Tuple      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_sequence_classification(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_token_classification(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[Any]      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_distilbert_for_multiple_choice(*_UpperCAmelCase)


     @slow
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __A						: List[str]      =       DistilBertModel.from_pretrained(_UpperCAmelCase)
               self.assertIsNotNone(_UpperCAmelCase)


     @slow
     @require_torch_gpu
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: Union[str, 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

               __A						: int      =       True
               __A						: List[str]      =       model_class(config=_UpperCAmelCase)

               __A						: str      =       self._prepare_for_class(_UpperCAmelCase						,	_UpperCAmelCase)
               __A						: Dict      =       torch.jit.trace(
                   _UpperCAmelCase						,	(inputs_dict['input_ids'].to('cpu'), inputs_dict['attention_mask'].to('cpu')))

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




@require_torch
class  SCREAMING_SNAKE_CASE    (unittest.TestCase							):


     @slow
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       DistilBertModel.from_pretrained('distilbert-base-uncased')
          __A						: Optional[int]      =       torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]])
          __A						: int      =       torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
          with torch.no_grad():
               __A						: List[str]      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase)[0]
          __A						: Any      =       torch.Size((1, 11, 768))
          self.assertEqual(output.shape						,	_UpperCAmelCase)
          __A						: int      =       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]						,	_UpperCAmelCase						,	atol=1e-4))






'''simple docstring'''




import random






def        _lowerCAmelCase			(  __snake_case  : int ,							__snake_case  : float ,							__snake_case  : bool = False			)     ->     dict:
     __A						: dict      =       {i: [] for i in range(__snake_case			)}

     # if probability is greater or equal than 1, then generate a complete graph
     if probability >= 1:
          return complete_graph(__snake_case			)
     # 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(__snake_case			):
          for j in range(i + 1 ,							__snake_case			):
               if random.random() < probability:
                    graph[i].append(__snake_case			)
                    if not directed:
                         # if the graph is undirected, add an edge in from j to i, either
                         graph[j].append(__snake_case			)
     return graph






def        _lowerCAmelCase			(  __snake_case  : int			)     ->     dict:
     return {
         i: [j for j in range(__snake_case			) if i != j] for i in range(__snake_case			)
     }


if __name__ == "__main__":
  import doctest

  doctest.testmod()