'''simple docstring'''






import re
import string
from collections import Counter

import sacrebleu
import sacremoses
from packaging import version

import datasets


__A	:	Optional[Any]     =     '\n@inproceedings{xu-etal-2016-optimizing,\n    title = {Optimizing Statistical Machine Translation for Text Simplification},\n    authors={Xu, Wei and Napoles, Courtney and Pavlick, Ellie and Chen, Quanze and Callison-Burch, Chris},\n    journal = {Transactions of the Association for Computational Linguistics},\n    volume = {4},\n    year={2016},\n    url = {https://www.aclweb.org/anthology/Q16-1029},\n    pages = {401--415\n},\n@inproceedings{post-2018-call,\n    title = "A Call for Clarity in Reporting {BLEU} Scores",\n    author = "Post, Matt",\n    booktitle = "Proceedings of the Third Conference on Machine Translation: Research Papers",\n    month = oct,\n    year = "2018",\n    address = "Belgium, Brussels",\n    publisher = "Association for Computational Linguistics",\n    url = "https://www.aclweb.org/anthology/W18-6319",\n    pages = "186--191",\n}\n'

__A	:	Dict     =     '\\nWIKI_SPLIT is the combination of three metrics SARI, EXACT and SACREBLEU\nIt can be used to evaluate the quality of machine-generated texts.\n'


__A	:	Union[str, Any]     =     '\nCalculates sari score (between 0 and 100) given a list of source and predicted\nsentences, and a list of lists of reference sentences. It also computes the BLEU score as well as the exact match score.\nArgs:\n    sources: list of source sentences where each sentence should be a string.\n    predictions: list of predicted sentences where each sentence should be a string.\n    references: list of lists of reference sentences where each sentence should be a string.\nReturns:\n    sari: sari score\n    sacrebleu: sacrebleu score\n    exact: exact score\n\nExamples:\n    >>> sources=["About 95 species are currently accepted ."]\n    >>> predictions=["About 95 you now get in ."]\n    >>> references=[["About 95 species are currently known ."]]\n    >>> wiki_split = datasets.load_metric("wiki_split")\n    >>> results = wiki_split.compute(sources=sources, predictions=predictions, references=references)\n    >>> print(results)\n    {\'sari\': 21.805555555555557, \'sacrebleu\': 14.535768424205482, \'exact\': 0.0}\n'





def      UpperCAmelCase      (		lowerCamelCase_  :Union[str, Any] ):







		'''simple docstring'''

		def remove_articles(lowerCamelCase_  :Tuple ):
				snake_case_		:							Optional[Any]			       =						re.compile(R"""\b(a|an|the)\b"""				,      re.UNICODE )
				return re.sub(lowerCamelCase_				,      """ """				,      lowerCamelCase_ )

		def white_space_fix(lowerCamelCase_  :Union[str, Any] ):
				return " ".join(text.split() )

		def remove_punc(lowerCamelCase_  :str ):
				snake_case_		:							Optional[int]			       =						set(string.punctuation )
				return "".join(ch for ch in text if ch not in exclude )

		def lower(lowerCamelCase_  :Any ):
				return text.lower()

		return white_space_fix(remove_articles(remove_punc(lower(lowerCamelCase_ ) ) ) )





def      UpperCAmelCase      (		lowerCamelCase_  :Union[str, Any]				,      lowerCamelCase_  :List[Any] ):







		'''simple docstring'''
		return int(normalize_answer(lowerCamelCase_ ) == normalize_answer(lowerCamelCase_ ) )





def      UpperCAmelCase      (		lowerCamelCase_  :Dict				,      lowerCamelCase_  :Tuple ):







		'''simple docstring'''
		snake_case_		:							Tuple			       =						[any(compute_exact(lowerCamelCase_				,      lowerCamelCase_ ) for ref in refs ) for pred, refs in zip(lowerCamelCase_				,      lowerCamelCase_ )]
		return (sum(lowerCamelCase_ ) / len(lowerCamelCase_ )) * 1_00





