



import argparse
import json
import os
from collections import OrderedDict

import torch

from transformers import LukeConfig, LukeForMaskedLM, MLukeTokenizer, XLMRobertaTokenizer
from transformers.tokenization_utils_base import AddedToken





@torch.no_grad()
def 							A			(  _UpperCAmelCase       :     Tuple							,     _UpperCAmelCase       :     List[Any]							,     _UpperCAmelCase       :     Optional[Any]							,     _UpperCAmelCase       :     List[Any]							,     _UpperCAmelCase       :     str					)       ->				int:







			'''simple docstring'''
			# Load configuration defined in the metadata file
			with open(_UpperCAmelCase					) as metadata_file:
						_UpperCAmelCase	 =     json.load(_UpperCAmelCase					)
			_UpperCAmelCase	 =     LukeConfig(use_entity_aware_attention=_UpperCAmelCase							,     **metadata['model_config']					)

			# Load in the weights from the checkpoint_path
			_UpperCAmelCase	 =     torch.load(_UpperCAmelCase							,     map_location='cpu'					)['module']

			# Load the entity vocab file
			_UpperCAmelCase	 =     load_original_entity_vocab(_UpperCAmelCase					)
			# add an entry for [MASK2]
			_UpperCAmelCase	 =     max(entity_vocab.values()					) + 1
			config.entity_vocab_size += 1

			_UpperCAmelCase	 =     XLMRobertaTokenizer.from_pretrained(metadata['model_config']['bert_model_name']					)

			# Add special tokens to the token vocabulary for downstream tasks
			_UpperCAmelCase	 =     AddedToken('<ent>'							,     lstrip=_UpperCAmelCase							,     rstrip=_UpperCAmelCase					)
			_UpperCAmelCase	 =     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							,     'tokenizer_config.json'					)							,     'r'					) as f:
						_UpperCAmelCase	 =     json.load(_UpperCAmelCase					)
			_UpperCAmelCase	 =     'MLukeTokenizer'
			with open(os.path.join(_UpperCAmelCase							,     'tokenizer_config.json'					)							,     'w'					) as f:
						json.dump(_UpperCAmelCase							,     _UpperCAmelCase					)

			with open(os.path.join(_UpperCAmelCase							,     MLukeTokenizer.vocab_files_names['entity_vocab_file']					)							,     'w'					) as f:
						json.dump(_UpperCAmelCase							,     _UpperCAmelCase					)

			_UpperCAmelCase	 =     MLukeTokenizer.from_pretrained(_UpperCAmelCase					)

			# Initialize the embeddings of the special tokens
			_UpperCAmelCase	 =     tokenizer.convert_tokens_to_ids(['@']					)[0]
			_UpperCAmelCase	 =     tokenizer.convert_tokens_to_ids(['#']					)[0]

			_UpperCAmelCase	 =     state_dict['embeddings.word_embeddings.weight']
			_UpperCAmelCase	 =     word_emb[ent_init_index].unsqueeze(0					)
			_UpperCAmelCase	 =     word_emb[enta_init_index].unsqueeze(0					)
			_UpperCAmelCase	 =     torch.cat([word_emb, ent_emb, enta_emb]					)
			# add special tokens for 'entity_predictions.bias'
			for bias_name in ["lm_head.decoder.bias", "lm_head.bias"]:
						_UpperCAmelCase	 =     state_dict[bias_name]
						_UpperCAmelCase	 =     decoder_bias[ent_init_index].unsqueeze(0					)
						_UpperCAmelCase	 =     decoder_bias[enta_init_index].unsqueeze(0					)
						_UpperCAmelCase	 =     torch.cat([decoder_bias, ent_decoder_bias, enta_decoder_bias]					)

			# 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"]:
									_UpperCAmelCase	 =     F"encoder.layer.{layer_index}.attention.self."
									_UpperCAmelCase	 =     state_dict[prefix + matrix_name]
									_UpperCAmelCase	 =     state_dict[prefix + matrix_name]
									_UpperCAmelCase	 =     state_dict[prefix + matrix_name]

    # Initialize the embedding of the [MASK2] entity using that of the [MASK] entity for downstream tasks
			_UpperCAmelCase	 =     state_dict['entity_embeddings.entity_embeddings.weight']
			_UpperCAmelCase	 =     entity_emb[entity_vocab['[MASK]']].unsqueeze(0					)
			_UpperCAmelCase	 =     torch.cat([entity_emb, entity_mask_emb]					)
			# add [MASK2] for 'entity_predictions.bias'
			_UpperCAmelCase	 =     state_dict['entity_predictions.bias']
			_UpperCAmelCase	 =     entity_prediction_bias[entity_vocab['[MASK]']].unsqueeze(0					)
			_UpperCAmelCase	 =     torch.cat([entity_prediction_bias, entity_mask_bias]					)

			_UpperCAmelCase	 =     LukeForMaskedLM(config=_UpperCAmelCase					).eval()

			state_dict.pop('entity_predictions.decoder.weight'					)
			state_dict.pop('lm_head.decoder.weight'					)
			state_dict.pop('lm_head.decoder.bias'					)
			_UpperCAmelCase	 =     OrderedDict()
			for key, value in state_dict.items():
						if not (key.startswith('lm_head'					) or key.startswith('entity_predictions'					)):
									_UpperCAmelCase	 =     state_dict[key]
						else:
									_UpperCAmelCase	 =     state_dict[key]

			_UpperCAmelCase ,  _UpperCAmelCase	 =     model.load_state_dict(_UpperCAmelCase							,     strict=_UpperCAmelCase					)

			if set(_UpperCAmelCase					) != {"luke.embeddings.position_ids"}:
						raise ValueError(F"Unexpected unexpected_keys: {unexpected_keys}"					)
			if set(_UpperCAmelCase					) != {
			    "lm_head.decoder.weight",
			    "lm_head.decoder.bias",
			    "entity_predictions.decoder.weight",
			}:
						raise ValueError(F"Unexpected missing_keys: {missing_keys}"					)

			model.tie_weights()
			assert (model.luke.embeddings.word_embeddings.weight == model.lm_head.decoder.weight).all()
			assert (model.luke.entity_embeddings.entity_embeddings.weight == model.entity_predictions.decoder.weight).all()

			# Check outputs
			_UpperCAmelCase	 =     MLukeTokenizer.from_pretrained(_UpperCAmelCase							,     task='entity_classification'					)

			_UpperCAmelCase	 =     'ISO 639-3 uses the code fas for the dialects spoken across Iran and アフガニスタン (Afghanistan).'
			_UpperCAmelCase	 =     (0, 9)
			_UpperCAmelCase	 =     tokenizer(_UpperCAmelCase							,     entity_spans=[span]							,     return_tensors='pt'					)

			_UpperCAmelCase	 =     model(**_UpperCAmelCase					)

			# Verify word hidden states
			if model_size == "large":
						raise NotImplementedError
			else:  # base
						_UpperCAmelCase	 =     torch.Size((1, 33, 768)					)
						_UpperCAmelCase	 =     torch.tensor([[0.0892, 0.0596, -0.2819], [0.0134, 0.1199, 0.0573], [-0.0169, 0.0927, 0.0644]]					)

			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":
						raise NotImplementedError
			else:  # base
						_UpperCAmelCase	 =     torch.Size((1, 1, 768)					)
						_UpperCAmelCase	 =     torch.tensor([[-0.1482, 0.0609, 0.0322]]					)

			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

			# Verify masked word/entity prediction
			_UpperCAmelCase	 =     MLukeTokenizer.from_pretrained(_UpperCAmelCase					)
			_UpperCAmelCase	 =     'Tokyo is the capital of <mask>.'
			_UpperCAmelCase	 =     (24, 30)
			_UpperCAmelCase	 =     tokenizer(_UpperCAmelCase							,     entity_spans=[span]							,     return_tensors='pt'					)

			_UpperCAmelCase	 =     model(**_UpperCAmelCase					)

			_UpperCAmelCase	 =     encoding['input_ids'][0].tolist()
			_UpperCAmelCase	 =     input_ids.index(tokenizer.convert_tokens_to_ids('<mask>'					)					)
			_UpperCAmelCase	 =     outputs.logits[0][mask_position_id].argmax(dim=-1					)
			assert "Japan" == tokenizer.decode(_UpperCAmelCase					)

			_UpperCAmelCase	 =     outputs.entity_logits[0][0].argmax().item()
			_UpperCAmelCase	 =     [
			    entity for entity, entity_id in tokenizer.entity_vocab.items() if entity_id == predicted_entity_id
			]
			assert [e for e in multilingual_predicted_entities if e.startswith('en:'					)][0] == "en:Japan"

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





