


"""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				,       lowerCamelCase__				,       lowerCamelCase__=1_3				,       lowerCamelCase__=7				,       lowerCamelCase__=True				,       lowerCamelCase__=True				,       lowerCamelCase__=False				,       lowerCamelCase__=True				,       lowerCamelCase__=9_9				,       lowerCamelCase__=3_2				,       lowerCamelCase__=5				,       lowerCamelCase__=4				,       lowerCamelCase__=3_7				,       lowerCamelCase__="gelu"				,       lowerCamelCase__=0.1				,       lowerCamelCase__=0.1				,       lowerCamelCase__=5_1_2				,       lowerCamelCase__=1_6				,       lowerCamelCase__=2				,       lowerCamelCase__=0.0_2				,       lowerCamelCase__=3				,       lowerCamelCase__=4				,       lowerCamelCase__=None				,       ):
									_lowerCamelCase        =				parent
									_lowerCamelCase        =				batch_size
									_lowerCamelCase        =				seq_length
									_lowerCamelCase        =				is_training
									_lowerCamelCase        =				use_input_mask
									_lowerCamelCase        =				use_token_type_ids
									_lowerCamelCase        =				use_labels
									_lowerCamelCase        =				vocab_size
									_lowerCamelCase        =				hidden_size
									_lowerCamelCase        =				num_hidden_layers
									_lowerCamelCase        =				num_attention_heads
									_lowerCamelCase        =				intermediate_size
									_lowerCamelCase        =				hidden_act
									_lowerCamelCase        =				hidden_dropout_prob
									_lowerCamelCase        =				attention_probs_dropout_prob
									_lowerCamelCase        =				max_position_embeddings
									_lowerCamelCase        =				type_vocab_size
									_lowerCamelCase        =				type_sequence_label_size
									_lowerCamelCase        =				initializer_range
									_lowerCamelCase        =				num_labels
									_lowerCamelCase        =				num_choices
									_lowerCamelCase        =				scope

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size   )

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

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

									_lowerCamelCase        =				self.get_config()

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

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

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
									_lowerCamelCase        =				DistilBertForMaskedLM(config=__SCREAMING_SNAKE_CASE   )
									model.to(__SCREAMING_SNAKE_CASE   )
									model.eval()
									_lowerCamelCase        =				model(__SCREAMING_SNAKE_CASE				,       attention_mask=__SCREAMING_SNAKE_CASE				,       labels=__SCREAMING_SNAKE_CASE   )
									self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)   )

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
									_lowerCamelCase        =				DistilBertForQuestionAnswering(config=__SCREAMING_SNAKE_CASE   )
									model.to(__SCREAMING_SNAKE_CASE   )
									model.eval()
									_lowerCamelCase        =				model(
									    __SCREAMING_SNAKE_CASE				,       attention_mask=__SCREAMING_SNAKE_CASE				,       start_positions=__SCREAMING_SNAKE_CASE				,       end_positions=__SCREAMING_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 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
									_lowerCamelCase        =				self.num_labels
									_lowerCamelCase        =				DistilBertForSequenceClassification(__SCREAMING_SNAKE_CASE   )
									model.to(__SCREAMING_SNAKE_CASE   )
									model.eval()
									_lowerCamelCase        =				model(__SCREAMING_SNAKE_CASE				,       attention_mask=__SCREAMING_SNAKE_CASE				,       labels=__SCREAMING_SNAKE_CASE   )
									self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)   )

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
									_lowerCamelCase        =				self.num_labels
									_lowerCamelCase        =				DistilBertForTokenClassification(config=__SCREAMING_SNAKE_CASE   )
									model.to(__SCREAMING_SNAKE_CASE   )
									model.eval()

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

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





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


@require_torch
class 							lowerCamelCase_(   __UpperCAmelCase,				__UpperCAmelCase,				unittest.TestCase    ):





					'''simple docstring'''



					lowercase__   :       Optional[int]										= (
					    (
					        DistilBertModel,
					        DistilBertForMaskedLM,
					        DistilBertForMultipleChoice,
					        DistilBertForQuestionAnswering,
					        DistilBertForSequenceClassification,
					        DistilBertForTokenClassification,
					    )
					    if is_torch_available()
					    else None
					)
					lowercase__   :       List[str]										= (
					    {
					        '''feature-extraction''': DistilBertModel,
					        '''fill-mask''': DistilBertForMaskedLM,
					        '''question-answering''': DistilBertForQuestionAnswering,
					        '''text-classification''': DistilBertForSequenceClassification,
					        '''token-classification''': DistilBertForTokenClassification,
					        '''zero-shot''': DistilBertForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					lowercase__   :       Any										= True
					lowercase__   :       List[Any]										= True
					lowercase__   :       str										= True
					lowercase__   :       Optional[int]										= True

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				DistilBertModelTester(self   )
									_lowerCamelCase        =				ConfigTester(self				,       config_class=__SCREAMING_SNAKE_CASE				,       dim=3_7   )

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

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_model(*__SCREAMING_SNAKE_CASE   )

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_masked_lm(*__SCREAMING_SNAKE_CASE   )

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_question_answering(*__SCREAMING_SNAKE_CASE   )

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_sequence_classification(*__SCREAMING_SNAKE_CASE   )

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_token_classification(*__SCREAMING_SNAKE_CASE   )

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_multiple_choice(*__SCREAMING_SNAKE_CASE   )

					@slow
					def 			snake_case__      (    self   ):

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





					@slow
					@require_torch_gpu
					def 			snake_case__      (    self   ):
									_lowerCamelCase  ,     _lowerCamelCase        =				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

													_lowerCamelCase        =				True
													_lowerCamelCase        =				model_class(config=__SCREAMING_SNAKE_CASE   )

													_lowerCamelCase        =				self._prepare_for_class(__SCREAMING_SNAKE_CASE				,       __SCREAMING_SNAKE_CASE   )
													_lowerCamelCase        =				torch.jit.trace(
													    __SCREAMING_SNAKE_CASE				,       (inputs_dict['''input_ids'''].to('''cpu'''   ), inputs_dict['''attention_mask'''].to('''cpu'''   ))   )

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


@require_torch
class 							lowerCamelCase_(   unittest.TestCase    ):





					'''simple docstring'''



					@slow
					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				DistilBertModel.from_pretrained('''distilbert-base-uncased'''   )
									_lowerCamelCase        =				torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]   )
									_lowerCamelCase        =				torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]   )
									with torch.no_grad():
													_lowerCamelCase        =				model(__SCREAMING_SNAKE_CASE				,       attention_mask=__SCREAMING_SNAKE_CASE   )[0]
									_lowerCamelCase        =				torch.Size((1, 1_1, 7_6_8)   )
									self.assertEqual(output.shape				,       __SCREAMING_SNAKE_CASE   )
									_lowerCamelCase        =				torch.tensor(
									    [[[-0.1_6_3_9, 0.3_2_9_9, 0.1_6_4_8], [-0.1_7_4_6, 0.3_2_8_9, 0.1_7_1_0], [-0.1_8_8_4, 0.3_3_5_7, 0.1_8_1_0]]]   )

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





"""simple docstring"""
import argparse
import json
import os

import torch

from transformers import LukeConfig, LukeModel, LukeTokenizer, RobertaTokenizer
from transformers.tokenization_utils_base import AddedToken



@torch.no_grad()
def        lowerCAmelCase_(   lowercase_     :							int						,   lowercase_     :							List[Any]						,   lowercase_     :							List[Any]						,   lowercase_     :							Tuple						,   lowercase_     :							List[str]							)	->       Dict:
				# Load configuration defined in the metadata file
				with open(lowercase_							) as metadata_file:
								_lowerCamelCase        =				json.load(lowercase_							)
				_lowerCamelCase        =				LukeConfig(use_entity_aware_attention=lowercase_						,   **metadata['''model_config''']							)

				# Load in the weights from the checkpoint_path
				_lowerCamelCase        =				torch.load(lowercase_						,   map_location='''cpu'''							)

				# Load the entity vocab file
				_lowerCamelCase        =				load_entity_vocab(lowercase_							)

				_lowerCamelCase        =				RobertaTokenizer.from_pretrained(metadata['''model_config''']['''bert_model_name''']							)

				# Add special tokens to the token vocabulary for downstream tasks
				_lowerCamelCase        =				AddedToken('''<ent>'''						,   lstrip=lowercase_						,   rstrip=lowercase_							)
				_lowerCamelCase        =				AddedToken('''<ent2>'''						,   lstrip=lowercase_						,   rstrip=lowercase_							)
				tokenizer.add_special_tokens({'''additional_special_tokens''': [entity_token_a, entity_token_a]}							)
				config.vocab_size += 2

				print(F"""Saving tokenizer to {pytorch_dump_folder_path}"""							)
				tokenizer.save_pretrained(lowercase_							)
				with open(os.path.join(lowercase_						,   LukeTokenizer.vocab_files_names['''entity_vocab_file''']							)						,   '''w'''							) as f:
								json.dump(lowercase_						,   lowercase_							)

				_lowerCamelCase        =				LukeTokenizer.from_pretrained(lowercase_							)

				# Initialize the embeddings of the special tokens
				_lowerCamelCase        =				state_dict['''embeddings.word_embeddings.weight''']
				_lowerCamelCase        =				word_emb[tokenizer.convert_tokens_to_ids(['''@''']							)[0]].unsqueeze(0							)
				_lowerCamelCase        =				word_emb[tokenizer.convert_tokens_to_ids(['''#''']							)[0]].unsqueeze(0							)
				_lowerCamelCase        =				torch.cat([word_emb, ent_emb, enta_emb]							)

				# Initialize the query layers of the entity-aware self-attention mechanism
				for layer_index in range(config.num_hidden_layers							):
								for matrix_name in ["query.weight", "query.bias"]:
												_lowerCamelCase        =				F"""encoder.layer.{layer_index}.attention.self."""
												_lowerCamelCase        =				state_dict[prefix + matrix_name]
												_lowerCamelCase        =				state_dict[prefix + matrix_name]
												_lowerCamelCase        =				state_dict[prefix + matrix_name]

    # Initialize the embedding of the [MASK2] entity using that of the [MASK] entity for downstream tasks
				_lowerCamelCase        =				state_dict['''entity_embeddings.entity_embeddings.weight''']
				_lowerCamelCase        =				entity_emb[entity_vocab['''[MASK]''']]

				_lowerCamelCase        =				LukeModel(config=lowercase_							).eval()

				_lowerCamelCase  ,     _lowerCamelCase        =				model.load_state_dict(lowercase_						,   strict=lowercase_							)
				if not (len(lowercase_							) == 1 and missing_keys[0] == "embeddings.position_ids"):
								raise ValueError(F"""Missing keys {", ".join(lowercase_							)}. Expected only missing embeddings.position_ids"""							)
				if not (all(key.startswith('''entity_predictions'''							) or key.startswith('''lm_head'''							) for key in unexpected_keys							)):
								raise ValueError(
								    '''Unexpected keys'''
								    F""" {", ".join([key for key in unexpected_keys if not (key.startswith("entity_predictions"							) or key.startswith("lm_head"							))]							)}"""							)

				# Check outputs
				_lowerCamelCase        =				LukeTokenizer.from_pretrained(lowercase_						,   task='''entity_classification'''							)

				_lowerCamelCase        =				(
				    '''Top seed Ana Ivanovic said on Thursday she could hardly believe her luck as a fortuitous netcord helped the'''
				    ''' new world number one avoid a humiliating second- round exit at Wimbledon .'''
				)
				_lowerCamelCase        =				(39, 42)
				_lowerCamelCase        =				tokenizer(lowercase_						,   entity_spans=[span]						,   add_prefix_space=lowercase_						,   return_tensors='''pt'''							)

				_lowerCamelCase        =				model(**lowercase_							)

				# Verify word hidden states
				if model_size == "large":
								_lowerCamelCase        =				torch.Size((1, 42, 10_24)							)
								_lowerCamelCase        =				torch.tensor(
								    [[0.0_1_3_3, 0.0_8_6_5, 0.0_0_9_5], [0.3_0_9_3, -0.2_5_7_6, -0.7_4_1_8], [-0.1_7_2_0, -0.2_1_1_7, -0.2_8_6_9]]							)
				else:  # base
								_lowerCamelCase        =				torch.Size((1, 42, 7_68)							)
								_lowerCamelCase        =				torch.tensor([[0.0_0_3_7, 0.1_3_6_8, -0.0_0_9_1], [0.1_0_9_9, 0.3_3_2_9, -0.1_0_9_5], [0.0_7_6_5, 0.5_3_3_5, 0.1_1_7_9]]							)

				if not (outputs.last_hidden_state.shape == expected_shape):
								raise ValueError(
								    F"""Outputs.last_hidden_state.shape is {outputs.last_hidden_state.shape}, Expected shape is {expected_shape}"""							)
				if not torch.allclose(outputs.last_hidden_state[0, :3, :3]						,   lowercase_						,   atol=1e-4							):
								raise ValueError

				# Verify entity hidden states
				if model_size == "large":
								_lowerCamelCase        =				torch.Size((1, 1, 10_24)							)
								_lowerCamelCase        =				torch.tensor([[0.0_4_6_6, -0.0_1_0_6, -0.0_1_7_9]]							)
				else:  # base
								_lowerCamelCase        =				torch.Size((1, 1, 7_68)							)
								_lowerCamelCase        =				torch.tensor([[0.1_4_5_7, 0.1_0_4_4, 0.0_1_7_4]]							)

				if not (outputs.entity_last_hidden_state.shape != expected_shape):
								raise ValueError(
								    F"""Outputs.entity_last_hidden_state.shape is {outputs.entity_last_hidden_state.shape}, Expected shape is"""
								    F""" {expected_shape}"""							)
				if not torch.allclose(outputs.entity_last_hidden_state[0, :3, :3]						,   lowercase_						,   atol=1e-4							):
								raise ValueError

				# Finally, save our PyTorch model and tokenizer
				print('''Saving PyTorch model to {}'''.format(lowercase_							)							)
				model.save_pretrained(lowercase_							)



def        lowerCAmelCase_(   lowercase_     :							Optional[Any]							)	->       Any:
				_lowerCamelCase        =				{}
				with open(lowercase_						,   '''r'''						,   encoding='''utf-8'''							) as f:
								for index, line in enumerate(lowercase_							):
												_lowerCamelCase  ,     _lowerCamelCase        =				line.rstrip().split('''\t'''							)
												_lowerCamelCase        =				index

				return entity_vocab


if __name__ == "__main__":
			__SCREAMING_SNAKE_CASE  :						List[str]    							=      argparse.ArgumentParser()
			# Required parameters
			parser.add_argument('''--checkpoint_path''', type=str, help='''Path to a pytorch_model.bin file.''')
			parser.add_argument(
			    '''--metadata_path''', default=None, type=str, help='''Path to a metadata.json file, defining the configuration.'''
			)
			parser.add_argument(
			    '''--entity_vocab_path''',
			    default=None,
			    type=str,
			    help='''Path to an entity_vocab.tsv file, containing the entity vocabulary.''',
			)
			parser.add_argument(
			    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to where to dump the output PyTorch model.'''
			)
			parser.add_argument(
			    '''--model_size''', default='''base''', type=str, choices=['''base''', '''large'''], help='''Size of the model to be converted.'''
			)
			__SCREAMING_SNAKE_CASE  :						int    							=      parser.parse_args()
			convert_luke_checkpoint(
			    args.checkpoint_path,
			    args.metadata_path,
			    args.entity_vocab_path,
			    args.pytorch_dump_folder_path,
			    args.model_size,
			)
