


from __future__ import absolute_import, division, print_function, unicode_literals

from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers import RobertaConfig
from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.roberta.modeling_roberta import (
    ROBERTA_INPUTS_DOCSTRING,
    ROBERTA_START_DOCSTRING,
    RobertaEmbeddings,
)

from .modeling_highway_bert import BertPreTrainedModel, DeeBertModel, HighwayException, entropy







@add_start_docstrings(
    "The RoBERTa Model transformer with early exiting (DeeRoBERTa). "							,  lowercase__							,  )
class 			_lowercase (							lowercase__):


       """simple docstring"""




       A__											=			RobertaConfig
       A__											=			"roberta"





       def __init__(   self							:			Optional[int]					,      __lowerCamelCase							:			Dict       ):


             '''simple docstring'''



             super().__init__(__lowerCamelCase       )

             lowerCamelCase__		:				Union[str, Any]    =  RobertaEmbeddings(__lowerCamelCase       )
             self.init_weights()







@add_start_docstrings(
    "RoBERTa Model (with early exiting - DeeRoBERTa) with a classifier on top,\n    also takes care of multi-layer training. "							,  lowercase__							,  )
class 			_lowercase (							lowercase__):


       """simple docstring"""




       A__											=			RobertaConfig
       A__											=			"roberta"





       def __init__(   self							:			int					,      __lowerCamelCase							:			List[Any]       ):


             '''simple docstring'''



             super().__init__(__lowerCamelCase       )
             lowerCamelCase__		:				Optional[int]    =  config.num_labels
             lowerCamelCase__		:				Optional[Any]    =  config.num_hidden_layers

             lowerCamelCase__		:				Optional[int]    =  DeeRobertaModel(__lowerCamelCase       )
             lowerCamelCase__		:				int    =  nn.Dropout(config.hidden_dropout_prob       )
             lowerCamelCase__		:				str    =  nn.Linear(config.hidden_size					,      self.config.num_labels       )





       @add_start_docstrings_to_model_forward(__lowerCamelCase       )
       def 						lowerCAmelCase     (   self							:			Tuple					,      __lowerCamelCase							:			Dict=None					,      __lowerCamelCase							:			Optional[Any]=None					,      __lowerCamelCase							:			Optional[int]=None					,      __lowerCamelCase							:			Tuple=None					,      __lowerCamelCase							:			Dict=None					,      __lowerCamelCase							:			Any=None					,      __lowerCamelCase							:			Optional[int]=None					,      __lowerCamelCase							:			Optional[Any]=-1					,      __lowerCamelCase							:			Optional[Any]=False					,      ):


             '''simple docstring'''



             lowerCamelCase__		:				Optional[Any]    =  self.num_layers
             try:
                   lowerCamelCase__		:				List[str]    =  self.roberta(
                       __lowerCamelCase					,      attention_mask=__lowerCamelCase					,      token_type_ids=__lowerCamelCase					,      position_ids=__lowerCamelCase					,      head_mask=__lowerCamelCase					,      inputs_embeds=__lowerCamelCase					,      )

                   lowerCamelCase__		:				str    =  outputs[1]

                   lowerCamelCase__		:				Optional[Any]    =  self.dropout(__lowerCamelCase       )
                   lowerCamelCase__		:				str    =  self.classifier(__lowerCamelCase       )
                   lowerCamelCase__		:				Tuple    =  (logits,) + outputs[2:]  # add hidden states and attention if they are here
             except HighwayException as e:
                   lowerCamelCase__		:				Any    =  e.message
                   lowerCamelCase__		:				Optional[int]    =  e.exit_layer
                   lowerCamelCase__		:				List[Any]    =  outputs[0]

             if not self.training:
                   lowerCamelCase__		:				int    =  entropy(__lowerCamelCase       )
                   lowerCamelCase__		:				Any    =  []
                   lowerCamelCase__		:				Union[str, Any]    =  []
             if labels is not None:
                   if self.num_labels == 1:
                         #  We are doing regression
                         lowerCamelCase__		:				Optional[Any]    =  MSELoss()
                         lowerCamelCase__		:				Optional[int]    =  loss_fct(logits.view(-1       )					,      labels.view(-1       )       )
                   else:
                         lowerCamelCase__		:				str    =  CrossEntropyLoss()
                         lowerCamelCase__		:				Any    =  loss_fct(logits.view(-1					,      self.num_labels       )					,      labels.view(-1       )       )

                   # work with highway exits
                   lowerCamelCase__		:				Any    =  []
                   for highway_exit in outputs[-1]:
                         lowerCamelCase__		:				Union[str, Any]    =  highway_exit[0]
                         if not self.training:
                               highway_logits_all.append(__lowerCamelCase       )
                               highway_entropy.append(highway_exit[2]       )
                         if self.num_labels == 1:
                               #  We are doing regression
                               lowerCamelCase__		:				List[Any]    =  MSELoss()
                               lowerCamelCase__		:				int    =  loss_fct(highway_logits.view(-1       )					,      labels.view(-1       )       )
                         else:
                               lowerCamelCase__		:				Optional[Any]    =  CrossEntropyLoss()
                               lowerCamelCase__		:				Any    =  loss_fct(highway_logits.view(-1					,      self.num_labels       )					,      labels.view(-1       )       )
                         highway_losses.append(__lowerCamelCase       )

                   if train_highway:
                         lowerCamelCase__		:				Tuple    =  (sum(highway_losses[:-1]       ),) + outputs
                         # exclude the final highway, of course
                   else:
                         lowerCamelCase__		:				List[str]    =  (loss,) + outputs
             if not self.training:
                   lowerCamelCase__		:				Dict    =  outputs + ((original_entropy, highway_entropy), exit_layer)
                   if output_layer >= 0:
                         lowerCamelCase__		:				Tuple    =  (
                             (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
                         )  # use the highway of the last layer

             return outputs  # (loss), logits, (hidden_states), (attentions), entropy






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


