'''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 				_lowercase     (				UpperCAmelCase__			):


 '''simple docstring'''






 def __init__(      self   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	Any      ,       SCREAMING_SNAKE_CASE__   :	Any=13      ,       SCREAMING_SNAKE_CASE__   :	Dict=7      ,       SCREAMING_SNAKE_CASE__   :	str=True      ,       SCREAMING_SNAKE_CASE__   :	List[Any]=True      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]=False      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]=True      ,       SCREAMING_SNAKE_CASE__   :	Dict=99      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]=32      ,       SCREAMING_SNAKE_CASE__   :	Any=5      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]=4      ,       SCREAMING_SNAKE_CASE__   :	List[Any]=37      ,       SCREAMING_SNAKE_CASE__   :	Any="gelu"      ,       SCREAMING_SNAKE_CASE__   :	Tuple=0.1      ,       SCREAMING_SNAKE_CASE__   :	int=0.1      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]=5_12      ,       SCREAMING_SNAKE_CASE__   :	List[str]=16      ,       SCREAMING_SNAKE_CASE__   :	Any=2      ,       SCREAMING_SNAKE_CASE__   :	List[Any]=0.0_2      ,       SCREAMING_SNAKE_CASE__   :	List[Any]=3      ,       SCREAMING_SNAKE_CASE__   :	Any=4      ,       SCREAMING_SNAKE_CASE__   :	int=None      ,       )  ->    Union[str, Any]:
  __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 a					(      self   :	Tuple							)  ->    Any:
  __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 a					(      self   :	Union[str, Any]							)  ->    int:
  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 a					(      self   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]      ,       SCREAMING_SNAKE_CASE__   :	str      ,       SCREAMING_SNAKE_CASE__   :	str      ,       SCREAMING_SNAKE_CASE__   :	Any      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]							)  ->    List[str]:
  __lowerCAmelCase											=						DistilBertModel(config=SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()
  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__							)
  self.parent.assertEqual(result.last_hidden_state.shape      ,       (self.batch_size, self.seq_length, self.hidden_size)							)
 def a					(      self   :	str      ,       SCREAMING_SNAKE_CASE__   :	Dict      ,       SCREAMING_SNAKE_CASE__   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]      ,       SCREAMING_SNAKE_CASE__   :	Dict      ,       SCREAMING_SNAKE_CASE__   :	int							)  ->    List[Any]:
  __lowerCAmelCase											=						DistilBertForMaskedLM(config=SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()
  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       labels=SCREAMING_SNAKE_CASE__							)
  self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.seq_length, self.vocab_size)							)
 def a					(      self   :	Any      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	Any      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	List[str]							)  ->    Optional[Any]:
  __lowerCAmelCase											=						DistilBertForQuestionAnswering(config=SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()
  __lowerCAmelCase											=						model(
      SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       start_positions=SCREAMING_SNAKE_CASE__      ,       end_positions=SCREAMING_SNAKE_CASE__							)
  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 a					(      self   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]							)  ->    List[Any]:
  __lowerCAmelCase											=						self.num_labels
  __lowerCAmelCase											=						DistilBertForSequenceClassification(SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()
  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       labels=SCREAMING_SNAKE_CASE__							)
  self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.num_labels)							)
 def a					(      self   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]      ,       SCREAMING_SNAKE_CASE__   :	Dict      ,       SCREAMING_SNAKE_CASE__   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tuple							)  ->    Tuple:
  __lowerCAmelCase											=						self.num_labels
  __lowerCAmelCase											=						DistilBertForTokenClassification(config=SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  model.eval()

  __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       labels=SCREAMING_SNAKE_CASE__							)
  self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.seq_length, self.num_labels)							)
 def a					(      self   :	Any      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	str      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	Any      ,       SCREAMING_SNAKE_CASE__   :	int							)  ->    List[Any]:
  __lowerCAmelCase											=						self.num_choices
  __lowerCAmelCase											=						DistilBertForMultipleChoice(config=SCREAMING_SNAKE_CASE__							)
  model.to(SCREAMING_SNAKE_CASE__							)
  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(
      SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__      ,       labels=SCREAMING_SNAKE_CASE__      ,       )
  self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.num_choices)							)





 def a					(      self   :	Union[str, Any]							)  ->    List[Any]:
  __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 				_lowercase     (				UpperCAmelCase__  ,							UpperCAmelCase__  ,							unittest.TestCase			):


 '''simple docstring'''






 _SCREAMING_SNAKE_CASE     :       int      	=					(
     (
         DistilBertModel,
         DistilBertForMaskedLM,
         DistilBertForMultipleChoice,
         DistilBertForQuestionAnswering,
         DistilBertForSequenceClassification,
         DistilBertForTokenClassification,
     )
     if is_torch_available()
     else None
 )
 _SCREAMING_SNAKE_CASE     :       str      	=					(
     {
         """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     :       Dict      	=					True
 _SCREAMING_SNAKE_CASE     :       Tuple      	=					True
 _SCREAMING_SNAKE_CASE     :       List[Any]      	=					True
 _SCREAMING_SNAKE_CASE     :       int      	=					True
 def a					(      self   :	List[Any]							)  ->    Union[str, Any]:
  __lowerCAmelCase											=						DistilBertModelTester(self							)
  __lowerCAmelCase											=						ConfigTester(self      ,       config_class=SCREAMING_SNAKE_CASE__      ,       dim=37							)
 def a					(      self   :	Tuple							)  ->    List[str]:
  self.config_tester.run_common_tests()
 def a					(      self   :	List[Any]							)  ->    List[str]:
  __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_distilbert_model(*SCREAMING_SNAKE_CASE__							)
 def a					(      self   :	str							)  ->    int:
  __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_distilbert_for_masked_lm(*SCREAMING_SNAKE_CASE__							)
 def a					(      self   :	Any							)  ->    Tuple:
  __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_distilbert_for_question_answering(*SCREAMING_SNAKE_CASE__							)
 def a					(      self   :	Tuple							)  ->    List[Any]:
  __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_distilbert_for_sequence_classification(*SCREAMING_SNAKE_CASE__							)
 def a					(      self   :	str							)  ->    List[str]:
  __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_distilbert_for_token_classification(*SCREAMING_SNAKE_CASE__							)
 def a					(      self   :	Union[str, Any]							)  ->    List[Any]:
  __lowerCAmelCase											=						self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_distilbert_for_multiple_choice(*SCREAMING_SNAKE_CASE__							)
 @slow
 def a					(      self   :	int							)  ->    Optional[Any]:
  for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
   __lowerCAmelCase											=						DistilBertModel.from_pretrained(SCREAMING_SNAKE_CASE__							)
   self.assertIsNotNone(SCREAMING_SNAKE_CASE__							)





 @slow
 @require_torch_gpu
 def a					(      self   :	List[str]							)  ->    str:
  __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=SCREAMING_SNAKE_CASE__							)

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

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







@require_torch
class 				_lowercase     (				unittest.TestCase			):


 '''simple docstring'''






 @slow
 def a					(      self   :	Union[str, Any]							)  ->    Any:
  __lowerCAmelCase											=						DistilBertModel.from_pretrained("""distilbert-base-uncased"""							)
  __lowerCAmelCase											=						torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]							)
  __lowerCAmelCase											=						torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]							)
  with torch.no_grad():
   __lowerCAmelCase											=						model(SCREAMING_SNAKE_CASE__      ,       attention_mask=SCREAMING_SNAKE_CASE__							)[0]
  __lowerCAmelCase											=						torch.Size((1, 11, 7_68)							)
  self.assertEqual(output.shape      ,       SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						torch.tensor(
      [[[-0.1_6_3_9, 0.3_2_9_9, 0.1_6_4_8], [-0.1_7_4_6, 0.3_2_8_9, 0.1_7_1_0], [-0.1_8_8_4, 0.3_3_5_7, 0.1_8_1_0]]]							)

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







'''simple docstring'''

from abc import ABC, abstractmethod
from argparse import ArgumentParser







class 				_lowercase     (				UpperCAmelCase__			):


 '''simple docstring'''






 @staticmethod
 @abstractmethod
 def a					(      SCREAMING_SNAKE_CASE__   :	ArgumentParser							)  ->    Tuple:
  raise NotImplementedError()





 @abstractmethod
 def a					(      self   :	int							)  ->    Union[str, Any]:
  raise NotImplementedError()






