





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     SCREAMING_SNAKE_CASE			(      lowerCAmelCase    ,       lowerCAmelCase    ,       lowerCAmelCase    ,       lowerCAmelCase    ,       lowerCAmelCase					):
       # Load configuration defined in the metadata file
       with open(lowerCAmelCase					) as metadata_file:
              _UpperCamelCase              =						json.load(lowerCAmelCase					)
       _UpperCamelCase              =						LukeConfig(use_entity_aware_attention=lowerCAmelCase    ,       **metadata['''model_config''']					)

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

       # Load the entity vocab file
       _UpperCamelCase              =						load_original_entity_vocab(lowerCAmelCase					)
       # 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=lowerCAmelCase    ,       rstrip=lowerCAmelCase					)
       _UpperCamelCase              =						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:
              _UpperCamelCase              =						json.load(lowerCAmelCase					)
       _UpperCamelCase              =						'''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					)

       _UpperCamelCase              =						MLukeTokenizer.from_pretrained(lowerCAmelCase					)

       # 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=lowerCAmelCase					).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(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
       _UpperCamelCase              =						MLukeTokenizer.from_pretrained(lowerCAmelCase    ,       task='''entity_classification'''					)

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

       _UpperCamelCase              =						model(**lowerCAmelCase					)

       # Verify word hidden states
       if model_size == "large":
              raise NotImplementedError
       else:  # base
              _UpperCamelCase              =						torch.Size((1, 33, 768)					)
              _UpperCamelCase              =						torch.tensor([[0.0_8_9_2, 0.0_5_9_6, -0.2_8_1_9], [0.0_1_3_4, 0.1_1_9_9, 0.0_5_7_3], [-0.0_1_6_9, 0.0_9_2_7, 0.0_6_4_4]]					)

       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
              _UpperCamelCase              =						torch.Size((1, 1, 768)					)
              _UpperCamelCase              =						torch.tensor([[-0.1_4_8_2, 0.0_6_0_9, 0.0_3_2_2]]					)

       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
       _UpperCamelCase              =						MLukeTokenizer.from_pretrained(lowerCAmelCase					)
       _UpperCamelCase              =						'''Tokyo is the capital of <mask>.'''
       _UpperCamelCase              =						(24, 30)
       _UpperCamelCase              =						tokenizer(lowerCAmelCase    ,       entity_spans=[span]    ,       return_tensors='''pt'''					)

       _UpperCamelCase              =						model(**lowerCAmelCase					)

       _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(lowerCAmelCase					)

       _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(lowerCAmelCase					)					)
       model.save_pretrained(lowerCAmelCase					)


def     SCREAMING_SNAKE_CASE			(      lowerCAmelCase					):
       _UpperCamelCase              =						['''[MASK]''', '''[PAD]''', '''[UNK]''']

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

       _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__":
       lowercase    : 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."""
       )
       lowercase    : Union[str, Any]    	=  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 datasets

from .evaluate import evaluate


lowercase    : str    	=  """\
@article{hendrycks2021cuad,
      title={CUAD: An Expert-Annotated NLP Dataset for Legal Contract Review},
      author={Dan Hendrycks and Collin Burns and Anya Chen and Spencer Ball},
      journal={arXiv preprint arXiv:2103.06268},
      year={2021}
}
"""

lowercase    : Tuple    	=  """
This metric wrap the official scoring script for version 1 of the Contract
Understanding Atticus Dataset (CUAD).
Contract Understanding Atticus Dataset (CUAD) v1 is a corpus of more than 13,000 labels in 510
commercial legal contracts that have been manually labeled to identify 41 categories of important
clauses that lawyers look for when reviewing contracts in connection with corporate transactions.
"""

lowercase    : List[str]    	=  """
Computes CUAD scores (EM, F1, AUPR, Precision@80%Recall, and Precision@90%Recall).
Args:
    predictions: List of question-answers dictionaries with the following key-values:
        - 'id': id of the question-answer pair as given in the references (see below)
        - 'prediction_text': list of possible texts for the answer, as a list of strings
        depending on a threshold on the confidence probability of each prediction.
    references: List of question-answers dictionaries with the following key-values:
        - 'id': id of the question-answer pair (see above),
        - 'answers': a Dict in the CUAD dataset format
            {
                'text': list of possible texts for the answer, as a list of strings
                'answer_start': list of start positions for the answer, as a list of ints
            }
            Note that answer_start values are not taken into account to compute the metric.
Returns:
    'exact_match': Exact match (the normalized answer exactly match the gold answer)
    'f1': The F-score of predicted tokens versus the gold answer
    'aupr': Area Under the Precision-Recall curve
    'prec_at_80_recall': Precision at 80% recall
    'prec_at_90_recall': Precision at 90% recall
Examples:
    >>> predictions = [{'prediction_text': ['The seller:', 'The buyer/End-User: Shenzhen LOHAS Supply Chain Management Co., Ltd.'], 'id': 'LohaCompanyltd_20191209_F-1_EX-10.16_11917878_EX-10.16_Supply Agreement__Parties'}]
    >>> references = [{'answers': {'answer_start': [143, 49], 'text': ['The seller:', 'The buyer/End-User: Shenzhen LOHAS Supply Chain Management Co., Ltd.']}, 'id': 'LohaCompanyltd_20191209_F-1_EX-10.16_11917878_EX-10.16_Supply Agreement__Parties'}]
    >>> cuad_metric = datasets.load_metric(\"cuad\")
    >>> results = cuad_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'exact_match': 100.0, 'f1': 100.0, 'aupr': 0.0, 'prec_at_80_recall': 1.0, 'prec_at_90_recall': 1.0}
"""

@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION						,			_KWARGS_DESCRIPTION )
class  __A(						datasets.Metric ):

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            return datasets.MetricInfo(
                description=_DESCRIPTION,       citation=_CITATION,       inputs_description=_KWARGS_DESCRIPTION,       features=datasets.Features(
                    {
                        '''predictions''': {
                            '''id''': datasets.Value('''string'''      ),
                            '''prediction_text''': datasets.features.Sequence(datasets.Value('''string'''      )      ),
                        },
                        '''references''': {
                            '''id''': datasets.Value('''string'''      ),
                            '''answers''': datasets.features.Sequence(
                                {
                                    '''text''': datasets.Value('''string'''      ),
                                    '''answer_start''': datasets.Value('''int32'''      ),
                                }      ),
                        },
                    }      ),       codebase_urls=['''https://www.atticusprojectai.org/cuad'''],       reference_urls=['''https://www.atticusprojectai.org/cuad'''],       )



     def 	_UpperCamelCase     (      self,       A,       A      ):





            """simple docstring"""

            _UpperCamelCase              =						{prediction['''id''']: prediction['''prediction_text'''] for prediction in predictions}
            _UpperCamelCase              =						[
                {
                    '''paragraphs''': [
                        {
                            '''qas''': [
                                {
                                    '''answers''': [{'''text''': answer_text} for answer_text in ref['''answers''']['''text''']],
                                    '''id''': ref['''id'''],
                                }
                                for ref in references
                            ]
                        }
                    ]
                }
            ]
            _UpperCamelCase              =						evaluate(dataset=A,       predictions=A      )
            return score






