



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      UpperCamelCase    (						lowerCamelCase__     ):



	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__=512			,	UpperCAmelCase__=16			,	UpperCAmelCase__=2			,	UpperCAmelCase__=0.02			,	UpperCAmelCase__=3			,	UpperCAmelCase__=4			,	UpperCAmelCase__=None			,	):
							A__							=       parent
							A__							=       batch_size
							A__							=       seq_length
							A__							=       is_training
							A__							=       use_input_mask
							A__							=       use_token_type_ids
							A__							=       use_labels
							A__							=       vocab_size
							A__							=       hidden_size
							A__							=       num_hidden_layers
							A__							=       num_attention_heads
							A__							=       intermediate_size
							A__							=       hidden_act
							A__							=       hidden_dropout_prob
							A__							=       attention_probs_dropout_prob
							A__							=       max_position_embeddings
							A__							=       type_vocab_size
							A__							=       type_sequence_label_size
							A__							=       initializer_range
							A__							=       num_labels
							A__							=       num_choices
							A__							=       scope



	def 			__A							(       self		):
							A__							=       ids_tensor([self.batch_size, self.seq_length]			,	self.vocab_size		)

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

							A__							=       None
							A__							=       None
							A__							=       None
							if self.use_labels:
													A__							=       ids_tensor([self.batch_size]			,	self.type_sequence_label_size		)
													A__							=       ids_tensor([self.batch_size, self.seq_length]			,	self.num_labels		)
													A__							=       ids_tensor([self.batch_size]			,	self.num_choices		)

							A__							=       self.get_config()

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



	def 			__A							(       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 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       DistilBertModel(config=UpperCAmelCase__		)
							model.to(UpperCAmelCase__		)
							model.eval()
							A__							=       model(UpperCAmelCase__			,	UpperCAmelCase__		)
							A__							=       model(UpperCAmelCase__		)
							self.parent.assertEqual(result.last_hidden_state.shape			,	(self.batch_size, self.seq_length, self.hidden_size)		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       DistilBertForMaskedLM(config=UpperCAmelCase__		)
							model.to(UpperCAmelCase__		)
							model.eval()
							A__							=       model(UpperCAmelCase__			,	attention_mask=UpperCAmelCase__			,	labels=UpperCAmelCase__		)
							self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.seq_length, self.vocab_size)		)



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



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       self.num_labels
							A__							=       DistilBertForTokenClassification(config=UpperCAmelCase__		)
							model.to(UpperCAmelCase__		)
							model.eval()

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



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       self.num_choices
							A__							=       DistilBertForMultipleChoice(config=UpperCAmelCase__		)
							model.to(UpperCAmelCase__		)
							model.eval()
							A__							=       input_ids.unsqueeze(1		).expand(-1			,	self.num_choices			,	-1		).contiguous()
							A__							=       input_mask.unsqueeze(1		).expand(-1			,	self.num_choices			,	-1		).contiguous()
							A__							=       model(
							    UpperCAmelCase__			,	attention_mask=UpperCAmelCase__			,	labels=UpperCAmelCase__			,	)
							self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.num_choices)		)



	def 			__A							(       self		):
							A__							=       self.prepare_config_and_inputs()
							(A__)							=       config_and_inputs
							A__							=       {"input_ids": input_ids, "attention_mask": input_mask}
							return config, inputs_dict






