





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_			(	__a     ,   __a     ,   __a     ,   __a     ,   __a		)					->		int:







		"""simple docstring"""





		with open(lowercase__		) as metadata_file:
				lowerCamelCase__:				str        =json.load(lowercase__		)
		lowerCamelCase__:				Optional[int]        =LukeConfig(use_entity_aware_attention=lowercase__     ,   **metadata["model_config"]		)

		# Load in the weights from the checkpoint_path
		lowerCamelCase__:				Tuple        =torch.load(lowercase__     ,   map_location="cpu"		)

		# Load the entity vocab file
		lowerCamelCase__:				int        =load_entity_vocab(lowercase__		)

		lowerCamelCase__:				Optional[int]        =RobertaTokenizer.from_pretrained(metadata["model_config"]["bert_model_name"]		)

		# Add special tokens to the token vocabulary for downstream tasks
		lowerCamelCase__:				Tuple        =AddedToken("<ent>"     ,   lstrip=lowercase__     ,   rstrip=lowercase__		)
		lowerCamelCase__:				Optional[Any]        =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__:				Optional[int]        =LukeTokenizer.from_pretrained(lowercase__		)

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

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

		lowerCamelCase__:				List[str]        =LukeModel(config=lowercase__		).eval()

		lowerCamelCase__  , lowerCamelCase__:				Any        =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__:				List[Any]        =LukeTokenizer.from_pretrained(lowercase__     ,   task="entity_classification"		)

		lowerCamelCase__:				Optional[int]        =(
		    "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__:				List[str]        =(39, 42)
		lowerCamelCase__:				List[str]        =tokenizer(lowercase__     ,   entity_spans=[span]     ,   add_prefix_space=lowercase__     ,   return_tensors="pt"		)

		lowerCamelCase__:				Dict        =model(**lowercase__		)

		# Verify word hidden states
		if model_size == "large":
				lowerCamelCase__:				str        =torch.Size((1, 42, 1024)		)
				lowerCamelCase__:				Union[str, Any]        =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__:				Dict        =torch.Size((1, 42, 768)		)
				lowerCamelCase__:				List[str]        =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__:				Union[str, Any]        =torch.Size((1, 1, 1024)		)
				lowerCamelCase__:				Dict        =torch.tensor([[0.0_4_6_6, -0.0_1_0_6, -0.0_1_7_9]]		)
		else:  # base
				lowerCamelCase__:				Union[str, Any]        =torch.Size((1, 1, 768)		)
				lowerCamelCase__:				str        =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_			(	__a		)					->		Any:







		"""simple docstring"""





		lowerCamelCase__:				Dict        ={}
		with open(lowercase__     ,   "r"     ,   encoding="utf-8"		) as f:
				for index, line in enumerate(lowercase__		):
						lowerCamelCase__  , lowerCamelCase__:				List[Any]        =line.rstrip().split("\t"		)
						lowerCamelCase__:				str        =index

		return entity_vocab


if __name__ == "__main__":
					__A     =    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."
					)
					__A     =    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 shutil
import tempfile
import unittest
from unittest.mock import patch

from transformers import (
    DefaultFlowCallback,
    IntervalStrategy,
    PrinterCallback,
    ProgressCallback,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    is_torch_available,
)
from transformers.testing_utils import require_torch