def      UpperCAmelCase      (		lowerCamelCase_  :Union[str, Any]				,      lowerCamelCase_  :Union[str, Any]				,      lowerCamelCase_  :Dict				,      lowerCamelCase_  :str ):







		'''simple docstring'''
		snake_case_		:							Optional[Any]			       =						[rgram for rgrams in rgramslist for rgram in rgrams]
		snake_case_		:							Optional[int]			       =						Counter(lowerCamelCase_ )

		snake_case_		:							Dict			       =						Counter(lowerCamelCase_ )
		snake_case_		:							Tuple			       =						Counter()
		for sgram, scount in sgramcounter.items():
				snake_case_		:							Union[str, Any]			       =						scount * numref

		snake_case_		:							List[str]			       =						Counter(lowerCamelCase_ )
		snake_case_		:							Optional[int]			       =						Counter()
		for cgram, ccount in cgramcounter.items():
				snake_case_		:							List[Any]			       =						ccount * numref

		# KEEP
		snake_case_		:							Tuple			       =						sgramcounter_rep & cgramcounter_rep
		snake_case_		:							Any			       =						keepgramcounter_rep & rgramcounter
		snake_case_		:							int			       =						sgramcounter_rep & rgramcounter

		snake_case_		:							List[str]			       =						0
		snake_case_		:							List[str]			       =						0
		for keepgram in keepgramcountergood_rep:
				keeptmpscorea += keepgramcountergood_rep[keepgram] / keepgramcounter_rep[keepgram]
				# Fix an alleged bug [2] in the keep score computation.
				# keeptmpscore2 += keepgramcountergood_rep[keepgram] / keepgramcounterall_rep[keepgram]
				keeptmpscorea += keepgramcountergood_rep[keepgram]
		# Define 0/0=1 instead of 0 to give higher scores for predictions that match
		#      a target exactly.
		snake_case_		:							Optional[int]			       =						1
		snake_case_		:							Dict			       =						1
		if len(lowerCamelCase_ ) > 0:
				snake_case_		:							Optional[int]			       =						keeptmpscorea / len(lowerCamelCase_ )
		if len(lowerCamelCase_ ) > 0:
				# Fix an alleged bug [2] in the keep score computation.
				# keepscore_recall = keeptmpscore2 / len(keepgramcounterall_rep)
				snake_case_		:							Any			       =						keeptmpscorea / sum(keepgramcounterall_rep.values() )
		snake_case_		:							Optional[Any]			       =						0
		if keepscore_precision > 0 or keepscore_recall > 0:
				snake_case_		:							Dict			       =						2 * keepscore_precision * keepscore_recall / (keepscore_precision + keepscore_recall)

		# DELETION
		snake_case_		:							List[str]			       =						sgramcounter_rep - cgramcounter_rep
		snake_case_		:							Optional[Any]			       =						delgramcounter_rep - rgramcounter
		snake_case_		:							str			       =						sgramcounter_rep - rgramcounter
		snake_case_		:							Union[str, Any]			       =						0
		snake_case_		:							List[Any]			       =						0
		for delgram in delgramcountergood_rep:
				deltmpscorea += delgramcountergood_rep[delgram] / delgramcounter_rep[delgram]
				deltmpscorea += delgramcountergood_rep[delgram] / delgramcounterall_rep[delgram]
		# Define 0/0=1 instead of 0 to give higher scores for predictions that match
		# a target exactly.
		snake_case_		:							int			       =						1
		if len(lowerCamelCase_ ) > 0:
				snake_case_		:							List[str]			       =						deltmpscorea / len(lowerCamelCase_ )

		# ADDITION
		snake_case_		:							Optional[Any]			       =						set(lowerCamelCase_ ) - set(lowerCamelCase_ )
		snake_case_		:							Union[str, Any]			       =						set(lowerCamelCase_ ) & set(lowerCamelCase_ )
		snake_case_		:							int			       =						set(lowerCamelCase_ ) - set(lowerCamelCase_ )

		snake_case_		:							Union[str, Any]			       =						0
		for addgram in addgramcountergood:
				addtmpscore += 1

		# Define 0/0=1 instead of 0 to give higher scores for predictions that match
		# a target exactly.
		snake_case_		:							Optional[Any]			       =						1
		snake_case_		:							List[str]			       =						1
		if len(lowerCamelCase_ ) > 0:
				snake_case_		:							str			       =						addtmpscore / len(lowerCamelCase_ )
		if len(lowerCamelCase_ ) > 0:
				snake_case_		:							List[Any]			       =						addtmpscore / len(lowerCamelCase_ )
		snake_case_		:							List[Any]			       =						0
		if addscore_precision > 0 or addscore_recall > 0:
				snake_case_		:							Optional[int]			       =						2 * addscore_precision * addscore_recall / (addscore_precision + addscore_recall)

		return (keepscore, delscore_precision, addscore)





