

'''simple docstring'''



from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization import (
    Adafactor,
    AdamW,
    get_constant_schedule,
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
		from fairscale.optim import OSS


__snake_case			:     Tuple										=      logging.get_logger(__name__)

__snake_case			:     Optional[int]										=      {
    'linear': get_linear_schedule_with_warmup,
    'cosine': get_cosine_schedule_with_warmup,
    'cosine_w_restarts': get_cosine_with_hard_restarts_schedule_with_warmup,
    'polynomial': get_polynomial_decay_schedule_with_warmup,
    'constant': get_constant_schedule,
    'constant_w_warmup': get_constant_schedule_with_warmup,
}



class lowerCamelCase			(  __SCREAMING_SNAKE_CASE						):



			'''simple docstring'''






			def __init__(				self	:						Tuple		,      lowerCAmelCase_	:						List[Any]=None		,      lowerCAmelCase_	:						int=None		,      *lowerCAmelCase_	:						List[Any]		,      **lowerCAmelCase_	:						Optional[int]			)      ->		Optional[Any]:





									'''simple docstring'''
									super().__init__(*_SCREAMING_SNAKE_CASE		,      **_SCREAMING_SNAKE_CASE			)

									if config is None:
															assert isinstance(self.model		,      _SCREAMING_SNAKE_CASE			), (
															    "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
															    f" {self.model.__class__}"
															)
															A__					:       List[Any]     =self.model.config
									else:
															A__					:       Union[str, Any]     =config

									A__					:       Dict     =data_args
									A__					:       Optional[int]     =self.config.tgt_vocab_size if isinstance(self.config		,      _SCREAMING_SNAKE_CASE			) else self.config.vocab_size

									if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
															assert self.config.pad_token_id is not None, (
															    "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
															    " calculation or doing label smoothing."
															)

									if self.config.pad_token_id is None and self.config.eos_token_id is not None:
															logger.warning(
															    f"The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for"
															    """ padding.."""			)

									if self.args.label_smoothing == 0:
															A__					:       List[str]     =torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id			)
									else:
															# dynamically import label_smoothed_nll_loss
															from utils import label_smoothed_nll_loss

															A__					:       Tuple     =label_smoothed_nll_loss


			def lowercase__						(				self	:						str		,      lowerCAmelCase_	:						int			)      ->		str:





									'''simple docstring'''

									if self.optimizer is None:
															A__					:       Dict     =["bias", "LayerNorm.weight"]
															A__					:       Optional[int]     =[
															    {
															        "params": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay			)],
															        "weight_decay": self.args.weight_decay,
															    },
															    {
															        "params": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay			)],
															        "weight_decay": 0.0,
															    },
															]
															A__					:       int     =Adafactor if self.args.adafactor else AdamW
															if self.args.adafactor:
																					A__					:       Any     =Adafactor
																					A__					:       int     ={"scale_parameter": False, "relative_step": False}
															else:
																					A__					:       Union[str, Any]     =AdamW
																					A__					:       List[Any]     ={
																					    "betas": (self.args.adam_betaa, self.args.adam_betaa),
																					    "eps": self.args.adam_epsilon,
																					}
															A__					:       Any     =self.args.learning_rate
															if self.sharded_ddp:
																					A__					:       List[str]     =OSS(
																					    params=_SCREAMING_SNAKE_CASE		,      optim=_SCREAMING_SNAKE_CASE		,      **_SCREAMING_SNAKE_CASE		,      )
															else:
																					A__					:       Any     =optimizer_cls(_SCREAMING_SNAKE_CASE		,      **_SCREAMING_SNAKE_CASE			)

									if self.lr_scheduler is None:
															A__					:       Optional[int]     =self._get_lr_scheduler(_SCREAMING_SNAKE_CASE			)
									else:  # ignoring --lr_scheduler
															logger.warning("""scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored."""			)


			def lowercase__						(				self	:						Dict		,      lowerCAmelCase_	:						Tuple			)      ->		str:





									'''simple docstring'''
									A__					:       List[str]     =arg_to_scheduler[self.args.lr_scheduler]
									if self.args.lr_scheduler == "constant":
															A__					:       Dict     =schedule_func(self.optimizer			)
									elif self.args.lr_scheduler == "constant_w_warmup":
															A__					:       Optional[Any]     =schedule_func(self.optimizer		,      num_warmup_steps=self.args.warmup_steps			)
									else:
															A__					:       str     =schedule_func(
															    self.optimizer		,      num_warmup_steps=self.args.warmup_steps		,      num_training_steps=_SCREAMING_SNAKE_CASE			)
									return scheduler


			def lowercase__						(				self	:						Optional[Any]			)      ->		Optional[torch.utils.data.Sampler]:





									'''simple docstring'''

									if isinstance(self.train_dataset		,      torch.utils.data.IterableDataset			):
															return None
									elif is_torch_tpu_available():
															return get_tpu_sampler(self.train_dataset			)
									else:
															if self.args.sortish_sampler:
																					self.train_dataset.make_sortish_sampler(
																					    self.args.per_device_train_batch_size		,      distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED)		,      )

															return (
															    RandomSampler(self.train_dataset			)
															    if self.args.local_rank == -1
															    else DistributedSampler(self.train_dataset			)
															)


			def lowercase__						(				self	:						Optional[Any]		,      lowerCAmelCase_	:						str		,      lowerCAmelCase_	:						int		,      lowerCAmelCase_	:						List[str]			)      ->		Optional[int]:





									'''simple docstring'''

									if self.args.label_smoothing == 0:
															if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
																					# force training to ignore pad token
																					A__					:       str     =model(**_SCREAMING_SNAKE_CASE		,      use_cache=_SCREAMING_SNAKE_CASE			)[0]
																					A__					:       Optional[Any]     =self.loss_fn(logits.view(-1		,      logits.shape[-1]			)		,      labels.view(-1			)			)
															else:
																					# compute usual loss via models
																					A__					:       Any     =model(**_SCREAMING_SNAKE_CASE		,      labels=_SCREAMING_SNAKE_CASE		,      use_cache=_SCREAMING_SNAKE_CASE			)[:2]
									else:
															# compute label smoothed loss
															A__					:       List[Any]     =model(**_SCREAMING_SNAKE_CASE		,      use_cache=_SCREAMING_SNAKE_CASE			)[0]
															A__					:       Tuple     =torch.nn.functional.log_softmax(_SCREAMING_SNAKE_CASE		,      dim=-1			)
															A__					:       str     =self.loss_fn(_SCREAMING_SNAKE_CASE		,      _SCREAMING_SNAKE_CASE		,      self.args.label_smoothing		,      ignore_index=self.config.pad_token_id			)
									return loss, logits


			def lowercase__						(				self	:						Optional[int]		,      lowerCAmelCase_	:						Dict		,      lowerCAmelCase_	:						Any			)      ->		Union[str, Any]:





									'''simple docstring'''
									A__					:       Union[str, Any]     =inputs.pop("""labels"""			)
									A__					:       List[Any]     =self._compute_loss(_SCREAMING_SNAKE_CASE		,      _SCREAMING_SNAKE_CASE		,      _SCREAMING_SNAKE_CASE			)
									return loss


			def lowercase__						(				self	:						str		,      lowerCAmelCase_	:						nn.Module		,      lowerCAmelCase_	:						Dict[str, Union[torch.Tensor, Any]]		,      lowerCAmelCase_	:						bool		,      lowerCAmelCase_	:						Optional[List[str]] = None		,      )      ->		Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:





									'''simple docstring'''
									A__					:       List[str]     =self._prepare_inputs(_SCREAMING_SNAKE_CASE			)

									A__					:       Any     ={
									    "max_length": self.data_args.val_max_target_length
									    if self.data_args is not None
									    else self.config.max_length,
									    "num_beams": self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
									}

									if self.args.predict_with_generate and not self.args.prediction_loss_only:
															A__					:       Tuple     =self.model.generate(
															    inputs["""input_ids"""]		,      attention_mask=inputs["""attention_mask"""]		,      **_SCREAMING_SNAKE_CASE		,      )
															# in case the batch is shorter than max length, the output should be padded
															if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
																					A__					:       List[Any]     =self._pad_tensors_to_max_len(_SCREAMING_SNAKE_CASE		,      gen_kwargs["""max_length"""]			)

									A__					:       List[str]     =inputs.pop("""labels"""			)
									with torch.no_grad():
															# compute loss on predict data
															A__					:       Dict     =self._compute_loss(_SCREAMING_SNAKE_CASE		,      _SCREAMING_SNAKE_CASE		,      _SCREAMING_SNAKE_CASE			)

									A__					:       str     =loss.mean().detach()
									if self.args.prediction_loss_only:
															return (loss, None, None)

									A__					:       Any     =generated_tokens if self.args.predict_with_generate else logits

									if labels.shape[-1] < gen_kwargs["max_length"]:
															A__					:       Optional[Any]     =self._pad_tensors_to_max_len(_SCREAMING_SNAKE_CASE		,      gen_kwargs["""max_length"""]			)

									return (loss, logits, labels)



			def lowercase__						(				self	:						Optional[int]		,      lowerCAmelCase_	:						Optional[Any]		,      lowerCAmelCase_	:						Dict			)      ->		List[str]:





									'''simple docstring'''
									A__					:       List[str]     =self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

									if pad_token_id is None:
															raise ValueError(
															    """Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be"""
															    f" padded to `max_length`={max_length}"			)

									A__					:       List[str]     =pad_token_id * torch.ones(
									    (tensor.shape[0], max_length)		,      dtype=tensor.dtype		,      device=tensor.device			)
									A__					:       Union[str, Any]     =tensor
									return padded_tensor




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



			'''simple docstring'''






			def __init__(				self	:						Dict		,      lowerCAmelCase_	:						Optional[int]		,      lowerCAmelCase_	:						List[str]=13		,      lowerCAmelCase_	:						Dict=7		,      lowerCAmelCase_	:						Union[str, Any]=True		,      lowerCAmelCase_	:						str=True		,      lowerCAmelCase_	:						int=False		,      lowerCAmelCase_	:						int=True		,      lowerCAmelCase_	:						int=99		,      lowerCAmelCase_	:						str=32		,      lowerCAmelCase_	:						str=5		,      lowerCAmelCase_	:						List[Any]=4		,      lowerCAmelCase_	:						Dict=37		,      lowerCAmelCase_	:						Optional[Any]="gelu"		,      lowerCAmelCase_	:						Tuple=0.1		,      lowerCAmelCase_	:						Dict=0.1		,      lowerCAmelCase_	:						str=5_12		,      lowerCAmelCase_	:						Dict=16		,      lowerCAmelCase_	:						Any=2		,      lowerCAmelCase_	:						Dict=0.02		,      lowerCAmelCase_	:						Optional[int]=3		,      lowerCAmelCase_	:						List[Any]=4		,      lowerCAmelCase_	:						Optional[int]=None		,      )      ->		List[str]:





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


			def lowercase__						(				self	:						int			)      ->		Tuple:





									'''simple docstring'''
									A__					:       Dict     =ids_tensor([self.batch_size, self.seq_length]		,      self.vocab_size			)

									A__					:       List[str]     =None
									if self.use_input_mask:
															A__					:       Union[str, Any]     =random_attention_mask([self.batch_size, self.seq_length]			)

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

									A__					:       str     =self.get_config()

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


			def lowercase__						(				self	:						Tuple			)      ->		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	:						Optional[Any]		,      lowerCAmelCase_	:						Dict		,      lowerCAmelCase_	:						Dict		,      lowerCAmelCase_	:						List[Any]		,      lowerCAmelCase_	:						Any		,      lowerCAmelCase_	:						Any		,      lowerCAmelCase_	:						List[Any]			)      ->		str:





									'''simple docstring'''
									A__					:       Union[str, Any]     =DistilBertModel(config=lowerCAmelCase_			)
									model.to(lowerCAmelCase_			)
									model.eval()
									A__					:       Dict     =model(lowerCAmelCase_		,      lowerCAmelCase_			)
									A__					:       List[str]     =model(lowerCAmelCase_			)
									self.parent.assertEqual(result.last_hidden_state.shape		,      (self.batch_size, self.seq_length, self.hidden_size)			)


			def lowercase__						(				self	:						Tuple		,      lowerCAmelCase_	:						Tuple		,      lowerCAmelCase_	:						Tuple		,      lowerCAmelCase_	:						Dict		,      lowerCAmelCase_	:						Optional[int]		,      lowerCAmelCase_	:						str		,      lowerCAmelCase_	:						Any			)      ->		List[str]:





									'''simple docstring'''
									A__					:       Union[str, Any]     =DistilBertForMaskedLM(config=lowerCAmelCase_			)
									model.to(lowerCAmelCase_			)
									model.eval()
									A__					:       Union[str, Any]     =model(lowerCAmelCase_		,      attention_mask=lowerCAmelCase_		,      labels=lowerCAmelCase_			)
									self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.vocab_size)			)


			def lowercase__						(				self	:						Dict		,      lowerCAmelCase_	:						List[str]		,      lowerCAmelCase_	:						str		,      lowerCAmelCase_	:						str		,      lowerCAmelCase_	:						List[Any]		,      lowerCAmelCase_	:						str		,      lowerCAmelCase_	:						Optional[Any]			)      ->		Dict:





									'''simple docstring'''
									A__					:       Dict     =DistilBertForQuestionAnswering(config=lowerCAmelCase_			)
									model.to(lowerCAmelCase_			)
									model.eval()
									A__					:       Tuple     =model(
									    lowerCAmelCase_		,      attention_mask=lowerCAmelCase_		,      start_positions=lowerCAmelCase_		,      end_positions=lowerCAmelCase_			)
									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	:						Optional[int]		,      lowerCAmelCase_	:						List[Any]		,      lowerCAmelCase_	:						Union[str, Any]		,      lowerCAmelCase_	:						Optional[Any]		,      lowerCAmelCase_	:						Optional[int]		,      lowerCAmelCase_	:						List[Any]		,      lowerCAmelCase_	:						Dict			)      ->		Tuple:





									'''simple docstring'''
									A__					:       List[Any]     =self.num_labels
									A__					:       Dict     =DistilBertForSequenceClassification(lowerCAmelCase_			)
									model.to(lowerCAmelCase_			)
									model.eval()
									A__					:       Union[str, Any]     =model(lowerCAmelCase_		,      attention_mask=lowerCAmelCase_		,      labels=lowerCAmelCase_			)
									self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_labels)			)


			def lowercase__						(				self	:						str		,      lowerCAmelCase_	:						List[str]		,      lowerCAmelCase_	:						Dict		,      lowerCAmelCase_	:						Tuple		,      lowerCAmelCase_	:						List[str]		,      lowerCAmelCase_	:						Tuple		,      lowerCAmelCase_	:						List[Any]			)      ->		int:





									'''simple docstring'''
									A__					:       Optional[int]     =self.num_labels
									A__					:       Union[str, Any]     =DistilBertForTokenClassification(config=lowerCAmelCase_			)
									model.to(lowerCAmelCase_			)
									model.eval()

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


			def lowercase__						(				self	:						Union[str, Any]		,      lowerCAmelCase_	:						Optional[int]		,      lowerCAmelCase_	:						int		,      lowerCAmelCase_	:						Tuple		,      lowerCAmelCase_	:						Dict		,      lowerCAmelCase_	:						List[str]		,      lowerCAmelCase_	:						Optional[int]			)      ->		Dict:





									'''simple docstring'''
									A__					:       int     =self.num_choices
									A__					:       Tuple     =DistilBertForMultipleChoice(config=lowerCAmelCase_			)
									model.to(lowerCAmelCase_			)
									model.eval()
									A__					:       int     =input_ids.unsqueeze(1			).expand(-1		,      self.num_choices		,      -1			).contiguous()
									A__					:       int     =input_mask.unsqueeze(1			).expand(-1		,      self.num_choices		,      -1			).contiguous()
									A__					:       Tuple     =model(
									    lowerCAmelCase_		,      attention_mask=lowerCAmelCase_		,      labels=lowerCAmelCase_		,      )
									self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_choices)			)



			def lowercase__						(				self	:						Tuple			)      ->		Optional[Any]:





									'''simple docstring'''
									A__					:       Optional[Any]     =self.prepare_config_and_inputs()
									((A__)   ,						(A__)   ,						(A__)   ,						(A__)   ,						(A__)   ,						(A__))					:       List[str]     =config_and_inputs
									A__					:       str     ={"""input_ids""": input_ids, """attention_mask""": input_mask}
									return config, inputs_dict



