

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				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : List[str]=1_3				,				UpperCAmelCase__    : int=7				,				UpperCAmelCase__    : int=True				,				UpperCAmelCase__    : List[str]=True				,				UpperCAmelCase__    : str=False				,				UpperCAmelCase__    : str=True				,				UpperCAmelCase__    : Tuple=9_9				,				UpperCAmelCase__    : List[str]=3_2				,				UpperCAmelCase__    : Optional[int]=5				,				UpperCAmelCase__    : Union[str, Any]=4				,				UpperCAmelCase__    : int=3_7				,				UpperCAmelCase__    : Union[str, Any]="gelu"				,				UpperCAmelCase__    : Tuple=0.1				,				UpperCAmelCase__    : Optional[int]=0.1				,				UpperCAmelCase__    : Union[str, Any]=5_1_2				,				UpperCAmelCase__    : Optional[Any]=1_6				,				UpperCAmelCase__    : List[str]=2				,				UpperCAmelCase__    : int=0.02				,				UpperCAmelCase__    : str=3				,				UpperCAmelCase__    : Optional[int]=4				,				UpperCAmelCase__    : Optional[int]=None				,				)					->Union[str, Any]:
        """simple docstring"""







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


  def  _lowercase      (			self    : Dict		)					->str:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      int       =   ids_tensor([self.batch_size, self.seq_length]				,				self.vocab_size		)

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

        SCREAMING_SNAKE_CASE :      Optional[int]       =   None
        SCREAMING_SNAKE_CASE :      List[str]       =   None
        SCREAMING_SNAKE_CASE :      Tuple       =   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 :      Union[str, Any]       =   ids_tensor([self.batch_size]				,				self.num_choices		)

        SCREAMING_SNAKE_CASE :      List[Any]       =   self.get_config()

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


  def  _lowercase      (			self    : Any		)					->List[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  _lowercase      (			self    : Dict				,				UpperCAmelCase__    : Optional[int]				,				UpperCAmelCase__    : Dict				,				UpperCAmelCase__    : Optional[int]				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Optional[int]		)					->Any:
        """simple docstring"""







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


  def  _lowercase      (			self    : Tuple				,				UpperCAmelCase__    : int				,				UpperCAmelCase__    : Tuple				,				UpperCAmelCase__    : Optional[Any]				,				UpperCAmelCase__    : str				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Optional[int]		)					->List[str]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[int]       =   DistilBertForMaskedLM(config=_lowercase		)
        model.to(_lowercase		)
        model.eval()
        SCREAMING_SNAKE_CASE :      str       =   model(_lowercase				,				attention_mask=_lowercase				,				labels=_lowercase		)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.seq_length, self.vocab_size)		)


  def  _lowercase      (			self    : Any				,				UpperCAmelCase__    : str				,				UpperCAmelCase__    : Dict				,				UpperCAmelCase__    : Optional[Any]				,				UpperCAmelCase__    : Dict				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Optional[Any]		)					->str:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   DistilBertForQuestionAnswering(config=_lowercase		)
        model.to(_lowercase		)
        model.eval()
        SCREAMING_SNAKE_CASE :      Tuple       =   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  _lowercase      (			self    : Any				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Optional[int]				,				UpperCAmelCase__    : List[Any]				,				UpperCAmelCase__    : Optional[Any]				,				UpperCAmelCase__    : int				,				UpperCAmelCase__    : Dict		)					->Optional[Any]:
        """simple docstring"""







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


  def  _lowercase      (			self    : Optional[int]				,				UpperCAmelCase__    : List[Any]				,				UpperCAmelCase__    : Dict				,				UpperCAmelCase__    : Optional[Any]				,				UpperCAmelCase__    : int				,				UpperCAmelCase__    : Dict				,				UpperCAmelCase__    : List[str]		)					->Tuple:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Any       =   self.num_labels
        SCREAMING_SNAKE_CASE :      List[Any]       =   DistilBertForTokenClassification(config=_lowercase		)
        model.to(_lowercase		)
        model.eval()

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


  def  _lowercase      (			self    : List[Any]				,				UpperCAmelCase__    : List[Any]				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Dict				,				UpperCAmelCase__    : Tuple				,				UpperCAmelCase__    : Any				,				UpperCAmelCase__    : Union[str, Any]		)					->List[str]:
        """simple docstring"""







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




  def  _lowercase      (			self    : int		)					->List[str]:
        """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[str]       =   config_and_inputs
        SCREAMING_SNAKE_CASE :      int       =   {"""input_ids""": input_ids, """attention_mask""": input_mask}
        return config, inputs_dict







@require_torch
class        a__       (   UpperCAmelCase_	,			UpperCAmelCase_	,			unittest.TestCase       ):







  """simple docstring"""




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


  def  _lowercase      (			self    : Dict		)					->Tuple:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      int       =   DistilBertModelTester(self		)
        SCREAMING_SNAKE_CASE :      List[Any]       =   ConfigTester(self				,				config_class=_lowercase				,				dim=3_7		)


  def  _lowercase      (			self    : int		)					->Tuple:
        """simple docstring"""







        self.config_tester.run_common_tests()


  def  _lowercase      (			self    : List[Any]		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[Any]       =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_model(*_lowercase		)


  def  _lowercase      (			self    : str		)					->int:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      str       =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_masked_lm(*_lowercase		)


  def  _lowercase      (			self    : Any		)					->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Tuple       =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_question_answering(*_lowercase		)


  def  _lowercase      (			self    : int		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[Any]       =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_sequence_classification(*_lowercase		)


  def  _lowercase      (			self    : str		)					->int:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Any       =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_token_classification(*_lowercase		)


  def  _lowercase      (			self    : str		)					->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      int       =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_multiple_choice(*_lowercase		)


  @slow
  def  _lowercase      (			self    : Tuple		)					->Optional[int]:
        """simple docstring"""







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




  @slow
  @require_torch_gpu
  def  _lowercase      (			self    : Optional[int]		)					->Tuple:
        """simple docstring"""







        SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      List[str]       =   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 :      int       =   True
              SCREAMING_SNAKE_CASE :      Optional[int]       =   model_class(config=_lowercase		)

              SCREAMING_SNAKE_CASE :      Any       =   self._prepare_for_class(_lowercase				,				_lowercase		)
              SCREAMING_SNAKE_CASE :      Dict       =   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"""		)		)
                    SCREAMING_SNAKE_CASE :      List[Any]       =   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  _lowercase      (			self    : int		)					->Tuple:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   DistilBertModel.from_pretrained("""distilbert-base-uncased"""		)
        SCREAMING_SNAKE_CASE :      Any       =   torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]		)
        SCREAMING_SNAKE_CASE :      Dict       =   torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
        with torch.no_grad():
              SCREAMING_SNAKE_CASE :      Optional[int]       =   model(_lowercase				,				attention_mask=_lowercase		)[0]
        SCREAMING_SNAKE_CASE :      Optional[int]       =   torch.Size((1, 1_1, 7_6_8)		)
        self.assertEqual(output.shape				,				_lowercase		)
        SCREAMING_SNAKE_CASE :      str       =   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]				,				_lowercase				,				atol=1e-4		)		)





def 					__lowercase      (     _A						,    _A			)     ->						int:
      return int((input_a, input_a).count(1			) != 0			)






def 					__lowercase      (     )     ->						None:
      assert or_gate(0						,    0			) == 0
      assert or_gate(0						,    1			) == 1
      assert or_gate(1						,    0			) == 1
      assert or_gate(1						,    1			) == 1


if __name__ == "__main__":
  print(or_gate(0, 1))
  print(or_gate(1, 0))
  print(or_gate(0, 0))
  print(or_gate(1, 1))




