


'''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 lowerCAmelCase_		(				_UpperCamelCase ):

 '''simple docstring'''

 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=5_12	,   _UpperCAmelCase=16	,   _UpperCAmelCase=2	,   _UpperCAmelCase=0.02	,   _UpperCAmelCase=3	,   _UpperCAmelCase=4	,   _UpperCAmelCase=None	,   ):
       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 		UpperCamelCase__  (			self      ):
       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 		UpperCamelCase__  (			self      ):
       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 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       snake_case_ =   DistilBertModel(config=_UpperCAmelCase      )
       model.to(_UpperCAmelCase      )
       model.eval()
       snake_case_ =   model(_UpperCAmelCase	,   _UpperCAmelCase      )
       snake_case_ =   model(_UpperCAmelCase      )
       self.parent.assertEqual(result.last_hidden_state.shape	,   (self.batch_size, self.seq_length, self.hidden_size)      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       snake_case_ =   DistilBertForMaskedLM(config=_UpperCAmelCase      )
       model.to(_UpperCAmelCase      )
       model.eval()
       snake_case_ =   model(_UpperCAmelCase	,   attention_mask=_UpperCAmelCase	,   labels=_UpperCAmelCase      )
       self.parent.assertEqual(result.logits.shape	,   (self.batch_size, self.seq_length, self.vocab_size)      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       snake_case_ =   DistilBertForQuestionAnswering(config=_UpperCAmelCase      )
       model.to(_UpperCAmelCase      )
       model.eval()
       snake_case_ =   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 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       snake_case_ =   self.num_labels
       snake_case_ =   DistilBertForSequenceClassification(_UpperCAmelCase      )
       model.to(_UpperCAmelCase      )
       model.eval()
       snake_case_ =   model(_UpperCAmelCase	,   attention_mask=_UpperCAmelCase	,   labels=_UpperCAmelCase      )
       self.parent.assertEqual(result.logits.shape	,   (self.batch_size, self.num_labels)      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       snake_case_ =   self.num_labels
       snake_case_ =   DistilBertForTokenClassification(config=_UpperCAmelCase      )
       model.to(_UpperCAmelCase      )
       model.eval()

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



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       snake_case_ =   self.num_choices
       snake_case_ =   DistilBertForMultipleChoice(config=_UpperCAmelCase      )
       model.to(_UpperCAmelCase      )
       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(
           _UpperCAmelCase	,   attention_mask=_UpperCAmelCase	,   labels=_UpperCAmelCase	,   )
       self.parent.assertEqual(result.logits.shape	,   (self.batch_size, self.num_choices)      )





 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   self.prepare_config_and_inputs()
       (snake_case_) =   config_and_inputs
       snake_case_ =   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict







@require_torch
class lowerCAmelCase_		(				_UpperCamelCase      ,		_UpperCamelCase      ,		unittest.TestCase ):

 '''simple docstring'''

 __snake_case				=						(
     (
         DistilBertModel,
         DistilBertForMaskedLM,
         DistilBertForMultipleChoice,
         DistilBertForQuestionAnswering,
         DistilBertForSequenceClassification,
         DistilBertForTokenClassification,
     )
     if is_torch_available()
     else None
 )
 __snake_case				=						(
     {
         "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				=						True
 __snake_case				=						True
 __snake_case				=						True
 __snake_case				=						True



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   DistilBertModelTester(self      )
       snake_case_ =   ConfigTester(self	,   config_class=_UpperCAmelCase	,   dim=37      )



 def 		UpperCamelCase__  (			self      ):
       self.config_tester.run_common_tests()



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_model(*_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_masked_lm(*_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_question_answering(*_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_sequence_classification(*_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_token_classification(*_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_multiple_choice(*_UpperCAmelCase      )



 @slow
 def 		UpperCamelCase__  (			self      ):
       for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             snake_case_ =   DistilBertModel.from_pretrained(_UpperCAmelCase      )
             self.assertIsNotNone(_UpperCAmelCase      )





 @slow
 @require_torch_gpu
 def 		UpperCamelCase__  (			self      ):
       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=_UpperCAmelCase      )

             snake_case_ =   self._prepare_for_class(_UpperCAmelCase	,   _UpperCAmelCase      )
             snake_case_ =   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'''      )      )
                   snake_case_ =   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 lowerCAmelCase_		(				unittest.TestCase ):

 '''simple docstring'''

 @slow
 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   DistilBertModel.from_pretrained('''distilbert-base-uncased'''      )
       snake_case_ =   torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]      )
       snake_case_ =   torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]      )
       with torch.no_grad():
             snake_case_ =   model(_UpperCAmelCase	,   attention_mask=_UpperCAmelCase      )[0]
       snake_case_ =   torch.Size((1, 11, 7_68)      )
       self.assertEqual(output.shape	,   _UpperCAmelCase      )
       snake_case_ =   torch.tensor(
           [[[-0.1_639, 0.3_299, 0.1_648], [-0.1_746, 0.3_289, 0.1_710], [-0.1_884, 0.3_357, 0.1_810]]]      )

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



from ...utils import is_note_seq_available, is_transformers_available, is_torch_available
from ...utils import OptionalDependencyNotAvailable


try:
   if not (is_transformers_available() and is_torch_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
   from .notes_encoder import SpectrogramNotesEncoder
   from .continous_encoder import SpectrogramContEncoder
   from .pipeline_spectrogram_diffusion import (
       SpectrogramContEncoder,
       SpectrogramDiffusionPipeline,
       TaFilmDecoder,
   )

try:
   if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
else:
   from .midi_utils import MidiProcessor