


import datasets

from .evaluate import evaluate


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

A      :    Any		=					"\nThis metric wrap the official scoring script for version 1 of the Contract\nUnderstanding Atticus Dataset (CUAD).\nContract Understanding Atticus Dataset (CUAD) v1 is a corpus of more than 13,000 labels in 510\ncommercial legal contracts that have been manually labeled to identify 41 categories of important\nclauses that lawyers look for when reviewing contracts in connection with corporate transactions.\n"

A      :    str		=					"\nComputes CUAD scores (EM, F1, AUPR, Precision@80%Recall, and Precision@90%Recall).\nArgs:\n    predictions: List of question-answers dictionaries with the following key-values:\n        - 'id': id of the question-answer pair as given in the references (see below)\n        - 'prediction_text': list of possible texts for the answer, as a list of strings\n        depending on a threshold on the confidence probability of each prediction.\n    references: List of question-answers dictionaries with the following key-values:\n        - 'id': id of the question-answer pair (see above),\n        - 'answers': a Dict in the CUAD dataset format\n            {\n                'text': list of possible texts for the answer, as a list of strings\n                'answer_start': list of start positions for the answer, as a list of ints\n            }\n            Note that answer_start values are not taken into account to compute the metric.\nReturns:\n    'exact_match': Exact match (the normalized answer exactly match the gold answer)\n    'f1': The F-score of predicted tokens versus the gold answer\n    'aupr': Area Under the Precision-Recall curve\n    'prec_at_80_recall': Precision at 80% recall\n    'prec_at_90_recall': Precision at 90% recall\nExamples:\n    >>> 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'}]\n    >>> 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'}]\n    >>> cuad_metric = datasets.load_metric(\"cuad\")\n    >>> results = cuad_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'exact_match': 100.0, 'f1': 100.0, 'aupr': 0.0, 'prec_at_80_recall': 1.0, 'prec_at_90_recall': 1.0}\n"







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


       """simple docstring"""




       def 						lowerCAmelCase     (   self							:			Optional[Any]       ):


             '''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 						lowerCAmelCase     (   self							:			Union[str, Any]					,      __lowerCamelCase							:			List[str]					,      __lowerCamelCase							:			Dict       ):


             '''simple docstring'''



             lowerCamelCase__		:				int    =  {prediction["id"]: prediction["prediction_text"] for prediction in predictions}
             lowerCamelCase__		:				Any    =  [
                 {
                     "paragraphs": [
                         {
                             "qas": [
                                 {
                                     "answers": [{"text": answer_text} for answer_text in ref["answers"]["text"]],
                                     "id": ref["id"],
                                 }
                                 for ref in references
                             ]
                         }
                     ]
                 }
             ]
             lowerCamelCase__		:				Optional[int]    =  evaluate(dataset=__lowerCamelCase					,      predictions=__lowerCamelCase       )
             return score






import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
    import jax
    import jax.numpy as jnp

    from transformers.modeling_flax_pytorch_utils import (
        convert_pytorch_state_dict_to_flax,
        load_flax_weights_in_pytorch_model,
    )
    from transformers.models.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
    import torch







class 			_lowercase :


       """simple docstring"""




       def __init__(   self							:			List[str]					,      __lowerCamelCase							:			Tuple					,      __lowerCamelCase							:			Union[str, Any]=14					,      __lowerCamelCase							:			str=7					,      __lowerCamelCase							:			Tuple=True					,      __lowerCamelCase							:			Tuple=True					,      __lowerCamelCase							:			Optional[int]=False					,      __lowerCamelCase							:			Tuple=True					,      __lowerCamelCase							:			List[Any]=99					,      __lowerCamelCase							:			Optional[Any]=32					,      __lowerCamelCase							:			List[Any]=4					,      __lowerCamelCase							:			List[Any]=4					,      __lowerCamelCase							:			List[str]=4					,      __lowerCamelCase							:			List[Any]=37					,      __lowerCamelCase							:			str="gelu"					,      __lowerCamelCase							:			Tuple=0.1					,      __lowerCamelCase							:			Any=0.1					,      __lowerCamelCase							:			Any=512					,      __lowerCamelCase							:			Dict=0.0_2					,      ):


             '''simple docstring'''



             lowerCamelCase__		:				int    =  parent
             lowerCamelCase__		:				Any    =  batch_size
             lowerCamelCase__		:				Tuple    =  seq_length
             lowerCamelCase__		:				List[str]    =  is_training
             lowerCamelCase__		:				List[Any]    =  use_input_mask
             lowerCamelCase__		:				Optional[Any]    =  use_token_type_ids
             lowerCamelCase__		:				Optional[Any]    =  use_labels
             lowerCamelCase__		:				List[Any]    =  vocab_size
             lowerCamelCase__		:				Any    =  hidden_size
             lowerCamelCase__		:				int    =  rotary_dim
             lowerCamelCase__		:				List[str]    =  num_hidden_layers
             lowerCamelCase__		:				int    =  num_attention_heads
             lowerCamelCase__		:				Optional[int]    =  intermediate_size
             lowerCamelCase__		:				Tuple    =  hidden_act
             lowerCamelCase__		:				Optional[int]    =  hidden_dropout_prob
             lowerCamelCase__		:				Optional[Any]    =  attention_probs_dropout_prob
             lowerCamelCase__		:				str    =  max_position_embeddings
             lowerCamelCase__		:				Tuple    =  initializer_range
             lowerCamelCase__		:				Any    =  None
             lowerCamelCase__		:				Optional[Any]    =  vocab_size - 1
             lowerCamelCase__		:				List[Any]    =  vocab_size - 1
             lowerCamelCase__		:				str    =  vocab_size - 1





       def 						lowerCAmelCase     (   self							:			Optional[int]       ):


             '''simple docstring'''



             lowerCamelCase__		:				Optional[int]    =  ids_tensor([self.batch_size, self.seq_length]					,      self.vocab_size       )

             lowerCamelCase__		:				List[str]    =  None
             if self.use_input_mask:
                   lowerCamelCase__		:				Any    =  random_attention_mask([self.batch_size, self.seq_length]       )

             lowerCamelCase__		:				Dict    =  GPTJConfig(
                 vocab_size=self.vocab_size					,      n_embd=self.hidden_size					,      n_layer=self.num_hidden_layers					,      n_head=self.num_attention_heads					,      n_positions=self.max_position_embeddings					,      use_cache=__lowerCamelCase					,      bos_token_id=self.bos_token_id					,      eos_token_id=self.eos_token_id					,      pad_token_id=self.pad_token_id					,      rotary_dim=self.rotary_dim					,      )

             return (config, input_ids, input_mask)





       def 						lowerCAmelCase     (   self							:			str       ):


             '''simple docstring'''



             lowerCamelCase__		:				Tuple    =  self.prepare_config_and_inputs()
             lowerCamelCase__    ,   lowerCamelCase__    ,   lowerCamelCase__		:				Tuple    =  config_and_inputs
             lowerCamelCase__		:				Tuple    =  {"input_ids": input_ids, "attention_mask": attention_mask}
             return config, inputs_dict





       def 						lowerCAmelCase     (   self							:			Tuple					,      __lowerCamelCase							:			Optional[int]					,      __lowerCamelCase							:			str					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			Optional[Any]       ):


             '''simple docstring'''



             lowerCamelCase__		:				str    =  20
             lowerCamelCase__		:				Tuple    =  model_class_name(__lowerCamelCase       )

             lowerCamelCase__		:				Union[str, Any]    =  model.init_cache(input_ids.shape[0]					,      __lowerCamelCase       )
             lowerCamelCase__		:				List[str]    =  jnp.ones((input_ids.shape[0], max_decoder_length)					,      dtype="i4"       )

             lowerCamelCase__		:				str    =  jnp.broadcast_to(
                 jnp.arange(input_ids.shape[-1] - 1       )[None, :]					,      (input_ids.shape[0], input_ids.shape[-1] - 1)       )
             lowerCamelCase__		:				Tuple    =  model(
                 input_ids[:, :-1]					,      attention_mask=__lowerCamelCase					,      past_key_values=__lowerCamelCase					,      position_ids=__lowerCamelCase					,      )

             lowerCamelCase__		:				Tuple    =  jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]					,      dtype="i4"       )
             lowerCamelCase__		:				str    =  model(
                 input_ids[:, -1:]					,      attention_mask=__lowerCamelCase					,      past_key_values=outputs_cache.past_key_values					,      position_ids=__lowerCamelCase					,      )

             lowerCamelCase__		:				List[Any]    =  model(__lowerCamelCase       )

             lowerCamelCase__		:				int    =  np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])       )       )
             self.parent.assertTrue(diff < 1E-3					,      msg=f"Max diff is {diff}"       )





       def 						lowerCAmelCase     (   self							:			Union[str, Any]					,      __lowerCamelCase							:			Any					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			str       ):


             '''simple docstring'''



             lowerCamelCase__		:				int    =  20
             lowerCamelCase__		:				int    =  model_class_name(__lowerCamelCase       )

             lowerCamelCase__		:				Dict    =  jnp.concatenate(
                 [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])       )]					,      axis=-1					,      )

             lowerCamelCase__		:				Dict    =  model.init_cache(input_ids.shape[0]					,      __lowerCamelCase       )
             lowerCamelCase__		:				Any    =  jnp.broadcast_to(
                 jnp.arange(input_ids.shape[-1] - 1       )[None, :]					,      (input_ids.shape[0], input_ids.shape[-1] - 1)       )

             lowerCamelCase__		:				List[Any]    =  model(
                 input_ids[:, :-1]					,      attention_mask=__lowerCamelCase					,      past_key_values=__lowerCamelCase					,      position_ids=__lowerCamelCase					,      )
             lowerCamelCase__		:				Union[str, Any]    =  jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]					,      dtype="i4"       )
             lowerCamelCase__		:				Any    =  model(
                 input_ids[:, -1:]					,      past_key_values=outputs_cache.past_key_values					,      attention_mask=__lowerCamelCase					,      position_ids=__lowerCamelCase					,      )

             lowerCamelCase__		:				Tuple    =  model(__lowerCamelCase					,      attention_mask=__lowerCamelCase       )

             lowerCamelCase__		:				Dict    =  np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])       )       )
             self.parent.assertTrue(diff < 1E-3					,      msg=f"Max diff is {diff}"       )







@require_flax
class 			_lowercase (							lowercase__							,  lowercase__							,  unittest.TestCase):


       """simple docstring"""




       A__											=			(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
       A__											=			(FlaxGPTJForCausalLM,) if is_flax_available() else ()





       def 						lowerCAmelCase     (   self							:			int       ):


             '''simple docstring'''



             lowerCamelCase__		:				Any    =  FlaxGPTJModelTester(self       )





       def 						lowerCAmelCase     (   self							:			str       ):


             '''simple docstring'''



             for model_class_name in self.all_model_classes:
                   lowerCamelCase__    ,   lowerCamelCase__    ,   lowerCamelCase__		:				Any    =  self.model_tester.prepare_config_and_inputs()
                   self.model_tester.check_use_cache_forward(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase       )





       def 						lowerCAmelCase     (   self							:			Any       ):


             '''simple docstring'''



             for model_class_name in self.all_model_classes:
                   lowerCamelCase__    ,   lowerCamelCase__    ,   lowerCamelCase__		:				List[Any]    =  self.model_tester.prepare_config_and_inputs()
                   self.model_tester.check_use_cache_forward_with_attn_mask(
                       __lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase       )





       @tooslow
       def 						lowerCAmelCase     (   self							:			str       ):


             '''simple docstring'''



             lowerCamelCase__		:				List[Any]    =  GPTaTokenizer.from_pretrained("gpt2"					,      pad_token="<|endoftext|>"					,      padding_side="left"       )
             lowerCamelCase__		:				List[str]    =  tokenizer(["Hello this is a long string", "Hey"]					,      return_tensors="np"					,      padding=__lowerCamelCase					,      truncation=__lowerCamelCase       )

             lowerCamelCase__		:				Any    =  FlaxGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"       )
             lowerCamelCase__		:				Optional[Any]    =  False
             lowerCamelCase__		:				str    =  model.config.eos_token_id

             lowerCamelCase__		:				int    =  jax.jit(model.generate       )

             lowerCamelCase__		:				str    =  jit_generate(
                 inputs["input_ids"]					,      attention_mask=inputs["attention_mask"]					,      pad_token_id=tokenizer.pad_token_id       ).sequences

             lowerCamelCase__		:				Tuple    =  tokenizer.batch_decode(__lowerCamelCase					,      skip_special_tokens=__lowerCamelCase       )

             lowerCamelCase__		:				Any    =  [
                 "Hello this is a long string of text.\n\nI'm trying to get the text of the",
                 "Hey, I'm a little late to the party. I'm going to",
             ]

             self.assertListEqual(__lowerCamelCase					,      __lowerCamelCase       )





       @is_pt_flax_cross_test
       def 						lowerCAmelCase     (   self							:			Optional[int]       ):


             '''simple docstring'''



             lowerCamelCase__    ,   lowerCamelCase__		:				Union[str, Any]    =  self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   with self.subTest(model_class.__name__       ):
                         # prepare inputs
                         lowerCamelCase__		:				Optional[Any]    =  self._prepare_for_class(__lowerCamelCase					,      __lowerCamelCase       )
                         lowerCamelCase__		:				int    =  {k: torch.tensor(v.tolist()       ) for k, v in prepared_inputs_dict.items()}

                         # load corresponding PyTorch class
                         lowerCamelCase__		:				List[str]    =  model_class.__name__[4:]  # Skip the "Flax" at the beginning
                         lowerCamelCase__		:				int    =  getattr(__lowerCamelCase					,      __lowerCamelCase       )

                         lowerCamelCase__    ,   lowerCamelCase__		:				Optional[int]    =  pt_inputs["input_ids"].shape
                         lowerCamelCase__		:				List[str]    =  np.random.randint(0					,      seq_length - 1					,      size=(batch_size,)       )
                         for batch_idx, start_index in enumerate(__lowerCamelCase       ):
                               lowerCamelCase__		:				Tuple    =  0
                               lowerCamelCase__		:				List[str]    =  1
                               lowerCamelCase__		:				Optional[Any]    =  0
                               lowerCamelCase__		:				Any    =  1
                         lowerCamelCase__		:				List[Any]    =  pt_model_class(__lowerCamelCase       ).eval()
                         lowerCamelCase__		:				List[str]    =  model_class(__lowerCamelCase					,      dtype=jnp.floataa       )

                         lowerCamelCase__		:				Optional[int]    =  convert_pytorch_state_dict_to_flax(pt_model.state_dict()					,      __lowerCamelCase       )
                         lowerCamelCase__		:				Optional[int]    =  fx_state

                         with torch.no_grad():
                               lowerCamelCase__		:				List[Any]    =  pt_model(**__lowerCamelCase       ).to_tuple()

                         lowerCamelCase__		:				Union[str, Any]    =  fx_model(**__lowerCamelCase       ).to_tuple()
                         self.assertEqual(len(__lowerCamelCase       )					,      len(__lowerCamelCase       )					,      "Output lengths differ between Flax and PyTorch"       )
                         for fx_output, pt_output in zip(__lowerCamelCase					,      __lowerCamelCase       ):
                               self.assert_almost_equals(fx_output[:, -1]					,      pt_output[:, -1].numpy()					,      4E-2       )

                         with tempfile.TemporaryDirectory() as tmpdirname:
                               pt_model.save_pretrained(__lowerCamelCase       )
                               lowerCamelCase__		:				List[Any]    =  model_class.from_pretrained(__lowerCamelCase					,      from_pt=__lowerCamelCase       )

                         lowerCamelCase__		:				int    =  fx_model_loaded(**__lowerCamelCase       ).to_tuple()
                         self.assertEqual(
                             len(__lowerCamelCase       )					,      len(__lowerCamelCase       )					,      "Output lengths differ between Flax and PyTorch"       )
                         for fx_output_loaded, pt_output in zip(__lowerCamelCase					,      __lowerCamelCase       ):
                               self.assert_almost_equals(fx_output_loaded[:, -1]					,      pt_output[:, -1].numpy()					,      4E-2       )





       @is_pt_flax_cross_test
       def 						lowerCAmelCase     (   self							:			str       ):


             '''simple docstring'''



             lowerCamelCase__    ,   lowerCamelCase__		:				Tuple    =  self.model_tester.prepare_config_and_inputs_for_common()
             for model_class in self.all_model_classes:
                   with self.subTest(model_class.__name__       ):
                         # prepare inputs
                         lowerCamelCase__		:				Union[str, Any]    =  self._prepare_for_class(__lowerCamelCase					,      __lowerCamelCase       )
                         lowerCamelCase__		:				Union[str, Any]    =  {k: torch.tensor(v.tolist()       ) for k, v in prepared_inputs_dict.items()}

                         # load corresponding PyTorch class
                         lowerCamelCase__		:				Union[str, Any]    =  model_class.__name__[4:]  # Skip the "Flax" at the beginning
                         lowerCamelCase__		:				List[Any]    =  getattr(__lowerCamelCase					,      __lowerCamelCase       )

                         lowerCamelCase__		:				Optional[Any]    =  pt_model_class(__lowerCamelCase       ).eval()
                         lowerCamelCase__		:				Tuple    =  model_class(__lowerCamelCase					,      dtype=jnp.floataa       )

                         lowerCamelCase__		:				Optional[Any]    =  load_flax_weights_in_pytorch_model(__lowerCamelCase					,      fx_model.params       )
                         lowerCamelCase__    ,   lowerCamelCase__		:				Optional[Any]    =  pt_inputs["input_ids"].shape
                         lowerCamelCase__		:				Optional[int]    =  np.random.randint(0					,      seq_length - 1					,      size=(batch_size,)       )
                         for batch_idx, start_index in enumerate(__lowerCamelCase       ):
                               lowerCamelCase__		:				Any    =  0
                               lowerCamelCase__		:				Any    =  1
                               lowerCamelCase__		:				Union[str, Any]    =  0
                               lowerCamelCase__		:				str    =  1

                         # make sure weights are tied in PyTorch
                         pt_model.tie_weights()

                         with torch.no_grad():
                               lowerCamelCase__		:				Dict    =  pt_model(**__lowerCamelCase       ).to_tuple()

                         lowerCamelCase__		:				List[str]    =  fx_model(**__lowerCamelCase       ).to_tuple()
                         self.assertEqual(len(__lowerCamelCase       )					,      len(__lowerCamelCase       )					,      "Output lengths differ between Flax and PyTorch"       )
                         for fx_output, pt_output in zip(__lowerCamelCase					,      __lowerCamelCase       ):
                               self.assert_almost_equals(fx_output[:, -1]					,      pt_output[:, -1].numpy()					,      4E-2       )

                         with tempfile.TemporaryDirectory() as tmpdirname:
                               fx_model.save_pretrained(__lowerCamelCase       )
                               lowerCamelCase__		:				Any    =  pt_model_class.from_pretrained(__lowerCamelCase					,      from_flax=__lowerCamelCase       )

                         with torch.no_grad():
                               lowerCamelCase__		:				Optional[int]    =  pt_model_loaded(**__lowerCamelCase       ).to_tuple()

                         self.assertEqual(
                             len(__lowerCamelCase       )					,      len(__lowerCamelCase       )					,      "Output lengths differ between Flax and PyTorch"       )
                         for fx_output, pt_output in zip(__lowerCamelCase					,      __lowerCamelCase       ):
                               self.assert_almost_equals(fx_output[:, -1]					,      pt_output[:, -1].numpy()					,      4E-2       )





       @tooslow
       def 						lowerCAmelCase     (   self							:			List[str]       ):


             '''simple docstring'''



             for model_class_name in self.all_model_classes:
                   lowerCamelCase__		:				Union[str, Any]    =  model_class_name.from_pretrained("EleutherAI/gpt-j-6B"       )
                   lowerCamelCase__		:				List[str]    =  model(np.ones((1, 1)       )       )
                   self.assertIsNotNone(__lowerCamelCase       )


