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

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


								'''simple docstring'''

								_UpperCamelCase						    =	parent
								_UpperCamelCase						    =	batch_size
								_UpperCamelCase						    =	seq_length
								_UpperCamelCase						    =	is_training
								_UpperCamelCase						    =	use_input_mask
								_UpperCamelCase						    =	use_token_type_ids
								_UpperCamelCase						    =	use_labels
								_UpperCamelCase						    =	vocab_size
								_UpperCamelCase						    =	hidden_size
								_UpperCamelCase						    =	num_hidden_layers
								_UpperCamelCase						    =	num_attention_heads
								_UpperCamelCase						    =	intermediate_size
								_UpperCamelCase						    =	hidden_act
								_UpperCamelCase						    =	hidden_dropout_prob
								_UpperCamelCase						    =	attention_probs_dropout_prob
								_UpperCamelCase						    =	max_position_embeddings
								_UpperCamelCase						    =	type_vocab_size
								_UpperCamelCase						    =	type_sequence_label_size
								_UpperCamelCase						    =	initializer_range
								_UpperCamelCase						    =	num_labels
								_UpperCamelCase						    =	num_choices
								_UpperCamelCase						    =	scope

						def  A_				(			self				)					->      Tuple:


								'''simple docstring'''

								_UpperCamelCase						    =	ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size				)

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

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

								_UpperCamelCase						    =	self.get_config()

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

						def  A_				(			self				)					->      Optional[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  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	DistilBertModel(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	model(a       ,							a				)
								_UpperCamelCase						    =	model(a				)
								self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)				)

						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Optional[int]:


								'''simple docstring'''

								_UpperCamelCase						    =	DistilBertForMaskedLM(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	model(a       ,							attention_mask=a       ,							labels=a				)
								self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.vocab_size)				)

						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Dict:


								'''simple docstring'''

								_UpperCamelCase						    =	DistilBertForQuestionAnswering(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	model(
								    a       ,							attention_mask=a       ,							start_positions=a       ,							end_positions=a				)
								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       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Optional[int]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.num_labels
								_UpperCamelCase						    =	DistilBertForSequenceClassification(a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	model(a       ,							attention_mask=a       ,							labels=a				)
								self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.num_labels)				)

						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.num_labels
								_UpperCamelCase						    =	DistilBertForTokenClassification(config=a				)
								model.to(a				)
								model.eval()

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

						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Any:


								'''simple docstring'''

								_UpperCamelCase						    =	self.num_choices
								_UpperCamelCase						    =	DistilBertForMultipleChoice(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	input_ids.unsqueeze(1				).expand(-1       ,							self.num_choices       ,							-1				).contiguous()
								_UpperCamelCase						    =	input_mask.unsqueeze(1				).expand(-1       ,							self.num_choices       ,							-1				).contiguous()
								_UpperCamelCase						    =	model(
								    a       ,							attention_mask=a       ,							labels=a       ,							)
								self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.num_choices)				)







						def  A_				(			self				)					->      str:


								'''simple docstring'''

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







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

						def  A_				(			self				)					->      Dict:


								'''simple docstring'''

								_UpperCamelCase						    =	DistilBertModelTester(self				)
								_UpperCamelCase						    =	ConfigTester(self       ,							config_class=a       ,							dim=37				)

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


								'''simple docstring'''

								self.config_tester.run_common_tests()

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


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_model(*a				)

						def  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_masked_lm(*a				)

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


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_question_answering(*a				)

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


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_sequence_classification(*a				)

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


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_token_classification(*a				)

						def  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_multiple_choice(*a				)

						@slow
						def  A_				(			self				)					->      Dict:


								'''simple docstring'''

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







						@slow
						@require_torch_gpu
						def  A_				(			self				)					->      Union[str, Any]:


								'''simple docstring'''

								_UpperCamelCase    ,	_UpperCamelCase						    =	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

										_UpperCamelCase						    =	True
										_UpperCamelCase						    =	model_class(config=a				)

										_UpperCamelCase						    =	self._prepare_for_class(a       ,							a				)
										_UpperCamelCase						    =	torch.jit.trace(
										    a       ,							(inputs_dict["""input_ids"""].to("""cpu"""				), inputs_dict["""attention_mask"""].to("""cpu"""				))				)

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







