


"""simple docstring"""







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      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_     ):




     '''simple docstring'''


     with open(lowerCAmelCase_     ) as metadata_file:
          __SCREAMING_SNAKE_CASE		     =       json.load(lowerCAmelCase_     )
     __SCREAMING_SNAKE_CASE		     =       LukeConfig(use_entity_aware_attention=lowerCAmelCase_						,		**metadata["model_config"]     )

     # Load in the weights from the checkpoint_path
     __SCREAMING_SNAKE_CASE		     =       torch.load(lowerCAmelCase_						,		map_location="cpu"     )["module"]

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

     __SCREAMING_SNAKE_CASE		     =       XLMRobertaTokenizer.from_pretrained(metadata["model_config"]["bert_model_name"]     )

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

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

     __SCREAMING_SNAKE_CASE		     =       MLukeTokenizer.from_pretrained(lowerCAmelCase_     )

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

     __SCREAMING_SNAKE_CASE		     =       state_dict["embeddings.word_embeddings.weight"]
     __SCREAMING_SNAKE_CASE		     =       word_emb[ent_init_index].unsqueeze(0     )
     __SCREAMING_SNAKE_CASE		     =       word_emb[enta_init_index].unsqueeze(0     )
     __SCREAMING_SNAKE_CASE		     =       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"]:
          __SCREAMING_SNAKE_CASE		     =       state_dict[bias_name]
          __SCREAMING_SNAKE_CASE		     =       decoder_bias[ent_init_index].unsqueeze(0     )
          __SCREAMING_SNAKE_CASE		     =       decoder_bias[enta_init_index].unsqueeze(0     )
          __SCREAMING_SNAKE_CASE		     =       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"]:
               __SCREAMING_SNAKE_CASE		     =       f"""encoder.layer.{layer_index}.attention.self."""
               __SCREAMING_SNAKE_CASE		     =       state_dict[prefix + matrix_name]
               __SCREAMING_SNAKE_CASE		     =       state_dict[prefix + matrix_name]
               __SCREAMING_SNAKE_CASE		     =       state_dict[prefix + matrix_name]

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

     __SCREAMING_SNAKE_CASE		     =       LukeForMaskedLM(config=lowerCAmelCase_     ).eval()

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

     __SCREAMING_SNAKE_CASE     ,						__SCREAMING_SNAKE_CASE		     =       model.load_state_dict(lowerCAmelCase_						,		strict=lowerCAmelCase_     )

     if set(lowerCAmelCase_     ) != {"luke.embeddings.position_ids"}:
          raise ValueError(f"""Unexpected unexpected_keys: {unexpected_keys}"""     )
     if set(lowerCAmelCase_     ) != {
         "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
     __SCREAMING_SNAKE_CASE		     =       MLukeTokenizer.from_pretrained(lowerCAmelCase_						,		task="entity_classification"     )

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

     __SCREAMING_SNAKE_CASE		     =       model(**lowerCAmelCase_     )

     # Verify word hidden states
     if model_size == "large":
          raise NotImplementedError
     else:  # base
          __SCREAMING_SNAKE_CASE		     =       torch.Size((1, 33, 768)     )
          __SCREAMING_SNAKE_CASE		     =       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]						,		lowerCAmelCase_						,		atol=1E-4     ):
          raise ValueError

     # Verify entity hidden states
     if model_size == "large":
          raise NotImplementedError
     else:  # base
          __SCREAMING_SNAKE_CASE		     =       torch.Size((1, 1, 768)     )
          __SCREAMING_SNAKE_CASE		     =       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]						,		lowerCAmelCase_						,		atol=1E-4     ):
          raise ValueError

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

     __SCREAMING_SNAKE_CASE		     =       model(**lowerCAmelCase_     )

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

     __SCREAMING_SNAKE_CASE		     =       outputs.entity_logits[0][0].argmax().item()
     __SCREAMING_SNAKE_CASE		     =       [
         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(lowerCAmelCase_     )     )
     model.save_pretrained(lowerCAmelCase_     )




def      UpperCAmelCase__ (lowerCAmelCase_     ):




     '''simple docstring'''


     __SCREAMING_SNAKE_CASE		     =       ["[MASK]", "[PAD]", "[UNK]"]

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

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


if __name__ == "__main__":
  a__       :							Union[str, Any]        =  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__       :							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,
  )