if is_torch_available():
     from transformers.trainer import DEFAULT_CALLBACKS

     from .test_trainer import RegressionDataset, RegressionModelConfig, RegressionPreTrainedModel



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''

 def __init__(self    :       str)	->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:				List[Any]        =[]

 def 							SCREAMING_SNAKE_CASE_ (self    :       Any	,			UpperCAmelCase_    :       int	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       Dict	,			**UpperCAmelCase_    :       Any)	->Dict:



   '''simple docstring'''



   self.events.append("on_init_end")

 def 							SCREAMING_SNAKE_CASE_ (self    :       Union[str, Any]	,			UpperCAmelCase_    :       Any	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       Dict	,			**UpperCAmelCase_    :       str)	->List[str]:



   '''simple docstring'''



   self.events.append("on_train_begin")

 def 							SCREAMING_SNAKE_CASE_ (self    :       List[str]	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       Optional[Any]	,			UpperCAmelCase_    :       List[Any]	,			**UpperCAmelCase_    :       str)	->int:



   '''simple docstring'''



   self.events.append("on_train_end")

 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple	,			UpperCAmelCase_    :       Union[str, Any]	,			UpperCAmelCase_    :       Any	,			UpperCAmelCase_    :       List[str]	,			**UpperCAmelCase_    :       Optional[int])	->List[Any]:



   '''simple docstring'''



   self.events.append("on_epoch_begin")

 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[int]	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       int	,			**UpperCAmelCase_    :       Any)	->Tuple:



   '''simple docstring'''



   self.events.append("on_epoch_end")

 def 							SCREAMING_SNAKE_CASE_ (self    :       List[Any]	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       Any	,			UpperCAmelCase_    :       Any	,			**UpperCAmelCase_    :       List[Any])	->Optional[int]:



   '''simple docstring'''



   self.events.append("on_step_begin")

 def 							SCREAMING_SNAKE_CASE_ (self    :       List[str]	,			UpperCAmelCase_    :       List[str]	,			UpperCAmelCase_    :       List[str]	,			UpperCAmelCase_    :       Optional[int]	,			**UpperCAmelCase_    :       List[str])	->Tuple:



   '''simple docstring'''



   self.events.append("on_step_end")

 def 							SCREAMING_SNAKE_CASE_ (self    :       Union[str, Any]	,			UpperCAmelCase_    :       Optional[Any]	,			UpperCAmelCase_    :       Dict	,			UpperCAmelCase_    :       Optional[Any]	,			**UpperCAmelCase_    :       str)	->Optional[int]:



   '''simple docstring'''



   self.events.append("on_evaluate")

 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple	,			UpperCAmelCase_    :       List[str]	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       Union[str, Any]	,			**UpperCAmelCase_    :       Any)	->int:



   '''simple docstring'''



   self.events.append("on_predict")

 def 							SCREAMING_SNAKE_CASE_ (self    :       Union[str, Any]	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       Optional[int]	,			UpperCAmelCase_    :       Tuple	,			**UpperCAmelCase_    :       List[Any])	->Any:



   '''simple docstring'''



   self.events.append("on_save")

 def 							SCREAMING_SNAKE_CASE_ (self    :       str	,			UpperCAmelCase_    :       Any	,			UpperCAmelCase_    :       List[str]	,			UpperCAmelCase_    :       Any	,			**UpperCAmelCase_    :       Optional[Any])	->str:



   '''simple docstring'''



   self.events.append("on_log")







 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any]	,			UpperCAmelCase_    :       Any	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       str	,			**UpperCAmelCase_    :       List[str])	->Optional[int]:



   '''simple docstring'''



   self.events.append("on_prediction_step")