@require_torch
class      UpperCamelCase    (						lowerCamelCase__       ,    lowerCamelCase__       ,    unittest.TestCase     ):
	lowerCAmelCase	:   List[Any]				=	(
	    (
	        DistilBertModel,
	        DistilBertForMaskedLM,
	        DistilBertForMultipleChoice,
	        DistilBertForQuestionAnswering,
	        DistilBertForSequenceClassification,
	        DistilBertForTokenClassification,
	    )
	    if is_torch_available()
	    else None
	)
	lowerCAmelCase	:   Optional[Any]				=	(
	    {
	        """feature-extraction""": DistilBertModel,
	        """fill-mask""": DistilBertForMaskedLM,
	        """question-answering""": DistilBertForQuestionAnswering,
	        """text-classification""": DistilBertForSequenceClassification,
	        """token-classification""": DistilBertForTokenClassification,
	        """zero-shot""": DistilBertForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	lowerCAmelCase	:   Any				=	True
	lowerCAmelCase	:   Dict				=	True
	lowerCAmelCase	:   List[str]				=	True
	lowerCAmelCase	:   Optional[int]				=	True



	def 			__A							(       self		):
							A__							=       DistilBertModelTester(self		)
							A__							=       ConfigTester(self			,	config_class=UpperCAmelCase__			,	dim=37		)



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



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_model(*UpperCAmelCase__		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_masked_lm(*UpperCAmelCase__		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_question_answering(*UpperCAmelCase__		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_sequence_classification(*UpperCAmelCase__		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_token_classification(*UpperCAmelCase__		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_multiple_choice(*UpperCAmelCase__		)



	@slow
	def 			__A							(       self		):




							for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													A__							=       DistilBertModel.from_pretrained(UpperCAmelCase__		)
													self.assertIsNotNone(UpperCAmelCase__		)



	@slow
	@require_torch_gpu
	def 			__A							(       self		):
							A__							=       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__							=       True
													A__							=       model_class(config=UpperCAmelCase__		)

													A__							=       self._prepare_for_class(UpperCAmelCase__			,	UpperCAmelCase__		)
													A__							=       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"		)		)
																			A__							=       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      UpperCamelCase    (						unittest.TestCase     ):



	@slow
	def 			__A							(       self		):
							A__							=       DistilBertModel.from_pretrained("distilbert-base-uncased"		)
							A__							=       torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]		)
							A__							=       torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
							with torch.no_grad():
													A__							=       model(UpperCAmelCase__			,	attention_mask=UpperCAmelCase__		)[0]
							A__							=       torch.Size((1, 11, 768)		)
							self.assertEqual(output.shape			,	UpperCAmelCase__		)
							A__							=       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		)		)










import math
from enum import Enum
from typing import Optional, Union

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR

from .utils import logging


UpperCAmelCase_							:      Optional[Any]      =  logging.get_logger(__name__)






class      UpperCamelCase    (						_UpperCAmelCase     ):
	lowerCAmelCase	:   List[str]				=	"""linear"""
	lowerCAmelCase	:   int				=	"""cosine"""
	lowerCAmelCase	:   Dict				=	"""cosine_with_restarts"""
	lowerCAmelCase	:   Optional[Any]				=	"""polynomial"""
	lowerCAmelCase	:   Dict				=	"""constant"""
	lowerCAmelCase	:   Any				=	"""constant_with_warmup"""
	lowerCAmelCase	:   Union[str, Any]				=	"""piecewise_constant"""





def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int = -1    )->       Dict:





						"""simple docstring"""





						return LambdaLR(_A			,							lambda _A    : 1			,							last_epoch=_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int			,							_A			:				int = -1    )->       Optional[Any]:





						"""simple docstring"""





						def lr_lambda(_A			:				int    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1.0			,							_A    )    )
												return 1.0

						return LambdaLR(_A			,							_A			,							last_epoch=_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				str			,							_A			:				int = -1    )->       Dict:





						"""simple docstring"""





						A__							=       {}
						A__							=       step_rules.split(","    )
						for rule_str in rule_list[:-1]:
												A__						,       A__							=       rule_str.split(":"    )
												A__							=       int(_A    )
												A__							=       float(_A    )
												A__							=       value
						A__							=       float(rule_list[-1]    )

						def create_rules_function(_A			:				Any			,							_A			:				Optional[int]    ):
												def rule_func(_A			:				int    ) -> float:
																		A__							=       sorted(rules_dict.keys()    )
																		for i, sorted_step in enumerate(_A    ):
																								if steps < sorted_step:
																														return rules_dict[sorted_steps[i]]
																		return last_lr_multiple

												return rule_func

						A__							=       create_rules_function(_A			,							_A    )

						return LambdaLR(_A			,							_A			,							last_epoch=_A    )
def  UpperCamelCase						(       _A			:				Any			,							_A			:				Union[str, Any]			,							_A			:				str			,							_A			:				str=-1    )->       Tuple:





						"""simple docstring"""





						def lr_lambda(_A			:				int    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												return max(
												    0.0			,							float(num_training_steps - current_step    ) / float(max(1			,							num_training_steps - num_warmup_steps    )    )    )

						return LambdaLR(_A			,							_A			,							_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int			,							_A			:				int			,							_A			:				float = 0.5			,							_A			:				int = -1    )->       Any:





						"""simple docstring"""





						def lr_lambda(_A			:				Tuple    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												A__							=       float(current_step - num_warmup_steps    ) / float(max(1			,							num_training_steps - num_warmup_steps    )    )
												return max(0.0			,							0.5 * (1.0 + math.cos(math.pi * float(_A    ) * 2.0 * progress    ))    )

						return LambdaLR(_A			,							_A			,							_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int			,							_A			:				int			,							_A			:				int = 1			,							_A			:				int = -1    )->       Any:





						"""simple docstring"""





						def lr_lambda(_A			:				Tuple    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												A__							=       float(current_step - num_warmup_steps    ) / float(max(1			,							num_training_steps - num_warmup_steps    )    )
												if progress >= 1.0:
																		return 0.0
												return max(0.0			,							0.5 * (1.0 + math.cos(math.pi * ((float(_A    ) * progress) % 1.0)    ))    )

						return LambdaLR(_A			,							_A			,							_A    )
def  UpperCamelCase						(       _A			:				Union[str, Any]			,							_A			:				Union[str, Any]			,							_A			:				List[str]			,							_A			:				Tuple=1E-7			,							_A			:				Dict=1.0			,							_A			:				Union[str, Any]=-1    )->       Any:





						"""simple docstring"""





						A__							=       optimizer.defaults["lr"]
						if not (lr_init > lr_end):
												raise ValueError(f"""lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})"""    )

						def lr_lambda(_A			:				int    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												elif current_step > num_training_steps:
																		return lr_end / lr_init  # as LambdaLR multiplies by lr_init
												else:
																		A__							=       lr_init - lr_end
																		A__							=       num_training_steps - num_warmup_steps
																		A__							=       1 - (current_step - num_warmup_steps) / decay_steps
																		A__							=       lr_range * pct_remaining**power + lr_end
																		return decay / lr_init  # as LambdaLR multiplies by lr_init

						return LambdaLR(_A			,							_A			,							_A    )


UpperCAmelCase_							:      Any      =  {
    SchedulerType.LINEAR: get_linear_schedule_with_warmup,
    SchedulerType.COSINE: get_cosine_schedule_with_warmup,
    SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
    SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
    SchedulerType.CONSTANT: get_constant_schedule,
    SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
    SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule,
}
def  UpperCamelCase						(       _A			:				Union[str, SchedulerType]			,							_A			:				Optimizer			,							_A			:				Optional[str] = None			,							_A			:				Optional[int] = None			,							_A			:				Optional[int] = None			,							_A			:				int = 1			,							_A			:				float = 1.0			,							_A			:				int = -1			,							)->       Union[str, Any]:





						"""simple docstring"""





						A__							=       SchedulerType(_A    )
						A__							=       TYPE_TO_SCHEDULER_FUNCTION[name]
						if name == SchedulerType.CONSTANT:
												return schedule_func(_A			,							last_epoch=_A    )

						if name == SchedulerType.PIECEWISE_CONSTANT:
												return schedule_func(_A			,							step_rules=_A			,							last_epoch=_A    )

						# All other schedulers require `num_warmup_steps`
						if num_warmup_steps is None:
												raise ValueError(f"""{name} requires `num_warmup_steps`, please provide that argument."""    )

						if name == SchedulerType.CONSTANT_WITH_WARMUP:
												return schedule_func(_A			,							num_warmup_steps=_A			,							last_epoch=_A    )

						# All other schedulers require `num_training_steps`
						if num_training_steps is None:
												raise ValueError(f"""{name} requires `num_training_steps`, please provide that argument."""    )

						if name == SchedulerType.COSINE_WITH_RESTARTS:
												return schedule_func(
												    _A			,							num_warmup_steps=_A			,							num_training_steps=_A			,							num_cycles=_A			,							last_epoch=_A			,							)

						if name == SchedulerType.POLYNOMIAL:
												return schedule_func(
												    _A			,							num_warmup_steps=_A			,							num_training_steps=_A			,							power=_A			,							last_epoch=_A			,							)

						return schedule_func(
						    _A			,							num_warmup_steps=_A			,							num_training_steps=_A			,							last_epoch=_A    )