"""simple docstring"""







import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
		from sagemaker import Session, TrainingJobAnalytics
		from sagemaker.huggingface import HuggingFace




@pytest.mark.skipif(
    literal_eval(os.getenv("TEST_SAGEMAKER"    ,						"False")) is not True    ,						reason="Skipping test because should only be run when releasing minor transformers version"    ,						)
@pytest.mark.usefixtures("sm_env")
@parameterized_class(
    [
        {
            "framework": "pytorch",
            "script": "run_glue.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 650, "eval_accuracy": 0.7, "eval_loss": 0.6},
        },
        {
            "framework": "pytorch",
            "script": "run_ddp.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 600, "eval_accuracy": 0.7, "eval_loss": 0.6},
        },
        {
            "framework": "tensorflow",
            "script": "run_tf_dist.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 600, "eval_accuracy": 0.6, "eval_loss": 0.7},
        },
    ])
class 				UpperCamelCase_			(				unittest.TestCase):

	"""simple docstring"""





	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			Any:
						if self.framework == "pytorch":
											subprocess.run(
											    F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()    ,    encoding="utf-8"    ,    check=UpperCAmelCase__    ,    )
						assert hasattr(self    ,    "env"     )




	def 						UpperCAmelCase_   (				self    :       Tuple    ,    UpperCAmelCase__    :       int     ) ->			Any:
						__SCREAMING_SNAKE_CASE		     =       F"""{self.env.base_job_name}-{instance_count}-{'ddp' if 'ddp' in self.script else 'smd'}"""
						# distributed data settings
						__SCREAMING_SNAKE_CASE		     =       {"smdistributed": {"dataparallel": {"enabled": True}}} if self.script != "run_ddp.py" else None

						# creates estimator
						return HuggingFace(
						    entry_point=self.script    ,    source_dir=self.env.test_path    ,    role=self.env.role    ,    image_uri=self.env.image_uri    ,    base_job_name=UpperCAmelCase__    ,    instance_count=UpperCAmelCase__    ,    instance_type=self.instance_type    ,    debugger_hook_config=UpperCAmelCase__    ,    hyperparameters={**self.env.distributed_hyperparameters, "model_name_or_path": self.model_name_or_path}    ,    metric_definitions=self.env.metric_definitions    ,    distribution=UpperCAmelCase__    ,    py_version="py36"    ,    )




	def 						UpperCAmelCase_   (				self    :       Optional[int]    ,    UpperCAmelCase__    :       Optional[int]     ) ->			Optional[Any]:
						TrainingJobAnalytics(UpperCAmelCase__     ).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv"""     )




	@parameterized.expand([(2,)]     )
	def 						UpperCAmelCase_   (				self    :       Dict    ,    UpperCAmelCase__    :       Optional[Any]     ) ->			List[str]:
						# create estimator
						__SCREAMING_SNAKE_CASE		     =       self.create_estimator(UpperCAmelCase__     )

						# run training
						estimator.fit()

						# result dataframe
						__SCREAMING_SNAKE_CASE		     =       TrainingJobAnalytics(estimator.latest_training_job.name     ).dataframe()

						# extract kpis
						__SCREAMING_SNAKE_CASE		     =       list(result_metrics_df[result_metrics_df.metric_name == "eval_accuracy"]["value"]     )
						__SCREAMING_SNAKE_CASE		     =       list(result_metrics_df[result_metrics_df.metric_name == "eval_loss"]["value"]     )
						# get train time from SageMaker job, this includes starting, preprocessing, stopping
						__SCREAMING_SNAKE_CASE		     =       (
						    Session().describe_training_job(estimator.latest_training_job.name     ).get("TrainingTimeInSeconds"    ,    9_9_9_9_9_9     )
						)

						# assert kpis
						assert train_runtime <= self.results["train_runtime"]
						assert all(t >= self.results["eval_accuracy"] for t in eval_accuracy     )
						assert all(t <= self.results["eval_loss"] for t in eval_loss     )

						# dump tests result into json file to share in PR
						with open(F"""{estimator.latest_training_job.name}.json"""    ,    "w"     ) as outfile:
											json.dump({"train_time": train_runtime, "eval_accuracy": eval_accuracy, "eval_loss": eval_loss}    ,    UpperCAmelCase__     )