def      UpperCAmelCase      (		lowerCamelCase_  :List[str]				,      lowerCamelCase_  :Any				,      lowerCamelCase_  :int ):







		'''simple docstring'''
		snake_case_		:							List[str]			       =						len(lowerCamelCase_ )

		snake_case_		:							Tuple			       =						ssent.split(""" """ )
		snake_case_		:							Union[str, Any]			       =						csent.split(""" """ )
		snake_case_		:							int			       =						[]
		snake_case_		:							Any			       =						[]
		snake_case_		:							Union[str, Any]			       =						[]
		snake_case_		:							Optional[Any]			       =						[]
		snake_case_		:							Any			       =						[]
		snake_case_		:							Tuple			       =						[]

		snake_case_		:							str			       =						[]
		snake_case_		:							Optional[Any]			       =						[]
		snake_case_		:							Tuple			       =						[]
		snake_case_		:							Tuple			       =						[]
		for rsent in rsents:
				snake_case_		:							List[Any]			       =						rsent.split(""" """ )
				snake_case_		:							List[str]			       =						[]
				snake_case_		:							Optional[int]			       =						[]
				snake_case_		:							Dict			       =						[]
				ragramslist.append(lowerCamelCase_ )
				for i in range(0				,      len(lowerCamelCase_ ) - 1 ):
						if i < len(lowerCamelCase_ ) - 1:
								snake_case_		:							List[Any]			       =						ragrams[i] + """ """ + ragrams[i + 1]
								ragrams.append(lowerCamelCase_ )
						if i < len(lowerCamelCase_ ) - 2:
								snake_case_		:							int			       =						ragrams[i] + """ """ + ragrams[i + 1] + """ """ + ragrams[i + 2]
								ragrams.append(lowerCamelCase_ )
						if i < len(lowerCamelCase_ ) - 3:
								snake_case_		:							str			       =						ragrams[i] + """ """ + ragrams[i + 1] + """ """ + ragrams[i + 2] + """ """ + ragrams[i + 3]
								ragrams.append(lowerCamelCase_ )
				ragramslist.append(lowerCamelCase_ )
				ragramslist.append(lowerCamelCase_ )
				ragramslist.append(lowerCamelCase_ )

		for i in range(0				,      len(lowerCamelCase_ ) - 1 ):
				if i < len(lowerCamelCase_ ) - 1:
						snake_case_		:							Optional[Any]			       =						sagrams[i] + """ """ + sagrams[i + 1]
						sagrams.append(lowerCamelCase_ )
				if i < len(lowerCamelCase_ ) - 2:
						snake_case_		:							Dict			       =						sagrams[i] + """ """ + sagrams[i + 1] + """ """ + sagrams[i + 2]
						sagrams.append(lowerCamelCase_ )
				if i < len(lowerCamelCase_ ) - 3:
						snake_case_		:							List[Any]			       =						sagrams[i] + """ """ + sagrams[i + 1] + """ """ + sagrams[i + 2] + """ """ + sagrams[i + 3]
						sagrams.append(lowerCamelCase_ )

		for i in range(0				,      len(lowerCamelCase_ ) - 1 ):
				if i < len(lowerCamelCase_ ) - 1:
						snake_case_		:							Tuple			       =						cagrams[i] + """ """ + cagrams[i + 1]
						cagrams.append(lowerCamelCase_ )
				if i < len(lowerCamelCase_ ) - 2:
						snake_case_		:							Union[str, Any]			       =						cagrams[i] + """ """ + cagrams[i + 1] + """ """ + cagrams[i + 2]
						cagrams.append(lowerCamelCase_ )
				if i < len(lowerCamelCase_ ) - 3:
						snake_case_		:							Any			       =						cagrams[i] + """ """ + cagrams[i + 1] + """ """ + cagrams[i + 2] + """ """ + cagrams[i + 3]
						cagrams.append(lowerCamelCase_ )

		(snake_case_)		:							Union[str, Any]			       =						SARIngram(lowerCamelCase_				,      lowerCamelCase_				,      lowerCamelCase_				,      lowerCamelCase_ )
		(snake_case_)		:							Union[str, Any]			       =						SARIngram(lowerCamelCase_				,      lowerCamelCase_				,      lowerCamelCase_				,      lowerCamelCase_ )
		(snake_case_)		:							Optional[Any]			       =						SARIngram(lowerCamelCase_				,      lowerCamelCase_				,      lowerCamelCase_				,      lowerCamelCase_ )
		(snake_case_)		:							Any			       =						SARIngram(lowerCamelCase_				,      lowerCamelCase_				,      lowerCamelCase_				,      lowerCamelCase_ )
		snake_case_		:							Dict			       =						sum([keepascore, keepascore, keepascore, keepascore] ) / 4
		snake_case_		:							str			       =						sum([delascore, delascore, delascore, delascore] ) / 4
		snake_case_		:							Dict			       =						sum([addascore, addascore, addascore, addascore] ) / 4
		snake_case_		:							List[str]			       =						(avgkeepscore + avgdelscore + avgaddscore) / 3
		return finalscore





