
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxSeqaSeqConfigWithPast
from ...utils import logging


if TYPE_CHECKING:
     from ...feature_extraction_utils import FeatureExtractionMixin
     from ...tokenization_utils_base import PreTrainedTokenizerBase
     from ...utils import TensorType

lowerCamelCase__     :		int        = logging.get_logger(__name__)

lowerCamelCase__     :		int        = {
    """openai/whisper-base""": """https://huggingface.co/openai/whisper-base/resolve/main/config.json""",
}

# fmt: off
lowerCamelCase__     :		Dict        = [
    1, 2, 7, 8, 9, 1_0, 1_4, 2_5,
    2_6, 2_7, 2_8, 2_9, 3_1, 5_8, 5_9, 6_0, 6_1, 6_2,
    6_3, 9_0, 9_1, 9_2, 9_3, 3_5_7, 3_6_6, 4_3_8, 5_3_2, 6_8_5,
    7_0_5, 7_9_6, 9_3_0, 1_0_5_8, 1_2_2_0, 1_2_6_7, 1_2_7_9, 1_3_0_3, 1_3_4_3, 1_3_7_7,
    1_3_9_1, 1_6_3_5, 1_7_8_2, 1_8_7_5, 2_1_6_2, 2_3_6_1, 2_4_8_8, 3_4_6_7, 4_0_0_8, 4_2_1_1,
    4_6_0_0, 4_8_0_8, 5_2_9_9, 5_8_5_5, 6_3_2_9, 7_2_0_3, 9_6_0_9, 9_9_5_9, 1_0_5_6_3, 1_0_7_8_6,
    1_1_4_2_0, 1_1_7_0_9, 1_1_9_0_7, 1_3_1_6_3, 1_3_6_9_7, 1_3_7_0_0, 1_4_8_0_8, 1_5_3_0_6, 1_6_4_1_0, 1_6_7_9_1,
    1_7_9_9_2, 1_9_2_0_3, 1_9_5_1_0, 2_0_7_2_4, 2_2_3_0_5, 2_2_9_3_5, 2_7_0_0_7, 3_0_1_0_9, 3_0_4_2_0, 3_3_4_0_9,
    3_4_9_4_9, 4_0_2_8_3, 4_0_4_9_3, 4_0_5_4_9, 4_7_2_8_2, 4_9_1_4_6, 5_0_2_5_7, 5_0_3_5_9, 5_0_3_6_0, 5_0_3_6_1
]
lowerCamelCase__     :		str        = [
    1, 2, 7, 8, 9, 1_0, 1_4, 2_5,
    2_6, 2_7, 2_8, 2_9, 3_1, 5_8, 5_9, 6_0, 6_1, 6_2,
    6_3, 9_0, 9_1, 9_2, 9_3, 3_5_9, 5_0_3, 5_2_2, 5_4_2, 8_7_3,
    8_9_3, 9_0_2, 9_1_8, 9_2_2, 9_3_1, 1_3_5_0, 1_8_5_3, 1_9_8_2, 2_4_6_0, 2_6_2_7,
    3_2_4_6, 3_2_5_3, 3_2_6_8, 3_5_3_6, 3_8_4_6, 3_9_6_1, 4_1_8_3, 4_6_6_7, 6_5_8_5, 6_6_4_7,
    7_2_7_3, 9_0_6_1, 9_3_8_3, 1_0_4_2_8, 1_0_9_2_9, 1_1_9_3_8, 1_2_0_3_3, 1_2_3_3_1, 1_2_5_6_2, 1_3_7_9_3,
    1_4_1_5_7, 1_4_6_3_5, 1_5_2_6_5, 1_5_6_1_8, 1_6_5_5_3, 1_6_6_0_4, 1_8_3_6_2, 1_8_9_5_6, 2_0_0_7_5, 2_1_6_7_5,
    2_2_5_2_0, 2_6_1_3_0, 2_6_1_6_1, 2_6_4_3_5, 2_8_2_7_9, 2_9_4_6_4, 3_1_6_5_0, 3_2_3_0_2, 3_2_4_7_0, 3_6_8_6_5,
    4_2_8_6_3, 4_7_4_2_5, 4_9_8_7_0, 5_0_2_5_4, 5_0_2_5_8, 5_0_3_6_0, 5_0_3_6_1, 5_0_3_6_2
]