@require_torch
class        _SCREAMING_SNAKE_CASE						(     unittest.TestCase   ):







 '''simple docstring'''

 def 							SCREAMING_SNAKE_CASE_ (self    :       List[Any])	->int:



   '''simple docstring'''



   lowerCamelCase__:				Tuple        =tempfile.mkdtemp()

 def 							SCREAMING_SNAKE_CASE_ (self    :       List[str])	->Tuple:



   '''simple docstring'''



   shutil.rmtree(self.output_dir)

 def 							SCREAMING_SNAKE_CASE_ (self    :       List[Any]	,			UpperCAmelCase_    :       int=0	,			UpperCAmelCase_    :       Tuple=0	,			UpperCAmelCase_    :       str=64	,			UpperCAmelCase_    :       List[Any]=64	,			UpperCAmelCase_    :       Union[str, Any]=None	,			UpperCAmelCase_    :       List[str]=False	,			**UpperCAmelCase_    :       Tuple)	->Any:



   '''simple docstring'''



   lowerCamelCase__:				Union[str, Any]        =RegressionDataset(length=UpperCAmelCase_)
   lowerCamelCase__:				int        =RegressionDataset(length=UpperCAmelCase_)
   lowerCamelCase__:				str        =RegressionModelConfig(a=UpperCAmelCase_	,			b=UpperCAmelCase_)
   lowerCamelCase__:				Union[str, Any]        =RegressionPreTrainedModel(UpperCAmelCase_)

   lowerCamelCase__:				int        =TrainingArguments(self.output_dir	,			disable_tqdm=UpperCAmelCase_	,			report_to=[]	,			**UpperCAmelCase_)
   return Trainer(
       UpperCAmelCase_	,			UpperCAmelCase_	,			train_dataset=UpperCAmelCase_	,			eval_dataset=UpperCAmelCase_	,			callbacks=UpperCAmelCase_	,			)

 def 							SCREAMING_SNAKE_CASE_ (self    :       List[str]	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       List[str])	->Dict:



   '''simple docstring'''



   self.assertEqual(len(UpperCAmelCase_)	,			len(UpperCAmelCase_))

   # Order doesn't matter
   lowerCamelCase__:				Dict        =sorted(UpperCAmelCase_	,			key=lambda UpperCAmelCase_: cb.__name__ if isinstance(UpperCAmelCase_	,			UpperCAmelCase_) else cb.__class__.__name__)
   lowerCamelCase__:				Optional[int]        =sorted(UpperCAmelCase_	,			key=lambda UpperCAmelCase_: cb.__name__ if isinstance(UpperCAmelCase_	,			UpperCAmelCase_) else cb.__class__.__name__)

   for cba, cba in zip(UpperCAmelCase_	,			UpperCAmelCase_):
     if isinstance(UpperCAmelCase_	,			UpperCAmelCase_) and isinstance(UpperCAmelCase_	,			UpperCAmelCase_):
       self.assertEqual(UpperCAmelCase_	,			UpperCAmelCase_)
     elif isinstance(UpperCAmelCase_	,			UpperCAmelCase_) and not isinstance(UpperCAmelCase_	,			UpperCAmelCase_):
       self.assertEqual(UpperCAmelCase_	,			cba.__class__)
     elif not isinstance(UpperCAmelCase_	,			UpperCAmelCase_) and isinstance(UpperCAmelCase_	,			UpperCAmelCase_):
       self.assertEqual(cba.__class__	,			UpperCAmelCase_)
     else:
       self.assertEqual(UpperCAmelCase_	,			UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_ (self    :       List[Any]	,			UpperCAmelCase_    :       Optional[Any])	->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:				Dict        =["on_init_end", "on_train_begin"]
   lowerCamelCase__:				List[str]        =0
   lowerCamelCase__:				List[Any]        =len(trainer.get_eval_dataloader())
   lowerCamelCase__:				Dict        =["on_prediction_step"] * len(trainer.get_eval_dataloader()) + ["on_log", "on_evaluate"]
   for _ in range(trainer.state.num_train_epochs):
     expected_events.append("on_epoch_begin")
     for _ in range(UpperCAmelCase_):
       step += 1
       expected_events += ["on_step_begin", "on_step_end"]
       if step % trainer.args.logging_steps == 0:
         expected_events.append("on_log")
       if trainer.args.evaluation_strategy == IntervalStrategy.STEPS and step % trainer.args.eval_steps == 0:
         expected_events += evaluation_events.copy()
       if step % trainer.args.save_steps == 0:
         expected_events.append("on_save")
     expected_events.append("on_epoch_end")
     if trainer.args.evaluation_strategy == IntervalStrategy.EPOCH:
       expected_events += evaluation_events.copy()
   expected_events += ["on_log", "on_train_end"]
   return expected_events

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



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =self.get_trainer()
   lowerCamelCase__:				Any        =DEFAULT_CALLBACKS.copy() + [ProgressCallback]
   self.check_callbacks_equality(trainer.callback_handler.callbacks	,			UpperCAmelCase_)

   # Callbacks passed at init are added to the default callbacks
   lowerCamelCase__:				Dict        =self.get_trainer(callbacks=[MyTestTrainerCallback])
   expected_callbacks.append(UpperCAmelCase_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks	,			UpperCAmelCase_)

   # TrainingArguments.disable_tqdm controls if use ProgressCallback or PrinterCallback
   lowerCamelCase__:				int        =self.get_trainer(disable_tqdm=UpperCAmelCase_)
   lowerCamelCase__:				Tuple        =DEFAULT_CALLBACKS.copy() + [PrinterCallback]
   self.check_callbacks_equality(trainer.callback_handler.callbacks	,			UpperCAmelCase_)

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



   '''simple docstring'''



   lowerCamelCase__:				Union[str, Any]        =DEFAULT_CALLBACKS.copy() + [ProgressCallback]
   lowerCamelCase__:				Optional[int]        =self.get_trainer()

   # We can add, pop, or remove by class name
   trainer.remove_callback(UpperCAmelCase_)
   expected_callbacks.remove(UpperCAmelCase_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks	,			UpperCAmelCase_)

   lowerCamelCase__:				Dict        =self.get_trainer()
   lowerCamelCase__:				str        =trainer.pop_callback(UpperCAmelCase_)
   self.assertEqual(cb.__class__	,			UpperCAmelCase_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks	,			UpperCAmelCase_)

   trainer.add_callback(UpperCAmelCase_)
   expected_callbacks.insert(0	,			UpperCAmelCase_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks	,			UpperCAmelCase_)

   # We can also add, pop, or remove by instance
   lowerCamelCase__:				List[str]        =self.get_trainer()
   lowerCamelCase__:				List[str]        =trainer.callback_handler.callbacks[0]
   trainer.remove_callback(UpperCAmelCase_)
   expected_callbacks.remove(UpperCAmelCase_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks	,			UpperCAmelCase_)

   lowerCamelCase__:				str        =self.get_trainer()
   lowerCamelCase__:				List[Any]        =trainer.callback_handler.callbacks[0]
   lowerCamelCase__:				Dict        =trainer.pop_callback(UpperCAmelCase_)
   self.assertEqual(UpperCAmelCase_	,			UpperCAmelCase_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks	,			UpperCAmelCase_)

   trainer.add_callback(UpperCAmelCase_)
   expected_callbacks.insert(0	,			UpperCAmelCase_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks	,			UpperCAmelCase_)







 def 							SCREAMING_SNAKE_CASE_ (self    :       List[Any])	->int:



   '''simple docstring'''



   import warnings

   # XXX: for now ignore scatter_gather warnings in this test since it's not relevant to what's being tested
   warnings.simplefilter(action="ignore"	,			category=UpperCAmelCase_)

   lowerCamelCase__:				Union[str, Any]        =self.get_trainer(callbacks=[MyTestTrainerCallback])
   trainer.train()
   lowerCamelCase__:				int        =trainer.callback_handler.callbacks[-2].events
   self.assertEqual(UpperCAmelCase_	,			self.get_expected_events(UpperCAmelCase_))

   # Independent log/save/eval
   lowerCamelCase__:				Dict        =self.get_trainer(callbacks=[MyTestTrainerCallback]	,			logging_steps=5)
   trainer.train()
   lowerCamelCase__:				Optional[int]        =trainer.callback_handler.callbacks[-2].events
   self.assertEqual(UpperCAmelCase_	,			self.get_expected_events(UpperCAmelCase_))

   lowerCamelCase__:				Any        =self.get_trainer(callbacks=[MyTestTrainerCallback]	,			save_steps=5)
   trainer.train()
   lowerCamelCase__:				List[Any]        =trainer.callback_handler.callbacks[-2].events
   self.assertEqual(UpperCAmelCase_	,			self.get_expected_events(UpperCAmelCase_))

   lowerCamelCase__:				int        =self.get_trainer(callbacks=[MyTestTrainerCallback]	,			eval_steps=5	,			evaluation_strategy="steps")
   trainer.train()
   lowerCamelCase__:				str        =trainer.callback_handler.callbacks[-2].events
   self.assertEqual(UpperCAmelCase_	,			self.get_expected_events(UpperCAmelCase_))

   lowerCamelCase__:				Dict        =self.get_trainer(callbacks=[MyTestTrainerCallback]	,			evaluation_strategy="epoch")
   trainer.train()
   lowerCamelCase__:				Tuple        =trainer.callback_handler.callbacks[-2].events
   self.assertEqual(UpperCAmelCase_	,			self.get_expected_events(UpperCAmelCase_))

   # A bit of everything
   lowerCamelCase__:				Tuple        =self.get_trainer(
       callbacks=[MyTestTrainerCallback]	,			logging_steps=3	,			save_steps=10	,			eval_steps=5	,			evaluation_strategy="steps"	,			)
   trainer.train()
   lowerCamelCase__:				int        =trainer.callback_handler.callbacks[-2].events
   self.assertEqual(UpperCAmelCase_	,			self.get_expected_events(UpperCAmelCase_))

   # warning should be emitted for duplicated callbacks
   with patch("transformers.trainer_callback.logger.warning") as warn_mock:
     lowerCamelCase__:				Optional[int]        =self.get_trainer(
         callbacks=[MyTestTrainerCallback, MyTestTrainerCallback]	,			)
     assert str(UpperCAmelCase_) in warn_mock.call_args[0][0]