@require_torch
class lowerCamelCase			(  lowercase_	,						lowercase_	,						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 lowercase__						(				self	:						List[Any]			)      ->		Any:





									'''simple docstring'''
									A__					:       Any     =DistilBertModelTester(self			)
									A__					:       Tuple     =ConfigTester(self		,      config_class=lowerCAmelCase_		,      dim=37			)


			def lowercase__						(				self	:						List[str]			)      ->		Union[str, Any]:





									'''simple docstring'''
									self.config_tester.run_common_tests()


			def lowercase__						(				self	:						Union[str, Any]			)      ->		Any:





									'''simple docstring'''
									A__					:       int     =self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_model(*lowerCAmelCase_			)


			def lowercase__						(				self	:						Any			)      ->		Tuple:





									'''simple docstring'''
									A__					:       Tuple     =self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_masked_lm(*lowerCAmelCase_			)


			def lowercase__						(				self	:						List[Any]			)      ->		str:





									'''simple docstring'''
									A__					:       str     =self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_question_answering(*lowerCAmelCase_			)


			def lowercase__						(				self	:						List[Any]			)      ->		List[str]:





									'''simple docstring'''
									A__					:       Any     =self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowerCAmelCase_			)


			def lowercase__						(				self	:						Optional[int]			)      ->		Optional[int]:





									'''simple docstring'''
									A__					:       Dict     =self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_token_classification(*lowerCAmelCase_			)


			def lowercase__						(				self	:						List[str]			)      ->		Optional[int]:





									'''simple docstring'''
									A__					:       List[Any]     =self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowerCAmelCase_			)


			@slow
			def lowercase__						(				self	:						Union[str, Any]			)      ->		Dict:





									'''simple docstring'''
									for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															A__					:       Optional[int]     =DistilBertModel.from_pretrained(lowerCAmelCase_			)
															self.assertIsNotNone(lowerCAmelCase_			)



			@slow
			@require_torch_gpu
			def lowercase__						(				self	:						Tuple			)      ->		Any:





									'''simple docstring'''
									A__   ,						A__					:       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

															A__					:       Any     =True
															A__					:       Optional[Any]     =model_class(config=lowerCAmelCase_			)

															A__					:       Dict     =self._prepare_for_class(lowerCAmelCase_		,      lowerCAmelCase_			)
															A__					:       Optional[int]     =torch.jit.trace(
															    lowerCAmelCase_		,      (inputs_dict["""input_ids"""].to("""cpu"""			), inputs_dict["""attention_mask"""].to("""cpu"""			))			)

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



@require_torch
class lowerCamelCase			(  unittest.TestCase						):



			'''simple docstring'''






			@slow
			def lowercase__						(				self	:						List[Any]			)      ->		List[Any]:





									'''simple docstring'''
									A__					:       List[str]     =DistilBertModel.from_pretrained("""distilbert-base-uncased"""			)
									A__					:       Optional[Any]     =torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]			)
									A__					:       int     =torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]			)
									with torch.no_grad():
															A__					:       Any     =model(lowerCAmelCase_		,      attention_mask=lowerCAmelCase_			)[0]
									A__					:       Tuple     =torch.Size((1, 11, 7_68)			)
									self.assertEqual(output.shape		,      lowerCAmelCase_			)
									A__					:       Optional[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]		,      lowerCAmelCase_		,      atol=1e-4			)			)