def      UpperCAmelCase      (		lowerCamelCase_  :Tuple				,      lowerCamelCase_  :bool = True				,      lowerCamelCase_  :str = "13a"				,      lowerCamelCase_  :bool = True ):







		'''simple docstring'''
		# Normalization is requried for the ASSET dataset (one of the primary
		# datasets in sentence simplification) to allow using space
		# to split the sentence. Even though Wiki-Auto and TURK datasets,
		# do not require normalization, we do it for consistency.
		# Code adapted from the EASSE library [1] written by the authors of the ASSET dataset.
		# [1] https://github.com/feralvam/easse/blob/580bba7e1378fc8289c663f864e0487188fe8067/easse/utils/preprocessing.py#L7

		if lowercase:
				snake_case_		:							str			       =						sentence.lower()

		if tokenizer in ["13a", "intl"]:
				if version.parse(sacrebleu.__version__ ).major >= 2:
						snake_case_		:							Any			       =						sacrebleu.metrics.bleu._get_tokenizer(lowerCamelCase_ )()(lowerCamelCase_ )
				else:
						snake_case_		:							List[Any]			       =						sacrebleu.TOKENIZERS[tokenizer]()(lowerCamelCase_ )
		elif tokenizer == "moses":
				snake_case_		:							str			       =						sacremoses.MosesTokenizer().tokenize(lowerCamelCase_				,      return_str=lowerCamelCase_				,      escape=lowerCamelCase_ )
		elif tokenizer == "penn":
				snake_case_		:							int			       =						sacremoses.MosesTokenizer().penn_tokenize(lowerCamelCase_				,      return_str=lowerCamelCase_ )
		else:
				snake_case_		:							Dict			       =						sentence

		if not return_str:
				snake_case_		:							Dict			       =						normalized_sent.split()

		return normalized_sent





def      UpperCAmelCase      (		lowerCamelCase_  :int				,      lowerCamelCase_  :Dict				,      lowerCamelCase_  :Union[str, Any] ):







		'''simple docstring'''
		if not (len(lowerCamelCase_ ) == len(lowerCamelCase_ ) == len(lowerCamelCase_ )):
				raise ValueError("""Sources length must match predictions and references lengths.""" )
		snake_case_		:							int			       =						0
		for src, pred, refs in zip(lowerCamelCase_				,      lowerCamelCase_				,      lowerCamelCase_ ):
				sari_score += SARIsent(normalize(lowerCamelCase_ )				,      normalize(lowerCamelCase_ )				,      [normalize(lowerCamelCase_ ) for sent in refs] )
		snake_case_		:							List[Any]			       =						sari_score / len(lowerCamelCase_ )
		return 1_00 * sari_score





