


import datasets

from .evaluate import (
    apply_no_ans_threshold,
    find_all_best_thresh,
    get_raw_scores,
    make_eval_dict,
    make_qid_to_has_ans,
    merge_eval,
)


__lowerCAmelCase     :			List[Any]      ='\\n@inproceedings{Rajpurkar2016SQuAD10,\n  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},\n  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},\n  booktitle={EMNLP},\n  year={2016}\n}\n'

__lowerCAmelCase     :			Union[str, Any]      ='\nThis metric wrap the official scoring script for version 2 of the Stanford Question\nAnswering Dataset (SQuAD).\n\nStanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by\ncrowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,\nfrom the corresponding reading passage, or the question might be unanswerable.\n\nSQuAD2.0 combines the 100,000 questions in SQuAD1.1 with over 50,000 unanswerable questions\nwritten adversarially by crowdworkers to look similar to answerable ones.\nTo do well on SQuAD2.0, systems must not only answer questions when possible, but also\ndetermine when no answer is supported by the paragraph and abstain from answering.\n'

__lowerCAmelCase     :			Tuple      ='\nComputes SQuAD v2 scores (F1 and EM).\nArgs:\n    predictions: List of triple for question-answers to score with the following elements:\n        - the question-answer \'id\' field as given in the references (see below)\n        - the text of the answer\n        - the probability that the question has no answer\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 list of Dict {\'text\': text of the answer as a string}\n    no_answer_threshold: float\n        Probability threshold to decide that a question has no answer.\nReturns:\n    \'exact\': Exact match (the normalized answer exactly match the gold answer)\n    \'f1\': The F-score of predicted tokens versus the gold answer\n    \'total\': Number of score considered\n    \'HasAns_exact\': Exact match (the normalized answer exactly match the gold answer)\n    \'HasAns_f1\': The F-score of predicted tokens versus the gold answer\n    \'HasAns_total\': Number of score considered\n    \'NoAns_exact\': Exact match (the normalized answer exactly match the gold answer)\n    \'NoAns_f1\': The F-score of predicted tokens versus the gold answer\n    \'NoAns_total\': Number of score considered\n    \'best_exact\': Best exact match (with varying threshold)\n    \'best_exact_thresh\': No-answer probability threshold associated to the best exact match\n    \'best_f1\': Best F1 (with varying threshold)\n    \'best_f1_thresh\': No-answer probability threshold associated to the best F1\nExamples:\n\n    >>> predictions = [{\'prediction_text\': \'1976\', \'id\': \'56e10a3be3433e1400422b22\', \'no_answer_probability\': 0.}]\n    >>> references = [{\'answers\': {\'answer_start\': [97], \'text\': [\'1976\']}, \'id\': \'56e10a3be3433e1400422b22\'}]\n    >>> squad_v2_metric = datasets.load_metric("squad_v2")\n    >>> results = squad_v2_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'exact\': 100.0, \'f1\': 100.0, \'total\': 1, \'HasAns_exact\': 100.0, \'HasAns_f1\': 100.0, \'HasAns_total\': 1, \'best_exact\': 100.0, \'best_exact_thresh\': 0.0, \'best_f1\': 100.0, \'best_f1_thresh\': 0.0}\n'
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION						,      _KWARGS_DESCRIPTION	)
class 	_lowercase				(      datasets.Metric	):






				'''simple docstring'''



				def 		__magic_name__(	self      :int     )		->      Optional[Any]:
										return datasets.MetricInfo(
										    description=_DESCRIPTION      ,      citation=_CITATION      ,      inputs_description=_KWARGS_DESCRIPTION      ,      features=datasets.Features(
										        {
										            '''predictions''': {
										                '''id''': datasets.Value('''string'''     ),
										                '''prediction_text''': datasets.Value('''string'''     ),
										                '''no_answer_probability''': datasets.Value('''float32'''     ),
										            },
										            '''references''': {
										                '''id''': datasets.Value('''string'''     ),
										                '''answers''': datasets.features.Sequence(
										                    {'''text''': datasets.Value('''string'''     ), '''answer_start''': datasets.Value('''int32'''     )}     ),
										            },
										        }     )      ,      codebase_urls=['''https://rajpurkar.github.io/SQuAD-explorer/''']      ,      reference_urls=['''https://rajpurkar.github.io/SQuAD-explorer/''']      ,      )



				def 		__magic_name__(	self      :Dict      ,      lowerCAmelCase__      :Optional[int]      ,      lowerCAmelCase__      :Tuple      ,      lowerCAmelCase__      :Union[str, Any]=1.0     )		->      str:
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		{p['''id''']: p['''no_answer_probability'''] for p in predictions}
										__SCREAMING_SNAKE_CASE						:					Dict		      =		[{'''paragraphs''': [{'''qas''': references}]}]
										__SCREAMING_SNAKE_CASE						:					Any		      =		{p['''id''']: p['''prediction_text'''] for p in predictions}

										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		make_qid_to_has_ans(lowerCAmelCase__     )  # maps qid to True/False
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		[k for k, v in qid_to_has_ans.items() if v]
										__SCREAMING_SNAKE_CASE						:					str		      =		[k for k, v in qid_to_has_ans.items() if not v]

										__SCREAMING_SNAKE_CASE      ,      __SCREAMING_SNAKE_CASE						:					Tuple		      =		get_raw_scores(lowerCAmelCase__      ,      lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		apply_no_ans_threshold(lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		apply_no_ans_threshold(lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		make_eval_dict(lowerCAmelCase__      ,      lowerCAmelCase__     )

										if has_ans_qids:
																__SCREAMING_SNAKE_CASE						:					Dict		      =		make_eval_dict(lowerCAmelCase__      ,      lowerCAmelCase__      ,      qid_list=lowerCAmelCase__     )
																merge_eval(lowerCAmelCase__      ,      lowerCAmelCase__      ,      '''HasAns'''     )
										if no_ans_qids:
																__SCREAMING_SNAKE_CASE						:					Any		      =		make_eval_dict(lowerCAmelCase__      ,      lowerCAmelCase__      ,      qid_list=lowerCAmelCase__     )
																merge_eval(lowerCAmelCase__      ,      lowerCAmelCase__      ,      '''NoAns'''     )
										find_all_best_thresh(lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__      ,      lowerCAmelCase__     )
										return dict(lowerCAmelCase__     )




import logging

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

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.albert.modeling_albert import (
    ALBERT_INPUTS_DOCSTRING,
    ALBERT_START_DOCSTRING,
    AlbertModel,
    AlbertPreTrainedModel,
    AlbertTransformer,
)


__lowerCAmelCase     :			Any      =logging.getLogger(__name__)
class 	_lowercase				(      A__	):






				'''simple docstring'''



				def 		__magic_name__(	self      :Optional[int]      ,      lowerCAmelCase__      :Dict      ,      lowerCAmelCase__      :List[str]      ,      lowerCAmelCase__      :List[Any]=None      ,      lowerCAmelCase__      :Optional[Any]=None     )		->      List[Any]:
										if current_layer == 0:
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		self.embedding_hidden_mapping_in(lowerCAmelCase__     )
										else:
																__SCREAMING_SNAKE_CASE						:					Tuple		      =		hidden_states[0]

										__SCREAMING_SNAKE_CASE						:					Tuple		      =		int(self.config.num_hidden_layers / self.config.num_hidden_groups     )

										# Index of the hidden group
										__SCREAMING_SNAKE_CASE						:					List[str]		      =		int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)     )

										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		self.albert_layer_groups[group_idx](
										    lowerCAmelCase__      ,      lowerCAmelCase__      ,      head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]      ,      )
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		layer_group_output[0]

										return (hidden_states,)




