





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  (	UpperCamelCase__   :   List[str]       ,			UpperCamelCase__   :   List[str]       ,			UpperCamelCase__   :   str       ,			UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   Any					)      ->    Optional[int]:

						"""simple docstring"""





						with open(UpperCamelCase__					) as metadata_file:
												__SCREAMING_SNAKE_CASE:				Any           =   json.load(UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				Tuple           =   LukeConfig(use_entity_aware_attention=UpperCamelCase__       ,			**metadata['''model_config''']					)

						# Load in the weights from the checkpoint_path
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   torch.load(UpperCamelCase__       ,			map_location='''cpu'''					)

						# Load the entity vocab file
						__SCREAMING_SNAKE_CASE:				Dict           =   load_entity_vocab(UpperCamelCase__					)

						__SCREAMING_SNAKE_CASE:				List[str]           =   RobertaTokenizer.from_pretrained(metadata['''model_config''']['''bert_model_name''']					)

						# Add special tokens to the token vocabulary for downstream tasks
						__SCREAMING_SNAKE_CASE:				int           =   AddedToken('''<ent>'''       ,			lstrip=UpperCamelCase__       ,			rstrip=UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				int           =   AddedToken('''<ent2>'''       ,			lstrip=UpperCamelCase__       ,			rstrip=UpperCamelCase__					)
						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(UpperCamelCase__					)
						with open(os.path.join(UpperCamelCase__       ,			LukeTokenizer.vocab_files_names['''entity_vocab_file''']					)       ,			'''w'''					) as f:
												json.dump(UpperCamelCase__       ,			UpperCamelCase__					)

						__SCREAMING_SNAKE_CASE:				Tuple           =   LukeTokenizer.from_pretrained(UpperCamelCase__					)

						# Initialize the embeddings of the special tokens
						__SCREAMING_SNAKE_CASE:				Dict           =   state_dict['''embeddings.word_embeddings.weight''']
						__SCREAMING_SNAKE_CASE:				Tuple           =   word_emb[tokenizer.convert_tokens_to_ids(['''@''']					)[0]].unsqueeze(0					)
						__SCREAMING_SNAKE_CASE:				Any           =   word_emb[tokenizer.convert_tokens_to_ids(['''#''']					)[0]].unsqueeze(0					)
						__SCREAMING_SNAKE_CASE:				List[str]           =   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"]:
																		__SCREAMING_SNAKE_CASE:				Dict           =   F"""encoder.layer.{layer_index}.attention.self."""
																		__SCREAMING_SNAKE_CASE:				int           =   state_dict[prefix + matrix_name]
																		__SCREAMING_SNAKE_CASE:				List[Any]           =   state_dict[prefix + matrix_name]
																		__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   state_dict[prefix + matrix_name]

    # Initialize the embedding of the [MASK2] entity using that of the [MASK] entity for downstream tasks
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   state_dict['''entity_embeddings.entity_embeddings.weight''']
						__SCREAMING_SNAKE_CASE:				str           =   entity_emb[entity_vocab['''[MASK]''']]

						__SCREAMING_SNAKE_CASE:				str           =   LukeModel(config=UpperCamelCase__					).eval()

						__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE:				Tuple           =   model.load_state_dict(UpperCamelCase__       ,			strict=UpperCamelCase__					)
						if not (len(UpperCamelCase__					) == 1 and missing_keys[0] == "embeddings.position_ids"):
												raise ValueError(F"""Missing keys {", ".join(UpperCamelCase__					)}. 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
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   LukeTokenizer.from_pretrained(UpperCamelCase__       ,			task='''entity_classification'''					)

						__SCREAMING_SNAKE_CASE:				List[Any]           =   (
						    '''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 .'''
						)
						__SCREAMING_SNAKE_CASE:				str           =   (39, 42)
						__SCREAMING_SNAKE_CASE:				str           =   tokenizer(UpperCamelCase__       ,			entity_spans=[span]       ,			add_prefix_space=UpperCamelCase__       ,			return_tensors='''pt'''					)

						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   model(**UpperCamelCase__					)

						# Verify word hidden states
						if model_size == "large":
												__SCREAMING_SNAKE_CASE:				List[str]           =   torch.Size((1, 42, 1_024)					)
												__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   torch.tensor(
												    [[0.01_33, 0.08_65, 0.00_95], [0.30_93, -0.25_76, -0.74_18], [-0.17_20, -0.21_17, -0.28_69]]					)
						else:  # base
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   torch.Size((1, 42, 768)					)
												__SCREAMING_SNAKE_CASE:				int           =   torch.tensor([[0.00_37, 0.13_68, -0.00_91], [0.10_99, 0.33_29, -0.10_95], [0.07_65, 0.53_35, 0.11_79]]					)

						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]       ,			UpperCamelCase__       ,			atol=1E-4					):
												raise ValueError

						# Verify entity hidden states
						if model_size == "large":
												__SCREAMING_SNAKE_CASE:				str           =   torch.Size((1, 1, 1_024)					)
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   torch.tensor([[0.04_66, -0.01_06, -0.01_79]]					)
						else:  # base
												__SCREAMING_SNAKE_CASE:				List[str]           =   torch.Size((1, 1, 768)					)
												__SCREAMING_SNAKE_CASE:				Dict           =   torch.tensor([[0.14_57, 0.10_44, 0.01_74]]					)

						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]       ,			UpperCamelCase__       ,			atol=1E-4					):
												raise ValueError

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






def 		lowerCAmelCase  (	UpperCamelCase__   :   Dict					)      ->    Tuple:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				str           =   {}
						with open(UpperCamelCase__       ,			'''r'''       ,			encoding='''utf-8'''					) as f:
												for index, line in enumerate(UpperCamelCase__					):
																		__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE:				Any           =   line.rstrip().split('''\t'''					)
																		__SCREAMING_SNAKE_CASE:				Tuple           =   index

						return entity_vocab


if __name__ == "__main__":
							lowerCAmelCase							:       Tuple					 =	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."""
							)
							lowerCAmelCase							:       List[str]					 =	parser.parse_args()
							convert_luke_checkpoint(
							    args.checkpoint_path,
							    args.metadata_path,
							    args.entity_vocab_path,
							    args.pytorch_dump_folder_path,
							    args.model_size,
							)










import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    CLIPConfig,
    CLIPImageProcessor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER


sys.path.append(str(Path(__file__).parent.parent.parent.parent / """utils"""))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_image_processing import CustomImageProcessor  # noqa E402





class 				a			(      unittest.TestCase			):






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   0






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[str]           =   AutoImageProcessor.from_pretrained('''openai/clip-vit-base-patch32'''					)
										self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   Path(_lowerCAmelCase					) / '''preprocessor_config.json'''
																__SCREAMING_SNAKE_CASE:				List[str]           =   Path(_lowerCAmelCase					) / '''config.json'''
																json.dump(
																    {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}			, open(_lowerCAmelCase			, '''w'''					)			, )
																json.dump({'''model_type''': '''clip'''}			, open(_lowerCAmelCase			, '''w'''					)					)

																__SCREAMING_SNAKE_CASE:				Dict           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					)
																self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																__SCREAMING_SNAKE_CASE:				int           =   Path(_lowerCAmelCase					) / '''preprocessor_config.json'''
																__SCREAMING_SNAKE_CASE:				Tuple           =   Path(_lowerCAmelCase					) / '''config.json'''
																json.dump(
																    {'''feature_extractor_type''': '''CLIPFeatureExtractor''', '''processor_class''': '''CLIPProcessor'''}			, open(_lowerCAmelCase			, '''w'''					)			, )
																json.dump({'''model_type''': '''clip'''}			, open(_lowerCAmelCase			, '''w'''					)					)

																__SCREAMING_SNAKE_CASE:				str           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					)
																self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																__SCREAMING_SNAKE_CASE:				int           =   CLIPConfig()

																# Create a dummy config file with image_proceesor_type
																__SCREAMING_SNAKE_CASE:				Any           =   Path(_lowerCAmelCase					) / '''preprocessor_config.json'''
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   Path(_lowerCAmelCase					) / '''config.json'''
																json.dump(
																    {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}			, open(_lowerCAmelCase			, '''w'''					)			, )
																json.dump({'''model_type''': '''clip'''}			, open(_lowerCAmelCase			, '''w'''					)					)

																# remove image_processor_type to make sure config.json alone is enough to load image processor locally
																__SCREAMING_SNAKE_CASE:				Any           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					).to_dict()

																config_dict.pop('''image_processor_type'''					)
																__SCREAMING_SNAKE_CASE:				List[str]           =   CLIPImageProcessor(**_lowerCAmelCase					)

																# save in new folder
																model_config.save_pretrained(_lowerCAmelCase					)
																config.save_pretrained(_lowerCAmelCase					)

																__SCREAMING_SNAKE_CASE:				Dict           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					)

																# make sure private variable is not incorrectly saved
																__SCREAMING_SNAKE_CASE:				int           =   json.loads(config.to_json_string()					)
																self.assertTrue('''_processor_class''' not in dict_as_saved					)

										self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																__SCREAMING_SNAKE_CASE:				Optional[int]           =   Path(_lowerCAmelCase					) / '''preprocessor_config.json'''
																json.dump(
																    {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}			, open(_lowerCAmelCase			, '''w'''					)			, )

																__SCREAMING_SNAKE_CASE:				int           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					)
																self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with self.assertRaisesRegex(
										    _lowerCAmelCase			, '''clip-base is not a local folder and is not a valid model identifier'''					):
																__SCREAMING_SNAKE_CASE:				Tuple           =   AutoImageProcessor.from_pretrained('''clip-base'''					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with self.assertRaisesRegex(
										    _lowerCAmelCase			, R'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''					):
																__SCREAMING_SNAKE_CASE:				Dict           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase			, revision='''aaaaaa'''					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with self.assertRaisesRegex(
										    _lowerCAmelCase			, '''hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json.'''			, ):
																__SCREAMING_SNAKE_CASE:				int           =   AutoImageProcessor.from_pretrained('''hf-internal-testing/config-no-model'''					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with self.assertRaises(_lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   AutoImageProcessor.from_pretrained('''hf-internal-testing/test_dynamic_image_processor'''					)
										# If remote code is disabled, we can't load this config.
										with self.assertRaises(_lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				Any           =   AutoImageProcessor.from_pretrained(
																    '''hf-internal-testing/test_dynamic_image_processor'''			, trust_remote_code=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   AutoImageProcessor.from_pretrained(
										    '''hf-internal-testing/test_dynamic_image_processor'''			, trust_remote_code=_lowerCAmelCase					)
										self.assertEqual(image_processor.__class__.__name__			, '''NewImageProcessor'''					)

										# Test image processor can be reloaded.
										with tempfile.TemporaryDirectory() as tmp_dir:
																image_processor.save_pretrained(_lowerCAmelCase					)
																__SCREAMING_SNAKE_CASE:				List[str]           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase			, trust_remote_code=_lowerCAmelCase					)
										self.assertEqual(reloaded_image_processor.__class__.__name__			, '''NewImageProcessor'''					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										try:
																AutoConfig.register('''custom'''			, _lowerCAmelCase					)
																AutoImageProcessor.register(_lowerCAmelCase			, _lowerCAmelCase					)
																# Trying to register something existing in the Transformers library will raise an error
																with self.assertRaises(_lowerCAmelCase					):
																						AutoImageProcessor.register(_lowerCAmelCase			, _lowerCAmelCase					)

																with tempfile.TemporaryDirectory() as tmpdirname:
																						__SCREAMING_SNAKE_CASE:				Optional[int]           =   Path(_lowerCAmelCase					) / '''preprocessor_config.json'''
																						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   Path(_lowerCAmelCase					) / '''config.json'''
																						json.dump(
																						    {'''feature_extractor_type''': '''CLIPFeatureExtractor''', '''processor_class''': '''CLIPProcessor'''}			, open(_lowerCAmelCase			, '''w'''					)			, )
																						json.dump({'''model_type''': '''clip'''}			, open(_lowerCAmelCase			, '''w'''					)					)

																						__SCREAMING_SNAKE_CASE:				List[str]           =   CustomImageProcessor.from_pretrained(_lowerCAmelCase					)

																# Now that the config is registered, it can be used as any other config with the auto-API
																with tempfile.TemporaryDirectory() as tmp_dir:
																						image_processor.save_pretrained(_lowerCAmelCase					)
																						__SCREAMING_SNAKE_CASE:				int           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					)
																						self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)

										finally:
																if "custom" in CONFIG_MAPPING._extra_content:
																						del CONFIG_MAPPING._extra_content["custom"]
																if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
																						del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]






				def 			snake_case_    (  self					):

										"""simple docstring"""




										class 				a			(      __lowercase			):
														SCREAMING_SNAKE_CASE__     :   Tuple		      =     True

										try:
																AutoConfig.register('''custom'''			, _lowerCAmelCase					)
																AutoImageProcessor.register(_lowerCAmelCase			, _lowerCAmelCase					)
																# If remote code is not set, the default is to use local
																__SCREAMING_SNAKE_CASE:				str           =   AutoImageProcessor.from_pretrained('''hf-internal-testing/test_dynamic_image_processor'''					)
																self.assertEqual(image_processor.__class__.__name__			, '''NewImageProcessor'''					)
																self.assertTrue(image_processor.is_local					)

																# If remote code is disabled, we load the local one.
																__SCREAMING_SNAKE_CASE:				Tuple           =   AutoImageProcessor.from_pretrained(
																    '''hf-internal-testing/test_dynamic_image_processor'''			, trust_remote_code=_lowerCAmelCase					)
																self.assertEqual(image_processor.__class__.__name__			, '''NewImageProcessor'''					)
																self.assertTrue(image_processor.is_local					)

																# If remote is enabled, we load from the Hub
																__SCREAMING_SNAKE_CASE:				Any           =   AutoImageProcessor.from_pretrained(
																    '''hf-internal-testing/test_dynamic_image_processor'''			, trust_remote_code=_lowerCAmelCase					)
																self.assertEqual(image_processor.__class__.__name__			, '''NewImageProcessor'''					)
																self.assertTrue(not hasattr(_lowerCAmelCase			, '''is_local'''					)					)

										finally:
																if "custom" in CONFIG_MAPPING._extra_content:
																						del CONFIG_MAPPING._extra_content["custom"]
																if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
																						del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]