def      UpperCAmelCase      (		lowerCamelCase_  :List[Any]				,      lowerCamelCase_  :Tuple				,      lowerCamelCase_  :List[str]="exp"				,      lowerCamelCase_  :Optional[int]=None				,      lowerCamelCase_  :List[str]=False				,      lowerCamelCase_  :Tuple=False				,      lowerCamelCase_  :int=False				,      ):







		'''simple docstring'''
		snake_case_		:							Optional[int]			       =						len(references[0] )
		if any(len(lowerCamelCase_ ) != references_per_prediction for refs in references ):
				raise ValueError("""Sacrebleu requires the same number of references for each prediction""" )
		snake_case_		:							int			       =						[[refs[i] for refs in references] for i in range(lowerCamelCase_ )]
		snake_case_		:							List[str]			       =						sacrebleu.corpus_bleu(
		    lowerCamelCase_				,      lowerCamelCase_				,      smooth_method=lowerCamelCase_				,      smooth_value=lowerCamelCase_				,      force=lowerCamelCase_				,      lowercase=lowerCamelCase_				,      use_effective_order=lowerCamelCase_				,      )
		return output.score






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





							def 		a__   (  self		:Tuple			):
									return datasets.MetricInfo(
									    description=_DESCRIPTION  ,citation=_CITATION  ,inputs_description=_KWARGS_DESCRIPTION  ,features=datasets.Features(
									        {
									            """predictions""": datasets.Value("""string"""  ,id="""sequence"""			),
									            """references""": datasets.Sequence(datasets.Value("""string"""  ,id="""sequence"""			)  ,id="""references"""			),
									        }			)  ,codebase_urls=[
									        """https://github.com/huggingface/transformers/blob/master/src/transformers/data/metrics/squad_metrics.py""",
									        """https://github.com/cocoxu/simplification/blob/master/SARI.py""",
									        """https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/utils/sari_hook.py""",
									        """https://github.com/mjpost/sacreBLEU""",
									    ]  ,reference_urls=[
									        """https://www.aclweb.org/anthology/Q16-1029.pdf""",
									        """https://github.com/mjpost/sacreBLEU""",
									        """https://en.wikipedia.org/wiki/BLEU""",
									        """https://towardsdatascience.com/evaluating-text-output-in-nlp-bleu-at-your-own-risk-e8609665a213""",
									    ]  ,)





							def 		a__   (  self		:Optional[Any]  ,_UpperCamelCase		:List[Any]  ,_UpperCamelCase		:List[Any]  ,_UpperCamelCase		:List[str]			):
									snake_case_		:							Tuple			       =						{}
									result.update({"""sari""": compute_sari(sources=_UpperCamelCase  ,predictions=_UpperCamelCase  ,references=_UpperCamelCase			)}			)
									result.update({"""sacrebleu""": compute_sacrebleu(predictions=_UpperCamelCase  ,references=_UpperCamelCase			)}			)
									result.update({"""exact""": compute_em(predictions=_UpperCamelCase  ,references=_UpperCamelCase			)}			)
									return result



'''simple docstring'''






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.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


