





'''simple docstring'''

import datasets

from .evaluate import evaluate


lowercase__		:     int			=							'\\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'

lowercase__		:     List[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'

lowercase__		:     int			=							'\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 							__lowerCAmelCase	(       datasets.Metric					):






       """simple docstring"""





       def snake_case__		(				self  :						str	)       ->	Union[str, 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 snake_case__		(				self  :						Dict			,					lowerCAmelCase__  :						Any			,					lowerCAmelCase__  :						List[str]	)       ->	Union[str, Any]:

            '''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=lowerCAmelCase__			,					predictions=lowerCAmelCase__	)
            return score










'''simple docstring'''






import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

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


if is_flax_available():
					import os

					# The slow tests are often failing with OOM error on GPU
					# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
					# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
					lowerCAmelCase									=   """platform"""

					import jax
					import jax.numpy as jnp

					from transformers import BlenderbotTokenizer
					from transformers.models.blenderbot.modeling_flax_blenderbot import (
					    FlaxBlenderbotForConditionalGeneration,
					    FlaxBlenderbotModel,
					    shift_tokens_right,
					)






def   __A    (				a_    :	Dict   ,a_    :	List[Any]   ,a_    :	List[str]=None   ,a_    :	Optional[int]=None   ,a_    :	Any=None   ,a_    :	Any=None   ,a_    :	str=None   ,a_    :	Union[str, Any]=None   ,):
		if attention_mask is None:
				lowerCAmelCase  :     List[Any]       						= np.where(input_ids != config.pad_token_id   ,1   ,0						)
		if decoder_attention_mask is None:
				lowerCAmelCase  :     Union[str, Any]       						= np.where(decoder_input_ids != config.pad_token_id   ,1   ,0						)
		if head_mask is None:
				lowerCAmelCase  :     Dict       						= np.ones((config.encoder_layers, config.encoder_attention_heads)						)
		if decoder_head_mask is None:
				lowerCAmelCase  :     int       						= np.ones((config.decoder_layers, config.decoder_attention_heads)						)
		if cross_attn_head_mask is None:
				lowerCAmelCase  :     Union[str, Any]       						= np.ones((config.decoder_layers, config.decoder_attention_heads)						)
		return {
		    "input_ids": input_ids,
		    "decoder_input_ids": decoder_input_ids,
		    "attention_mask": attention_mask,
		    "decoder_attention_mask": attention_mask,
		}
class 		lowerCamelCase	:



	def __init__( self		,   a_		,   a_=13		,   a_=7		,   a_=True		,   a_=False		,   a_=99		,   a_=16		,   a_=2		,   a_=4		,   a_=4		,   a_="gelu"		,   a_=0.1		,   a_=0.1		,   a_=32		,   a_=2		,   a_=1		,   a_=0		,   a_=0.02		,   ):
			lowerCAmelCase  :     Optional[int]       						= parent
			lowerCAmelCase  :     Tuple       						= batch_size
			lowerCAmelCase  :     Any       						= seq_length
			lowerCAmelCase  :     int       						= is_training
			lowerCAmelCase  :     int       						= use_labels
			lowerCAmelCase  :     Any       						= vocab_size
			lowerCAmelCase  :     Tuple       						= hidden_size
			lowerCAmelCase  :     Optional[int]       						= num_hidden_layers
			lowerCAmelCase  :     int       						= num_attention_heads
			lowerCAmelCase  :     int       						= intermediate_size
			lowerCAmelCase  :     str       						= hidden_act
			lowerCAmelCase  :     List[Any]       						= hidden_dropout_prob
			lowerCAmelCase  :     List[str]       						= attention_probs_dropout_prob
			lowerCAmelCase  :     Dict       						= max_position_embeddings
			lowerCAmelCase  :     Optional[int]       						= eos_token_id
			lowerCAmelCase  :     Union[str, Any]       						= pad_token_id
			lowerCAmelCase  :     Tuple       						= bos_token_id
			lowerCAmelCase  :     Tuple       						= initializer_range



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     Optional[Any]       						= np.clip(ids_tensor([self.batch_size, self.seq_length - 1]		,   self.vocab_size						)		,   3		,   self.vocab_size						)
			lowerCAmelCase  :     str       						= np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)		,   dtype=np.intaa						))		,   -1						)

			lowerCAmelCase  :     List[str]       						= shift_tokens_right(a_		,   1		,   2						)

			lowerCAmelCase  :     Union[str, Any]       						= BlenderbotConfig(
			    vocab_size=self.vocab_size		,   d_model=self.hidden_size		,   encoder_layers=self.num_hidden_layers		,   decoder_layers=self.num_hidden_layers		,   encoder_attention_heads=self.num_attention_heads		,   decoder_attention_heads=self.num_attention_heads		,   encoder_ffn_dim=self.intermediate_size		,   decoder_ffn_dim=self.intermediate_size		,   dropout=self.hidden_dropout_prob		,   attention_dropout=self.attention_probs_dropout_prob		,   max_position_embeddings=self.max_position_embeddings		,   eos_token_id=self.eos_token_id		,   bos_token_id=self.bos_token_id		,   pad_token_id=self.pad_token_id		,   initializer_range=self.initializer_range		,   use_cache=a_		,   )
			lowerCAmelCase  :     Tuple       						= prepare_blenderbot_inputs_dict(a_		,   a_		,   a_						)
			return config, inputs_dict



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase						,		lowerCAmelCase  :     Union[str, Any]       						= self.prepare_config_and_inputs()
			return config, inputs_dict



	def 	_lowerCamelCase     ( self		,   a_		,   a_		,   a_						):
			lowerCAmelCase  :     Any       						= 20
			lowerCAmelCase  :     int       						= model_class_name(a_						)

			lowerCAmelCase  :     Tuple       						= model.encode(inputs_dict["input_ids"]						)

			lowerCAmelCase						,		lowerCAmelCase  :     List[Any]       						= (
			    inputs_dict["decoder_input_ids"],
			    inputs_dict["decoder_attention_mask"],
			)

			lowerCAmelCase  :     Tuple       						= model.init_cache(decoder_input_ids.shape[0]		,   a_		,   a_						)
			lowerCAmelCase  :     Any       						= jnp.ones((decoder_input_ids.shape[0], max_decoder_length)		,   dtype="i4"						)

			lowerCAmelCase  :     Dict       						= jnp.broadcast_to(
			    jnp.arange(decoder_input_ids.shape[-1] - 1						)[None, :]		,   (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)		,   )
			lowerCAmelCase  :     Union[str, Any]       						= model.decode(
			    decoder_input_ids[:, :-1]		,   a_		,   decoder_attention_mask=a_		,   past_key_values=a_		,   decoder_position_ids=a_		,   )

			lowerCAmelCase  :     Tuple       						= jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]		,   dtype="i4"						)
			lowerCAmelCase  :     Optional[Any]       						= model.decode(
			    decoder_input_ids[:, -1:]		,   a_		,   decoder_attention_mask=a_		,   past_key_values=outputs_cache.past_key_values		,   decoder_position_ids=a_		,   )

			lowerCAmelCase  :     int       						= model.decode(a_		,   a_						)

			lowerCAmelCase  :     List[str]       						= 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		,   a_		,   a_		,   a_						):
			lowerCAmelCase  :     Union[str, Any]       						= 20
			lowerCAmelCase  :     Tuple       						= model_class_name(a_						)

			lowerCAmelCase  :     Dict       						= model.encode(inputs_dict["input_ids"]						)

			lowerCAmelCase						,		lowerCAmelCase  :     str       						= (
			    inputs_dict["decoder_input_ids"],
			    inputs_dict["decoder_attention_mask"],
			)

			lowerCAmelCase  :     Any       						= jnp.concatenate(
			    [
			        decoder_attention_mask,
			        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])						),
			    ]		,   axis=-1		,   )

			lowerCAmelCase  :     Union[str, Any]       						= model.init_cache(decoder_input_ids.shape[0]		,   a_		,   a_						)
			lowerCAmelCase  :     Any       						= jnp.broadcast_to(
			    jnp.arange(decoder_input_ids.shape[-1] - 1						)[None, :]		,   (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)		,   )

			lowerCAmelCase  :     Any       						= model.decode(
			    decoder_input_ids[:, :-1]		,   a_		,   decoder_attention_mask=a_		,   past_key_values=a_		,   decoder_position_ids=a_		,   )
			lowerCAmelCase  :     int       						= jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]		,   dtype="i4"						)
			lowerCAmelCase  :     Any       						= model.decode(
			    decoder_input_ids[:, -1:]		,   a_		,   past_key_values=outputs_cache.past_key_values		,   decoder_attention_mask=a_		,   decoder_position_ids=a_		,   )

			lowerCAmelCase  :     List[str]       						= model.decode(a_		,   a_		,   decoder_attention_mask=a_						)

			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}'''						)






@require_flax
class 		lowerCamelCase	(   unittest.TestCase							):
	snake_case_   = 99



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     List[Any]       						= np.array(
			    [
			        [71, 82, 18, 33, 46, 91, 2],
			        [68, 34, 26, 58, 30, 82, 2],
			        [5, 97, 17, 39, 94, 40, 2],
			        [76, 83, 94, 25, 70, 78, 2],
			        [87, 59, 41, 35, 48, 66, 2],
			        [55, 13, 16, 58, 5, 2, 1],  # note padding
			        [64, 27, 31, 51, 12, 75, 2],
			        [52, 64, 86, 17, 83, 39, 2],
			        [48, 61, 9, 24, 71, 82, 2],
			        [26, 1, 60, 48, 22, 13, 2],
			        [21, 5, 62, 28, 14, 76, 2],
			        [45, 98, 37, 86, 59, 48, 2],
			        [70, 70, 50, 9, 28, 0, 2],
			    ]		,   dtype=np.intaa		,   )

			lowerCAmelCase  :     Union[str, Any]       						= input_ids.shape[0]
			lowerCAmelCase  :     Dict       						= BlenderbotConfig(
			    vocab_size=self.vocab_size		,   d_model=24		,   encoder_layers=2		,   decoder_layers=2		,   encoder_attention_heads=2		,   decoder_attention_heads=2		,   encoder_ffn_dim=32		,   decoder_ffn_dim=32		,   max_position_embeddings=48		,   eos_token_id=2		,   pad_token_id=1		,   bos_token_id=0		,   )
			return config, input_ids, batch_size



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase  :     List[str]       						= self._get_config_and_data()
			lowerCAmelCase  :     Tuple       						= FlaxBlenderbotForConditionalGeneration(a_						)
			lowerCAmelCase  :     str       						= lm_model(input_ids=a_						)
			lowerCAmelCase  :     Tuple       						= (batch_size, input_ids.shape[1], config.vocab_size)
			self.assertEqual(outputs["logits"].shape		,   a_						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     Union[str, Any]       						= BlenderbotConfig(
			    vocab_size=self.vocab_size		,   d_model=14		,   encoder_layers=2		,   decoder_layers=2		,   encoder_attention_heads=2		,   decoder_attention_heads=2		,   encoder_ffn_dim=8		,   decoder_ffn_dim=8		,   max_position_embeddings=48		,   )
			lowerCAmelCase  :     int       						= FlaxBlenderbotForConditionalGeneration(a_						)
			lowerCAmelCase  :     Optional[Any]       						= np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]		,   dtype=np.intaa						)
			lowerCAmelCase  :     List[str]       						= np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]		,   dtype=np.intaa						)
			lowerCAmelCase  :     List[str]       						= lm_model(input_ids=a_		,   decoder_input_ids=a_						)
			lowerCAmelCase  :     List[str]       						= (*summary.shape, config.vocab_size)
			self.assertEqual(outputs["logits"].shape		,   a_						)





	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     int       						= np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]		,   dtype=np.intaa						)
			lowerCAmelCase  :     Dict       						= shift_tokens_right(a_		,   1		,   2						)
			lowerCAmelCase  :     Optional[Any]       						= np.equal(a_		,   1						).astype(np.floataa						).sum()
			lowerCAmelCase  :     Any       						= np.equal(a_		,   1						).astype(np.floataa						).sum()
			self.assertEqual(shifted.shape		,   input_ids.shape						)
			self.assertEqual(a_		,   n_pad_before - 1						)
			self.assertTrue(np.equal(shifted[:, 0]		,   2						).all()						)