@add_start_docstrings(
    '''The bare ALBERT Model transformer with PABEE outputting raw hidden-states without any specific head on top.'''						,      A__						,      )
class 	_lowercase				(      A__	):






				'''simple docstring'''



				def __init__(	self      :int      ,      lowerCAmelCase__      :List[str]     )		->      List[Any]:
										super().__init__(lowerCAmelCase__     )

										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		AlbertTransformerWithPabee(lowerCAmelCase__     )

										self.init_weights()
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		0
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		0
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		0

										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		0


				def 		__magic_name__(	self      :int      ,      lowerCAmelCase__      :Optional[int]     )		->      str:
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		threshold


				def 		__magic_name__(	self      :List[Any]      ,      lowerCAmelCase__      :Optional[int]     )		->      List[Any]:
										__SCREAMING_SNAKE_CASE						:					Dict		      =		patience


				def 		__magic_name__(	self      :Union[str, Any]     )		->      str:
										__SCREAMING_SNAKE_CASE						:					Any		      =		0
										__SCREAMING_SNAKE_CASE						:					int		      =		0


				def 		__magic_name__(	self      :Optional[int]     )		->      Dict:
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		self.inference_layers_num / self.inference_instances_num
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		(
										    f'''*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='''
										    f''' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'''
										)
										print(lowerCAmelCase__     )



				@add_start_docstrings_to_model_forward(lowerCAmelCase__     )
				def 		__magic_name__(	self      :Dict      ,      lowerCAmelCase__      :List[str]=None      ,      lowerCAmelCase__      :List[str]=None      ,      lowerCAmelCase__      :Optional[Any]=None      ,      lowerCAmelCase__      :Any=None      ,      lowerCAmelCase__      :Dict=None      ,      lowerCAmelCase__      :Optional[int]=None      ,      lowerCAmelCase__      :str=None      ,      lowerCAmelCase__      :str=None      ,      lowerCAmelCase__      :str=False      ,      )		->      Union[str, Any]:

										if input_ids is not None and inputs_embeds is not None:
																raise ValueError('''You cannot specify both input_ids and inputs_embeds at the same time'''     )
										elif input_ids is not None:
																__SCREAMING_SNAKE_CASE						:					List[Any]		      =		input_ids.size()
										elif inputs_embeds is not None:
																__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		inputs_embeds.size()[:-1]
										else:
																raise ValueError('''You have to specify either input_ids or inputs_embeds'''     )

										__SCREAMING_SNAKE_CASE						:					List[str]		      =		input_ids.device if input_ids is not None else inputs_embeds.device

										if attention_mask is None:
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		torch.ones(lowerCAmelCase__      ,      device=lowerCAmelCase__     )
										if token_type_ids is None:
																__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		torch.zeros(lowerCAmelCase__      ,      dtype=torch.long      ,      device=lowerCAmelCase__     )

										__SCREAMING_SNAKE_CASE						:					Any		      =		attention_mask.unsqueeze(1     ).unsqueeze(2     )
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		extended_attention_mask.to(dtype=self.dtype     )  # fp16 compatibility
										__SCREAMING_SNAKE_CASE						:					Dict		      =		(1.0 - extended_attention_mask) * -1_0000.0
										__SCREAMING_SNAKE_CASE						:					str		      =		self.get_head_mask(lowerCAmelCase__      ,      self.config.num_hidden_layers     )

										__SCREAMING_SNAKE_CASE						:					List[str]		      =		self.embeddings(
										    lowerCAmelCase__      ,      position_ids=lowerCAmelCase__      ,      token_type_ids=lowerCAmelCase__      ,      inputs_embeds=lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		embedding_output

										if self.training:
																__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		[]
																for i in range(self.config.num_hidden_layers     ):
																						__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		self.encoder.adaptive_forward(
																						    lowerCAmelCase__      ,      current_layer=lowerCAmelCase__      ,      attention_mask=lowerCAmelCase__      ,      head_mask=lowerCAmelCase__      ,      )

																						__SCREAMING_SNAKE_CASE						:					Any		      =		self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]     )     )
																						__SCREAMING_SNAKE_CASE						:					Tuple		      =		output_layers[i](output_dropout(lowerCAmelCase__     )     )
																						res.append(lowerCAmelCase__     )
										elif self.patience == 0:  # Use all layers for inference
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		self.encoder(lowerCAmelCase__      ,      lowerCAmelCase__      ,      head_mask=lowerCAmelCase__     )
																__SCREAMING_SNAKE_CASE						:					List[Any]		      =		self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]     )     )
																__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		[output_layers[self.config.num_hidden_layers - 1](lowerCAmelCase__     )]
										else:
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		0
																__SCREAMING_SNAKE_CASE						:					Any		      =		None
																__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		0
																for i in range(self.config.num_hidden_layers     ):
																						calculated_layer_num += 1
																						__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		self.encoder.adaptive_forward(
																						    lowerCAmelCase__      ,      current_layer=lowerCAmelCase__      ,      attention_mask=lowerCAmelCase__      ,      head_mask=lowerCAmelCase__      ,      )

																						__SCREAMING_SNAKE_CASE						:					List[str]		      =		self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]     )     )
																						__SCREAMING_SNAKE_CASE						:					str		      =		output_layers[i](lowerCAmelCase__     )
																						if regression:
																												__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		logits.detach()
																												if patient_result is not None:
																																		__SCREAMING_SNAKE_CASE						:					int		      =		patient_result.detach()
																												if (patient_result is not None) and torch.abs(patient_result - labels     ) < self.regression_threshold:
																																		patient_counter += 1
																												else:
																																		__SCREAMING_SNAKE_CASE						:					List[str]		      =		0
																						else:
																												__SCREAMING_SNAKE_CASE						:					List[str]		      =		logits.detach().argmax(dim=1     )
																												if patient_result is not None:
																																		__SCREAMING_SNAKE_CASE						:					List[str]		      =		patient_result.detach().argmax(dim=1     )
																												if (patient_result is not None) and torch.all(labels.eq(lowerCAmelCase__     )     ):
																																		patient_counter += 1
																												else:
																																		__SCREAMING_SNAKE_CASE						:					str		      =		0

																						__SCREAMING_SNAKE_CASE						:					str		      =		logits
																						if patient_counter == self.patience:
																												break
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		[patient_result]
																self.inference_layers_num += calculated_layer_num
																self.inference_instances_num += 1

										return res