__A	:	Any     =     logging.getLogger(__name__)
class 						__UpperCamelCase	(						lowercase__							):





							def 		a__   (  self		:int  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Optional[Any]  ,_UpperCamelCase		:List[Any]=None  ,_UpperCamelCase		:Any=None			):
									snake_case_		:							List[str]			       =						self.layer[current_layer](_UpperCamelCase  ,_UpperCamelCase  ,head_mask[current_layer]			)

									snake_case_		:							Tuple			       =						layer_outputs[0]

									return hidden_states





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





							def __init__(  self		:Dict  ,_UpperCamelCase		:Optional[Any]			):
									super().__init__(_UpperCamelCase			)

									snake_case_		:							List[Any]			       =						BertEncoderWithPabee(_UpperCamelCase			)

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

									snake_case_		:							str			       =						0





							def 		a__   (  self		:Tuple  ,_UpperCamelCase		:int			):
									snake_case_		:							Optional[Any]			       =						threshold





							def 		a__   (  self		:Optional[Any]  ,_UpperCamelCase		:Union[str, Any]			):
									snake_case_		:							str			       =						patience





							def 		a__   (  self		:Optional[int]			):
									snake_case_		:							List[str]			       =						0
									snake_case_		:							List[Any]			       =						0





							def 		a__   (  self		:List[Any]			):
									snake_case_		:							str			       =						self.inference_layers_num / self.inference_instances_num
									snake_case_		:							List[str]			       =						(
									    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(_UpperCamelCase			)





							@add_start_docstrings_to_model_forward(_UpperCamelCase			)
							def 		a__   (  self		:Optional[Any]  ,_UpperCamelCase		:Any=None  ,_UpperCamelCase		:List[str]=None  ,_UpperCamelCase		:Any=None  ,_UpperCamelCase		:Optional[int]=None  ,_UpperCamelCase		:Optional[int]=None  ,_UpperCamelCase		:int=None  ,_UpperCamelCase		:Optional[int]=None  ,_UpperCamelCase		:Any=None  ,_UpperCamelCase		:int=None  ,_UpperCamelCase		:Union[str, Any]=None  ,_UpperCamelCase		:int=False  ,):

									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:
											snake_case_		:							Any			       =						input_ids.size()
									elif inputs_embeds is not None:
											snake_case_		:							Union[str, Any]			       =						inputs_embeds.size()[:-1]
									else:
											raise ValueError("""You have to specify either input_ids or inputs_embeds"""			)

									snake_case_		:							Tuple			       =						input_ids.device if input_ids is not None else inputs_embeds.device

									if attention_mask is None:
											snake_case_		:							Any			       =						torch.ones(_UpperCamelCase  ,device=_UpperCamelCase			)
									if token_type_ids is None:
											snake_case_		:							str			       =						torch.zeros(_UpperCamelCase  ,dtype=torch.long  ,device=_UpperCamelCase			)

									# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
									# ourselves in which case we just need to make it broadcastable to all heads.
									snake_case_		:							torch.Tensor			       =						self.get_extended_attention_mask(_UpperCamelCase  ,_UpperCamelCase  ,_UpperCamelCase			)

									# If a 2D ou 3D attention mask is provided for the cross-attention
									# we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
									if self.config.is_decoder and encoder_hidden_states is not None:
											snake_case_       ,      snake_case_       ,      snake_case_		:							List[Any]			       =						encoder_hidden_states.size()
											snake_case_		:							str			       =						(encoder_batch_size, encoder_sequence_length)
											if encoder_attention_mask is None:
													snake_case_		:							List[str]			       =						torch.ones(_UpperCamelCase  ,device=_UpperCamelCase			)
											snake_case_		:							Optional[Any]			       =						self.invert_attention_mask(_UpperCamelCase			)
									else:
											snake_case_		:							List[str]			       =						None

									# Prepare head mask if needed
									# 1.0 in head_mask indicate we keep the head
									# attention_probs has shape bsz x n_heads x N x N
									# input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
									# and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
									snake_case_		:							List[Any]			       =						self.get_head_mask(_UpperCamelCase  ,self.config.num_hidden_layers			)

									snake_case_		:							Dict			       =						self.embeddings(
									    input_ids=_UpperCamelCase  ,position_ids=_UpperCamelCase  ,token_type_ids=_UpperCamelCase  ,inputs_embeds=_UpperCamelCase			)
									snake_case_		:							Optional[int]			       =						embedding_output

									if self.training:
											snake_case_		:							Any			       =						[]
											for i in range(self.config.num_hidden_layers			):
													snake_case_		:							Any			       =						self.encoder.adaptive_forward(
													    _UpperCamelCase  ,current_layer=_UpperCamelCase  ,attention_mask=_UpperCamelCase  ,head_mask=_UpperCamelCase			)

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

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

													snake_case_		:							List[Any]			       =						logits
													if patient_counter == self.patience:
															break
											snake_case_		:							int			       =						[patient_result]
											self.inference_layers_num += calculated_layer_num
											self.inference_instances_num += 1

									return res





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





							def __init__(  self		:Tuple  ,_UpperCamelCase		:List[str]			):
									super().__init__(_UpperCamelCase			)
									snake_case_		:							List[str]			       =						config.num_labels

									snake_case_		:							Dict			       =						BertModelWithPabee(_UpperCamelCase			)
									snake_case_		:							List[str]			       =						nn.Dropout(config.hidden_dropout_prob			)
									snake_case_		:							str			       =						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(_UpperCamelCase			)
							def 		a__   (  self		:Union[str, Any]  ,_UpperCamelCase		:Tuple=None  ,_UpperCamelCase		:Dict=None  ,_UpperCamelCase		:Union[str, Any]=None  ,_UpperCamelCase		:Union[str, Any]=None  ,_UpperCamelCase		:Tuple=None  ,_UpperCamelCase		:int=None  ,_UpperCamelCase		:Union[str, Any]=None  ,):
									snake_case_		:							int			       =						self.bert(
									    input_ids=_UpperCamelCase  ,attention_mask=_UpperCamelCase  ,token_type_ids=_UpperCamelCase  ,position_ids=_UpperCamelCase  ,head_mask=_UpperCamelCase  ,inputs_embeds=_UpperCamelCase  ,output_dropout=self.dropout  ,output_layers=self.classifiers  ,regression=self.num_labels == 1  ,)

									snake_case_		:							Optional[int]			       =						(logits[-1],)

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

									return outputs