@require_torch
class        lowerCAmelCase__      (   unittest.TestCase	):

						@slow
						def  A_				(			self				)					->      Optional[Any]:


								'''simple docstring'''

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





import argparse

import torch

from transformers import YosoConfig, YosoForMaskedLM


def 							__A(lowerCAmelCase )		->				List[str]:

		"""simple docstring"""






		if "model" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""model."""						,  """""" )
		if "norm1" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""norm1"""						,  """attention.output.LayerNorm""" )
		if "norm2" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""norm2"""						,  """output.LayerNorm""" )
		if "norm" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""norm"""						,  """LayerNorm""" )
		if "transformer" in orig_key:
				_UpperCamelCase						    =	orig_key.split(""".""" )[0].split("""_""" )[-1]
				_UpperCamelCase						    =	orig_key.replace(F'transformer_{layer_num}'						,  F'encoder.layer.{layer_num}' )
		if "mha.attn" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""mha.attn"""						,  """attention.self""" )
		if "mha" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""mha"""						,  """attention""" )
		if "W_q" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""W_q"""						,  """self.query""" )
		if "W_k" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""W_k"""						,  """self.key""" )
		if "W_v" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""W_v"""						,  """self.value""" )
		if "ff1" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""ff1"""						,  """intermediate.dense""" )
		if "ff2" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""ff2"""						,  """output.dense""" )
		if "ff" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""ff"""						,  """output.dense""" )
		if "mlm_class" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""mlm.mlm_class"""						,  """cls.predictions.decoder""" )
		if "mlm" in orig_key:
				_UpperCamelCase						    =	orig_key.replace("""mlm"""						,  """cls.predictions.transform""" )
		if "cls" not in orig_key:
				_UpperCamelCase						    =	"""yoso.""" + orig_key

		return orig_key


def 							__A(lowerCAmelCase						,  lowerCAmelCase )		->				List[str]:

		"""simple docstring"""






		for key in orig_state_dict.copy().keys():
				_UpperCamelCase						    =	orig_state_dict.pop(lowerCAmelCase )

				if ("pooler" in key) or ("sen_class" in key):
						continue
				else:
						_UpperCamelCase						    =	val

		_UpperCamelCase						    =	orig_state_dict["""cls.predictions.decoder.bias"""]
		_UpperCamelCase						    =	torch.arange(lowerCAmelCase ).expand((1, -1) ) + 2

		return orig_state_dict


def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				Optional[int]:

		"""simple docstring"""






		_UpperCamelCase						    =	torch.load(lowerCAmelCase						,  map_location="""cpu""" )["""model_state_dict"""]
		_UpperCamelCase						    =	YosoConfig.from_json_file(lowerCAmelCase )
		_UpperCamelCase						    =	YosoForMaskedLM(lowerCAmelCase )

		_UpperCamelCase						    =	convert_checkpoint_helper(config.max_position_embeddings						,  lowerCAmelCase )

		print(model.load_state_dict(lowerCAmelCase ) )
		model.eval()
		model.save_pretrained(lowerCAmelCase )

		print(F'Checkpoint successfuly converted. Model saved at {pytorch_dump_path}' )


if __name__ == "__main__":
		lowerCamelCase__	=							argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    "--pytorch_model_path", default=None, type=str, required=True, help="Path to YOSO pytorch checkpoint."
		)
		parser.add_argument(
		    "--config_file",
		    default=None,
		    type=str,
		    required=True,
		    help="The json file for YOSO model config.",
		)
		parser.add_argument(
		    "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
		)
		lowerCamelCase__	=							parser.parse_args()
		convert_yoso_checkpoint(args.pytorch_model_path, args.config_file, args.pytorch_dump_path)