class    _snake_case (						UpperCAmelCase_				):
 __lowerCAmelCase	:  List[str]	   =      'whisper'
 __lowerCAmelCase	:  str	   =      ['past_key_values']
 __lowerCAmelCase	:  Optional[int]	   =      {'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}


 def __init__(       self       ,	SCREAMING_SNAKE_CASE_=5_18_65       ,	SCREAMING_SNAKE_CASE_=80       ,	SCREAMING_SNAKE_CASE_=6       ,	SCREAMING_SNAKE_CASE_=4       ,	SCREAMING_SNAKE_CASE_=6       ,	SCREAMING_SNAKE_CASE_=4       ,	SCREAMING_SNAKE_CASE_=15_36       ,	SCREAMING_SNAKE_CASE_=15_36       ,	SCREAMING_SNAKE_CASE_=0.0       ,	SCREAMING_SNAKE_CASE_=0.0       ,	SCREAMING_SNAKE_CASE_=5_02_57       ,	SCREAMING_SNAKE_CASE_=True       ,	SCREAMING_SNAKE_CASE_=True       ,	SCREAMING_SNAKE_CASE_="gelu"       ,	SCREAMING_SNAKE_CASE_=2_56       ,	SCREAMING_SNAKE_CASE_=0.0       ,	SCREAMING_SNAKE_CASE_=0.0       ,	SCREAMING_SNAKE_CASE_=0.0       ,	SCREAMING_SNAKE_CASE_=0.0_2       ,	SCREAMING_SNAKE_CASE_=False       ,	SCREAMING_SNAKE_CASE_=15_00       ,	SCREAMING_SNAKE_CASE_=4_48       ,	SCREAMING_SNAKE_CASE_=5_02_56       ,	SCREAMING_SNAKE_CASE_=5_02_56       ,	SCREAMING_SNAKE_CASE_=5_02_56       ,	SCREAMING_SNAKE_CASE_=None       ,	SCREAMING_SNAKE_CASE_=[2_20, 5_02_56]       ,	SCREAMING_SNAKE_CASE_=False       ,	SCREAMING_SNAKE_CASE_=2_56       ,	SCREAMING_SNAKE_CASE_=False       ,	SCREAMING_SNAKE_CASE_=0.0_5       ,	SCREAMING_SNAKE_CASE_=10       ,	SCREAMING_SNAKE_CASE_=2       ,	SCREAMING_SNAKE_CASE_=0.0       ,	SCREAMING_SNAKE_CASE_=10       ,	SCREAMING_SNAKE_CASE_=0       ,	SCREAMING_SNAKE_CASE_=7       ,	**SCREAMING_SNAKE_CASE_       ,	):


   '''simple docstring'''




   lowercase__	: str     		=   vocab_size
   lowercase__	: Tuple     		=   num_mel_bins
   lowercase__	: int     		=   d_model
   lowercase__	: str     		=   encoder_layers
   lowercase__	: str     		=   encoder_attention_heads
   lowercase__	: Any     		=   decoder_layers
   lowercase__	: Optional[int]     		=   decoder_attention_heads
   lowercase__	: List[str]     		=   decoder_ffn_dim
   lowercase__	: Dict     		=   encoder_ffn_dim
   lowercase__	: Dict     		=   dropout
   lowercase__	: Dict     		=   attention_dropout
   lowercase__	: List[str]     		=   activation_dropout
   lowercase__	: List[str]     		=   activation_function
   lowercase__	: Dict     		=   init_std
   lowercase__	: int     		=   encoder_layerdrop
   lowercase__	: List[str]     		=   decoder_layerdrop
   lowercase__	: List[str]     		=   use_cache
   lowercase__	: Tuple     		=   encoder_layers
   lowercase__	: List[Any]     		=   scale_embedding  # scale factor will be sqrt(d_model) if True
   lowercase__	: int     		=   max_source_positions
   lowercase__	: Tuple     		=   max_target_positions

   # Audio Classification-specific parameters. Feel free to ignore for other classes.
   lowercase__	: Tuple     		=   classifier_proj_size
   lowercase__	: Optional[int]     		=   use_weighted_layer_sum

   # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
   lowercase__	: Optional[int]     		=   apply_spec_augment
   lowercase__	: Any     		=   mask_time_prob
   lowercase__	: str     		=   mask_time_length
   lowercase__	: int     		=   mask_time_min_masks
   lowercase__	: List[Any]     		=   mask_feature_prob
   lowercase__	: List[Any]     		=   mask_feature_length
   lowercase__	: Optional[int]     		=   mask_feature_min_masks

   lowercase__	: Dict     		=   median_filter_width

   super().__init__(
       pad_token_id=SCREAMING_SNAKE_CASE_       ,	bos_token_id=SCREAMING_SNAKE_CASE_       ,	eos_token_id=SCREAMING_SNAKE_CASE_       ,	is_encoder_decoder=SCREAMING_SNAKE_CASE_       ,	decoder_start_token_id=SCREAMING_SNAKE_CASE_       ,	suppress_tokens=SCREAMING_SNAKE_CASE_       ,	begin_suppress_tokens=SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_       ,	)



class    _snake_case (						UpperCAmelCase_				):


 @property
 def   lowercase__  (       self):


   '''simple docstring'''




   lowercase__	: int     		=   OrderedDict(
       [
           ("""input_features""", {0: """batch""", 1: """feature_size""", 2: """encoder_sequence"""}),
       ])
   if self.use_past:
     lowercase__	: List[str]     		=   {0: """batch"""}
   else:
     lowercase__	: int     		=   {0: """batch""", 1: """decoder_sequence"""}

   if self.use_past:
     self.fill_with_past_key_values_(SCREAMING_SNAKE_CASE_       ,	direction="""inputs""")

   return common_inputs


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_ = -1       ,	SCREAMING_SNAKE_CASE_ = -1       ,	SCREAMING_SNAKE_CASE_ = False       ,	SCREAMING_SNAKE_CASE_ = None       ,	SCREAMING_SNAKE_CASE_ = 2_20_50       ,	SCREAMING_SNAKE_CASE_ = 5.0       ,	SCREAMING_SNAKE_CASE_ = 2_20       ,	):


   '''simple docstring'''




   lowercase__	: str     		=   OrderedDict()
   lowercase__	: List[Any]     		=   OnnxConfig.generate_dummy_inputs(
       self       ,	preprocessor=preprocessor.feature_extractor       ,	batch_size=SCREAMING_SNAKE_CASE_       ,	framework=SCREAMING_SNAKE_CASE_       ,	sampling_rate=SCREAMING_SNAKE_CASE_       ,	time_duration=SCREAMING_SNAKE_CASE_       ,	frequency=SCREAMING_SNAKE_CASE_       ,	)
   lowercase__	: Dict     		=   encoder_inputs["""input_features"""].shape[2]
   lowercase__	: List[Any]     		=   encoder_sequence_length // 2 if self.use_past else seq_length

   lowercase__	: Optional[Any]     		=   super().generate_dummy_inputs(
       preprocessor.tokenizer       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)

   lowercase__	: List[str]     		=   encoder_inputs.pop("""input_features""")
   lowercase__	: Optional[Any]     		=   decoder_inputs.pop("""decoder_input_ids""")

   if "past_key_values" in decoder_inputs:
     lowercase__	: Dict     		=   decoder_inputs.pop("""past_key_values""")

   return dummy_inputs



 @property
 def   lowercase__  (       self):


   '''simple docstring'''




   return 1E-3








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






	def __init__(self				: Dict      ,	snake_case				: str      ,	snake_case				: str=13      ,	snake_case				: Union[str, Any]=7      ,	snake_case				: int=True      ,	snake_case				: Any=True      ,	snake_case				: str=False      ,	snake_case				: Optional[Any]=True      ,	snake_case				: Optional[Any]=99      ,	snake_case				: Dict=32      ,	snake_case				: Union[str, Any]=5      ,	snake_case				: List[str]=4      ,	snake_case				: Optional[int]=37      ,	snake_case				: Optional[int]="gelu"      ,	snake_case				: Optional[Any]=0.1      ,	snake_case				: Optional[int]=0.1      ,	snake_case				: List[Any]=512      ,	snake_case				: List[Any]=16      ,	snake_case				: Optional[int]=2      ,	snake_case				: Tuple=0.02      ,	snake_case				: Union[str, Any]=3      ,	snake_case				: Any=4      ,	snake_case				: Any=None      ,	) ->			List[Any]:
							_lowercase       :	Dict           =   parent
							_lowercase       :	int           =   batch_size
							_lowercase       :	Optional[Any]           =   seq_length
							_lowercase       :	int           =   is_training
							_lowercase       :	Dict           =   use_input_mask
							_lowercase       :	Union[str, Any]           =   use_token_type_ids
							_lowercase       :	Tuple           =   use_labels
							_lowercase       :	int           =   vocab_size
							_lowercase       :	Union[str, Any]           =   hidden_size
							_lowercase       :	int           =   num_hidden_layers
							_lowercase       :	Optional[Any]           =   num_attention_heads
							_lowercase       :	Tuple           =   intermediate_size
							_lowercase       :	Dict           =   hidden_act
							_lowercase       :	int           =   hidden_dropout_prob
							_lowercase       :	int           =   attention_probs_dropout_prob
							_lowercase       :	Tuple           =   max_position_embeddings
							_lowercase       :	str           =   type_vocab_size
							_lowercase       :	int           =   type_sequence_label_size
							_lowercase       :	Any           =   initializer_range
							_lowercase       :	Optional[Any]           =   num_labels
							_lowercase       :	Optional[Any]           =   num_choices
							_lowercase       :	str           =   scope






	def       _a(self				: int    ) ->			Dict:
							_lowercase       :	str           =   ids_tensor([self.batch_size, self.seq_length]      ,	self.vocab_size    )

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

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

							_lowercase       :	Dict           =   self.get_config()

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






	def       _a(self				: Any    ) ->			Any:
							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				: int      ,	snake_case				: Optional[Any]      ,	snake_case				: Tuple      ,	snake_case				: List[str]      ,	snake_case				: Tuple      ,	snake_case				: Any      ,	snake_case				: Dict    ) ->			Optional[int]:
							_lowercase       :	Optional[int]           =   DistilBertModel(config=snake_case    )
							model.to(snake_case    )
							model.eval()
							_lowercase       :	List[Any]           =   model(snake_case      ,	snake_case    )
							_lowercase       :	int           =   model(snake_case    )
							self.parent.assertEqual(result.last_hidden_state.shape      ,	(self.batch_size, self.seq_length, self.hidden_size)    )






	def       _a(self				: int      ,	snake_case				: Optional[Any]      ,	snake_case				: Tuple      ,	snake_case				: Optional[Any]      ,	snake_case				: List[str]      ,	snake_case				: Optional[Any]      ,	snake_case				: Optional[Any]    ) ->			Dict:
							_lowercase       :	Optional[int]           =   DistilBertForMaskedLM(config=snake_case    )
							model.to(snake_case    )
							model.eval()
							_lowercase       :	int           =   model(snake_case      ,	attention_mask=snake_case      ,	labels=snake_case    )
							self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.seq_length, self.vocab_size)    )






	def       _a(self				: Tuple      ,	snake_case				: List[str]      ,	snake_case				: Any      ,	snake_case				: List[str]      ,	snake_case				: Dict      ,	snake_case				: str      ,	snake_case				: str    ) ->			Any:
							_lowercase       :	Dict           =   DistilBertForQuestionAnswering(config=snake_case    )
							model.to(snake_case    )
							model.eval()
							_lowercase       :	List[str]           =   model(
							    snake_case      ,	attention_mask=snake_case      ,	start_positions=snake_case      ,	end_positions=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				: Union[str, Any]      ,	snake_case				: str      ,	snake_case				: Dict      ,	snake_case				: Dict      ,	snake_case				: Union[str, Any]      ,	snake_case				: Optional[int]      ,	snake_case				: Dict    ) ->			Dict:
							_lowercase       :	str           =   self.num_labels
							_lowercase       :	Any           =   DistilBertForSequenceClassification(snake_case    )
							model.to(snake_case    )
							model.eval()
							_lowercase       :	Optional[Any]           =   model(snake_case      ,	attention_mask=snake_case      ,	labels=snake_case    )
							self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.num_labels)    )






	def       _a(self				: int      ,	snake_case				: Union[str, Any]      ,	snake_case				: Optional[Any]      ,	snake_case				: int      ,	snake_case				: Optional[int]      ,	snake_case				: Optional[int]      ,	snake_case				: str    ) ->			str:
							_lowercase       :	str           =   self.num_labels
							_lowercase       :	List[str]           =   DistilBertForTokenClassification(config=snake_case    )
							model.to(snake_case    )
							model.eval()

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






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






	def       _a(self				: List[str]    ) ->			List[str]:
							_lowercase       :	Union[str, Any]           =   self.prepare_config_and_inputs()
							((_lowercase)      ,    (_lowercase)      ,    (_lowercase)      ,    (_lowercase)      ,    (_lowercase)      ,    (_lowercase))       :	Union[str, Any]           =   config_and_inputs
							_lowercase       :	Union[str, Any]           =   {"input_ids": input_ids, "attention_mask": input_mask}
							return config, inputs_dict