@add_start_docstrings(
    '''Albert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. '''						,      A__						,      )
class 	_lowercase				(      A__	):






				'''simple docstring'''



				def __init__(	self      :List[Any]      ,      lowerCAmelCase__      :int     )		->      Optional[int]:
										super().__init__(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Any		      =		config.num_labels

										__SCREAMING_SNAKE_CASE						:					str		      =		AlbertModelWithPabee(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		nn.Dropout(config.classifier_dropout_prob     )
										__SCREAMING_SNAKE_CASE						:					Dict		      =		nn.ModuleList(
										    [nn.Linear(config.hidden_size      ,      self.config.num_labels     ) for _ in range(config.num_hidden_layers     )]     )

										self.init_weights()



				@add_start_docstrings_to_model_forward(lowerCAmelCase__     )
				def 		__magic_name__(	self      :Optional[int]      ,      lowerCAmelCase__      :str=None      ,      lowerCAmelCase__      :List[str]=None      ,      lowerCAmelCase__      :Dict=None      ,      lowerCAmelCase__      :str=None      ,      lowerCAmelCase__      :Tuple=None      ,      lowerCAmelCase__      :int=None      ,      lowerCAmelCase__      :Any=None      ,      )		->      Dict:
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		self.albert(
										    input_ids=lowerCAmelCase__      ,      attention_mask=lowerCAmelCase__      ,      token_type_ids=lowerCAmelCase__      ,      position_ids=lowerCAmelCase__      ,      head_mask=lowerCAmelCase__      ,      inputs_embeds=lowerCAmelCase__      ,      output_dropout=self.dropout      ,      output_layers=self.classifiers      ,      regression=self.num_labels == 1      ,      )

										__SCREAMING_SNAKE_CASE						:					int		      =		(logits[-1],)

										if labels is not None:
																__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		None
																__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		0
																for ix, logits_item in enumerate(lowerCAmelCase__     ):
																						if self.num_labels == 1:
																												#  We are doing regression
																												__SCREAMING_SNAKE_CASE						:					str		      =		MSELoss()
																												__SCREAMING_SNAKE_CASE						:					List[Any]		      =		loss_fct(logits_item.view(-1     )      ,      labels.view(-1     )     )
																						else:
																												__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		CrossEntropyLoss()
																												__SCREAMING_SNAKE_CASE						:					List[Any]		      =		loss_fct(logits_item.view(-1      ,      self.num_labels     )      ,      labels.view(-1     )     )
																						if total_loss is None:
																												__SCREAMING_SNAKE_CASE						:					Dict		      =		loss
																						else:
																												total_loss += loss * (ix + 1)
																						total_weights += ix + 1
																__SCREAMING_SNAKE_CASE						:					Any		      =		(total_loss / total_weights,) + outputs

										return outputs
