






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     (							__lowerCAmelCase							:					int					,							__lowerCAmelCase							:					str					,							__lowerCAmelCase							:					Union[str, Any]					,							__lowerCAmelCase							:					Any					,							__lowerCAmelCase							:					Optional[Any]							)     ->  Optional[Any]:
   # Load configuration defined in the metadata file
   with open(__lowerCAmelCase							) as metadata_file:
      __UpperCamelCase	:			Any					=					json.load(__lowerCAmelCase							)
   __UpperCamelCase	:			int					=					LukeConfig(use_entity_aware_attention=__lowerCAmelCase					,							**metadata["""model_config"""]							)

   # Load in the weights from the checkpoint_path
   __UpperCamelCase	:			List[Any]					=					torch.load(__lowerCAmelCase					,							map_location="""cpu"""							)

   # Load the entity vocab file
   __UpperCamelCase	:			Dict					=					load_entity_vocab(__lowerCAmelCase							)

   __UpperCamelCase	:			List[Any]					=					RobertaTokenizer.from_pretrained(metadata["""model_config"""]["""bert_model_name"""]							)

   # Add special tokens to the token vocabulary for downstream tasks
   __UpperCamelCase	:			List[str]					=					AddedToken("""<ent>"""					,							lstrip=__lowerCAmelCase					,							rstrip=__lowerCAmelCase							)
   __UpperCamelCase	:			Optional[Any]					=					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					,							LukeTokenizer.vocab_files_names["""entity_vocab_file"""]							)					,							"""w"""							) as f:
      json.dump(__lowerCAmelCase					,							__lowerCAmelCase							)

   __UpperCamelCase	:			List[Any]					=					LukeTokenizer.from_pretrained(__lowerCAmelCase							)

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

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

   __UpperCamelCase	:			int					=					LukeModel(config=__lowerCAmelCase							).eval()

   __UpperCamelCase	,     __UpperCamelCase	:			Tuple					=					model.load_state_dict(__lowerCAmelCase					,							strict=__lowerCAmelCase							)
   if not (len(__lowerCAmelCase							) == 1 and missing_keys[0] == "embeddings.position_ids"):
      raise ValueError(f'Missing keys {", ".join(__lowerCAmelCase							)}. 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
   __UpperCamelCase	:			Optional[int]					=					LukeTokenizer.from_pretrained(__lowerCAmelCase					,							task="""entity_classification"""							)

   __UpperCamelCase	:			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 ."""
   )
   __UpperCamelCase	:			int					=					(39, 42)
   __UpperCamelCase	:			List[Any]					=					tokenizer(__lowerCAmelCase					,							entity_spans=[span]					,							add_prefix_space=__lowerCAmelCase					,							return_tensors="""pt"""							)

   __UpperCamelCase	:			Optional[int]					=					model(**__lowerCAmelCase							)

   # Verify word hidden states
   if model_size == "large":
      __UpperCamelCase	:			Union[str, Any]					=					torch.Size((1, 42, 1024)							)
      __UpperCamelCase	:			Dict					=					torch.tensor(
          [[0.0133, 0.0865, 0.0095], [0.3093, -0.2576, -0.7418], [-0.1720, -0.2117, -0.2869]]							)
   else:  # base
      __UpperCamelCase	:			Union[str, Any]					=					torch.Size((1, 42, 768)							)
      __UpperCamelCase	:			int					=					torch.tensor([[0.0037, 0.1368, -0.0091], [0.1099, 0.3329, -0.1095], [0.0765, 0.5335, 0.1179]]							)

   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":
      __UpperCamelCase	:			Optional[Any]					=					torch.Size((1, 1, 1024)							)
      __UpperCamelCase	:			int					=					torch.tensor([[0.0466, -0.0106, -0.0179]]							)
   else:  # base
      __UpperCamelCase	:			Dict					=					torch.Size((1, 1, 768)							)
      __UpperCamelCase	:			Dict					=					torch.tensor([[0.1457, 0.1044, 0.0174]]							)

   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

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


def 				__lowerCamelCase     (							__lowerCAmelCase							:					Optional[Any]							)     ->  Optional[Any]:
   __UpperCamelCase	:			Union[str, Any]					=					{}
   with open(__lowerCAmelCase					,							"""r"""					,							encoding="""utf-8"""							) as f:
      for index, line in enumerate(__lowerCAmelCase							):
         __UpperCamelCase	,     __UpperCamelCase	:			Dict					=					line.rstrip().split("""\t"""							)
         __UpperCamelCase	:			str					=					index

   return entity_vocab


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 re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


UpperCamelCase    							=      logging.get_logger(__name__)


def 				__lowerCamelCase     (							__lowerCAmelCase							:					Tuple							)     ->  List[str]:
   __UpperCamelCase	:			List[str]					=					R"""\w+[.]\d+"""
   __UpperCamelCase	:			Optional[int]					=					re.findall(__lowerCAmelCase					,							__lowerCAmelCase							)
   for pat in pats:
      __UpperCamelCase	:			str					=					key.replace(__lowerCAmelCase					,							"""_""".join(pat.split("""."""							)							)							)
   return key


def 				__lowerCamelCase     (							__lowerCAmelCase							:					int					,							__lowerCAmelCase							:					Optional[int]					,							__lowerCAmelCase							:					Union[str, Any]							)     ->  str:
   __UpperCamelCase	:			Dict					=					pt_tuple_key[:-1] + ("""scale""",)
   if (
       any("""norm""" in str_ for str_ in pt_tuple_key							)
       and (pt_tuple_key[-1] == "bias")
       and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
       and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
   ):
      __UpperCamelCase	:			List[Any]					=					pt_tuple_key[:-1] + ("""scale""",)
      return renamed_pt_tuple_key, pt_tensor
   elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
      __UpperCamelCase	:			Dict					=					pt_tuple_key[:-1] + ("""scale""",)
      return renamed_pt_tuple_key, pt_tensor

   # embedding
   if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
      __UpperCamelCase	:			Any					=					pt_tuple_key[:-1] + ("""embedding""",)
      return renamed_pt_tuple_key, pt_tensor

   # conv layer
   __UpperCamelCase	:			List[Any]					=					pt_tuple_key[:-1] + ("""kernel""",)
   if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
      __UpperCamelCase	:			Dict					=					pt_tensor.transpose(2					,							3					,							1					,							0							)
      return renamed_pt_tuple_key, pt_tensor

   # linear layer
   __UpperCamelCase	:			Union[str, Any]					=					pt_tuple_key[:-1] + ("""kernel""",)
   if pt_tuple_key[-1] == "weight":
      __UpperCamelCase	:			int					=					pt_tensor.T
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm weight
   __UpperCamelCase	:			List[Any]					=					pt_tuple_key[:-1] + ("""weight""",)
   if pt_tuple_key[-1] == "gamma":
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm bias
   __UpperCamelCase	:			Tuple					=					pt_tuple_key[:-1] + ("""bias""",)
   if pt_tuple_key[-1] == "beta":
      return renamed_pt_tuple_key, pt_tensor

   return pt_tuple_key, pt_tensor



def 				__lowerCamelCase     (							__lowerCAmelCase							:					Dict					,							__lowerCAmelCase							:					List[Any]					,							__lowerCAmelCase							:					Optional[int]=42							)     ->  int:
   # Step 1: Convert pytorch tensor to numpy
   __UpperCamelCase	:			Dict					=					{k: v.numpy() for k, v in pt_state_dict.items()}

   # Step 2: Since the model is stateless, get random Flax params
   __UpperCamelCase	:			str					=					flax_model.init_weights(PRNGKey(__lowerCAmelCase							)							)

   __UpperCamelCase	:			Union[str, Any]					=					flatten_dict(__lowerCAmelCase							)
   __UpperCamelCase	:			List[str]					=					{}

   # Need to change some parameters name to match Flax names
   for pt_key, pt_tensor in pt_state_dict.items():
      __UpperCamelCase	:			str					=					rename_key(__lowerCAmelCase							)
      __UpperCamelCase	:			Any					=					tuple(renamed_pt_key.split("""."""							)							)

      # Correctly rename weight parameters
      __UpperCamelCase	,     __UpperCamelCase	:			Tuple					=					rename_key_and_reshape_tensor(__lowerCAmelCase					,							__lowerCAmelCase					,							__lowerCAmelCase							)

      if flax_key in random_flax_state_dict:
         if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
            raise ValueError(
                f'PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '
                f'{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'							)

        # also add unexpected weight so that warning is thrown
      __UpperCamelCase	:			Any					=					jnp.asarray(__lowerCAmelCase							)

   return unflatten_dict(__lowerCAmelCase							)