def 							A			(  _UpperCAmelCase       :     Any					)       ->				List[str]:







			'''simple docstring'''
			_UpperCAmelCase	 =     ['[MASK]', '[PAD]', '[UNK]']

			_UpperCAmelCase	 =     [json.loads(_UpperCAmelCase					) for line in open(_UpperCAmelCase					)]

			_UpperCAmelCase	 =     {}
			for entry in data:
						_UpperCAmelCase	 =     entry['id']
						for entity_name, language in entry["entities"]:
									if entity_name in SPECIAL_TOKENS:
												_UpperCAmelCase	 =     entity_id
												break
									_UpperCAmelCase	 =     F"{language}:{entity_name}"
									_UpperCAmelCase	 =     entity_id
			return new_mapping


if __name__ == "__main__":
							UpperCAmelCase__           = 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."
							)
							UpperCAmelCase__           = 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 unittest

from transformers import is_vision_available
from transformers.pipelines import pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
							from PIL import Image
else:







							class     __lowerCAmelCase					:



										@staticmethod
										def 					_lowerCamelCase (		*A     :						Union[str, Any]				,	**A     :						List[Any])     ->					Union[str, Any]:



													"""simple docstring"""






													pass







@is_pipeline_test
@require_vision
class     __lowerCAmelCase					(    unittest.TestCase		):



			@require_torch
			def 					_lowerCamelCase (		self     :						List[str])     ->					Tuple:



						"""simple docstring"""






						_UpperCAmelCase	 =     pipeline(
						    model='hf-internal-testing/tiny-random-clip-zero-shot-image-classification'				,	)
						_UpperCAmelCase	 =     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
						_UpperCAmelCase	 =     image_classifier(A				,	candidate_labels=['a', 'b', 'c'])

						# The floating scores are so close, we enter floating error approximation and the order is not guaranteed across
						# python and torch versions.
						self.assertIn(
						    nested_simplify(A)				,	[
						        [{'score': 0.3_3_3, 'label': 'a'}, {'score': 0.3_3_3, 'label': 'b'}, {'score': 0.3_3_3, 'label': 'c'}],
						        [{'score': 0.3_3_3, 'label': 'a'}, {'score': 0.3_3_3, 'label': 'c'}, {'score': 0.3_3_3, 'label': 'b'}],
						    ]				,	)

						_UpperCAmelCase	 =     image_classifier([image] * 5				,	candidate_labels=['A', 'B', 'C']				,	batch_size=2)
						self.assertEqual(
						    nested_simplify(A)				,	[
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						    ]				,	)



			@require_tf
			def 					_lowerCamelCase (		self     :						str)     ->					Tuple:



						"""simple docstring"""






						_UpperCAmelCase	 =     pipeline(
						    model='hf-internal-testing/tiny-random-clip-zero-shot-image-classification'				,	framework='tf')
						_UpperCAmelCase	 =     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
						_UpperCAmelCase	 =     image_classifier(A				,	candidate_labels=['a', 'b', 'c'])

						self.assertEqual(
						    nested_simplify(A)				,	[{'score': 0.3_3_3, 'label': 'a'}, {'score': 0.3_3_3, 'label': 'b'}, {'score': 0.3_3_3, 'label': 'c'}]				,	)

						_UpperCAmelCase	 =     image_classifier([image] * 5				,	candidate_labels=['A', 'B', 'C']				,	batch_size=2)
						self.assertEqual(
						    nested_simplify(A)				,	[
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						        [
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						            {'score': 0.3_3_3, 'label': ANY(A)},
						        ],
						    ]				,	)



			@slow
			@require_torch
			def 					_lowerCamelCase (		self     :						Tuple)     ->					Optional[Any]:



						"""simple docstring"""






						_UpperCAmelCase	 =     pipeline(
						    task='zero-shot-image-classification'				,	model='openai/clip-vit-base-patch32'				,	)
						# This is an image of 2 cats with remotes and no planes
						_UpperCAmelCase	 =     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
						_UpperCAmelCase	 =     image_classifier(A				,	candidate_labels=['cat', 'plane', 'remote'])

						self.assertEqual(
						    nested_simplify(A)				,	[
						        {'score': 0.5_1_1, 'label': 'remote'},
						        {'score': 0.4_8_5, 'label': 'cat'},
						        {'score': 0.0_0_4, 'label': 'plane'},
						    ]				,	)

						_UpperCAmelCase	 =     image_classifier([image] * 5				,	candidate_labels=['cat', 'plane', 'remote']				,	batch_size=2)
						self.assertEqual(
						    nested_simplify(A)				,	[
						        [
						            {'score': 0.5_1_1, 'label': 'remote'},
						            {'score': 0.4_8_5, 'label': 'cat'},
						            {'score': 0.0_0_4, 'label': 'plane'},
						        ],
						    ]
						    * 5				,	)



			@slow
			@require_tf
			def 					_lowerCamelCase (		self     :						List[str])     ->					Any:



						"""simple docstring"""






						_UpperCAmelCase	 =     pipeline(
						    task='zero-shot-image-classification'				,	model='openai/clip-vit-base-patch32'				,	framework='tf')
						# This is an image of 2 cats with remotes and no planes
						_UpperCAmelCase	 =     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
						_UpperCAmelCase	 =     image_classifier(A				,	candidate_labels=['cat', 'plane', 'remote'])
						self.assertEqual(
						    nested_simplify(A)				,	[
						        {'score': 0.5_1_1, 'label': 'remote'},
						        {'score': 0.4_8_5, 'label': 'cat'},
						        {'score': 0.0_0_4, 'label': 'plane'},
						    ]				,	)

						_UpperCAmelCase	 =     image_classifier([image] * 5				,	candidate_labels=['cat', 'plane', 'remote']				,	batch_size=2)
						self.assertEqual(
						    nested_simplify(A)				,	[
						        [
						            {'score': 0.5_1_1, 'label': 'remote'},
						            {'score': 0.4_8_5, 'label': 'cat'},
						            {'score': 0.0_0_4, 'label': 'plane'},
						        ],
						    ]
						    * 5				,	)