@require_flax
class 		lowerCamelCase	(   _A   , unittest.TestCase   , _A							):
	snake_case_   = True
	snake_case_   = (
	    (
	        FlaxBlenderbotModel,
	        FlaxBlenderbotForConditionalGeneration,
	    )
	    if is_flax_available()
	    else ()
	)
	snake_case_   = (FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     Optional[Any]       						= FlaxBlenderbotModelTester(self						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase						,		lowerCAmelCase  :     Optional[int]       						= self.model_tester.prepare_config_and_inputs()
			for model_class in self.all_model_classes:
					self.model_tester.check_use_cache_forward(a_		,   a_		,   a_						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase						,		lowerCAmelCase  :     List[str]       						= self.model_tester.prepare_config_and_inputs()
			for model_class in self.all_model_classes:
					self.model_tester.check_use_cache_forward_with_attn_mask(a_		,   a_		,   a_						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase						,		lowerCAmelCase  :     Optional[Any]       						= self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__						):
							lowerCAmelCase  :     List[str]       						= self._prepare_for_class(a_		,   a_						)
							lowerCAmelCase  :     str       						= model_class(a_						)

							@jax.jit
							def encode_jitted(a_		,   a_=None		,   **a_						):
									return model.encode(input_ids=a_		,   attention_mask=a_						)

							with self.subTest("JIT Enabled"						):
									lowerCAmelCase  :     List[Any]       						= encode_jitted(**a_						).to_tuple()

							with self.subTest("JIT Disabled"						):
									with jax.disable_jit():
											lowerCAmelCase  :     Tuple       						= encode_jitted(**a_						).to_tuple()

							self.assertEqual(len(a_						)		,   len(a_						)						)
							for jitted_output, output in zip(a_		,   a_						):
									self.assertEqual(jitted_output.shape		,   output.shape						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase						,		lowerCAmelCase  :     Optional[int]       						= self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__						):
							lowerCAmelCase  :     int       						= model_class(a_						)
							lowerCAmelCase  :     Optional[int]       						= model.encode(inputs_dict["input_ids"]		,   inputs_dict["attention_mask"]						)

							lowerCAmelCase  :     Optional[Any]       						= {
							    "decoder_input_ids": inputs_dict["decoder_input_ids"],
							    "decoder_attention_mask": inputs_dict["decoder_attention_mask"],
							    "encoder_outputs": encoder_outputs,
							}

							@jax.jit
							def decode_jitted(a_		,   a_		,   a_						):
									return model.decode(
									    decoder_input_ids=a_		,   decoder_attention_mask=a_		,   encoder_outputs=a_		,   )

							with self.subTest("JIT Enabled"						):
									lowerCAmelCase  :     str       						= decode_jitted(**a_						).to_tuple()

							with self.subTest("JIT Disabled"						):
									with jax.disable_jit():
											lowerCAmelCase  :     Optional[Any]       						= decode_jitted(**a_						).to_tuple()

							self.assertEqual(len(a_						)		,   len(a_						)						)
							for jitted_output, output in zip(a_		,   a_						):
									self.assertEqual(jitted_output.shape		,   output.shape						)



	@slow
	def 	_lowerCamelCase     ( self						):
			for model_class_name in self.all_model_classes:
					lowerCAmelCase  :     List[str]       						= model_class_name.from_pretrained("facebook/blenderbot-400M-distill"						)
					# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
					lowerCAmelCase  :     Union[str, Any]       						= np.ones((1, 1)						) * model.config.eos_token_id
					lowerCAmelCase  :     Dict       						= model(a_						)
					self.assertIsNotNone(a_						)





	@unittest.skipUnless(jax_device != "cpu"		,   "3B test too slow on CPU."						)
	@slow
	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     Dict       						= {"num_beams": 1, "early_stopping": True, "min_length": 15, "max_length": 25}
			lowerCAmelCase  :     Any       						= {"skip_special_tokens": True, "clean_up_tokenization_spaces": True}

			lowerCAmelCase  :     Optional[int]       						= FlaxBlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-3B"		,   from_pt=a_						)
			lowerCAmelCase  :     List[str]       						= BlenderbotTokenizer.from_pretrained("facebook/blenderbot-3B"						)

			lowerCAmelCase  :     Any       						= ["Sam"]
			lowerCAmelCase  :     int       						= tokenizer(a_		,   return_tensors="jax"						)

			lowerCAmelCase  :     List[Any]       						= model.generate(**a_		,   **a_						)
			lowerCAmelCase  :     Optional[Any]       						= "Sam is a great name. It means \"sun\" in Gaelic."

			lowerCAmelCase  :     Optional[Any]       						= tokenizer.batch_decode(a_		,   **a_						)
			assert generated_txt[0].strip() == tgt_text