@require_torch
class   __lowercase						(						__snake_case   , __snake_case   , unittest.TestCase       ):
	_A       =   (
	    (
	        DistilBertModel,
	        DistilBertForMaskedLM,
	        DistilBertForMultipleChoice,
	        DistilBertForQuestionAnswering,
	        DistilBertForSequenceClassification,
	        DistilBertForTokenClassification,
	    )
	    if is_torch_available()
	    else None
	)
	_A       =   (
	    {
	        "feature-extraction": DistilBertModel,
	        "fill-mask": DistilBertForMaskedLM,
	        "question-answering": DistilBertForQuestionAnswering,
	        "text-classification": DistilBertForSequenceClassification,
	        "token-classification": DistilBertForTokenClassification,
	        "zero-shot": DistilBertForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	_A       =   True
	_A       =   True
	_A       =   True
	_A       =   True






	def       _a(self				: Dict    ) ->			List[Any]:
							_lowercase       :	Optional[Any]           =   DistilBertModelTester(self    )
							_lowercase       :	str           =   ConfigTester(self      ,	config_class=snake_case      ,	dim=37    )






	def       _a(self				: int    ) ->			List[str]:
							self.config_tester.run_common_tests()






	def       _a(self				: Optional[Any]    ) ->			Optional[int]:
							_lowercase       :	Tuple           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_model(*snake_case    )






	def       _a(self				: Any    ) ->			int:
							_lowercase       :	str           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_masked_lm(*snake_case    )






	def       _a(self				: Dict    ) ->			List[Any]:
							_lowercase       :	List[str]           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_question_answering(*snake_case    )






	def       _a(self				: str    ) ->			Tuple:
							_lowercase       :	Tuple           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_sequence_classification(*snake_case    )






	def       _a(self				: Any    ) ->			List[Any]:
							_lowercase       :	Optional[Any]           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_token_classification(*snake_case    )






	def       _a(self				: Optional[int]    ) ->			Optional[int]:
							_lowercase       :	Any           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_multiple_choice(*snake_case    )






	@slow
	def       _a(self				: Optional[Any]    ) ->			Dict:
							for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													_lowercase       :	Tuple           =   DistilBertModel.from_pretrained(snake_case    )
													self.assertIsNotNone(snake_case    )






	@slow
	@require_torch_gpu
	def       _a(self				: Optional[int]    ) ->			Optional[int]:
							_lowercase      ,    _lowercase       :	Optional[int]           =   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

													_lowercase       :	str           =   True
													_lowercase       :	Tuple           =   model_class(config=snake_case    )

													_lowercase       :	str           =   self._prepare_for_class(snake_case      ,	snake_case    )
													_lowercase       :	Optional[int]           =   torch.jit.trace(
													    snake_case      ,	(inputs_dict["input_ids"].to("cpu"    ), inputs_dict["attention_mask"].to("cpu"    ))    )

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




@require_torch
class   __lowercase						(						unittest.TestCase       ):






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




