





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


__UpperCamelCase					    =	logging.getLogger(__name__)




class  _A (					__lowercase		):

		def 	lowercase__     (   self					:   int  ,      __magic_name__					:   Dict  ,      __magic_name__					:   Any  ,      __magic_name__					:   Union[str, Any]=None  ,      __magic_name__					:   Optional[Any]=None			)					->    str:


					"""simple docstring"""

					__snake_case					:      List[Any]							      =     self.layer[current_layer](__magic_name__  ,      __magic_name__  ,      head_mask[current_layer]			)

					__snake_case					:      Any							      =     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  _A (					__lowercase		):

		def __init__(   self					:   Optional[Any]  ,      __magic_name__					:   str			)					->    Dict:


					"""simple docstring"""

					super().__init__(__magic_name__			)

					__snake_case					:      Optional[Any]							      =     BertEncoderWithPabee(__magic_name__			)

					self.init_weights()
					__snake_case					:      Union[str, Any]							      =     0
					__snake_case					:      Union[str, Any]							      =     0
					__snake_case					:      Optional[int]							      =     0

					__snake_case					:      List[str]							      =     0

		def 	lowercase__     (   self					:   Union[str, Any]  ,      __magic_name__					:   Tuple			)					->    Dict:


					"""simple docstring"""

					__snake_case					:      Tuple							      =     threshold

		def 	lowercase__     (   self					:   List[str]  ,      __magic_name__					:   List[Any]			)					->    Tuple:


					"""simple docstring"""

					__snake_case					:      List[Any]							      =     patience

		def 	lowercase__     (   self					:   Optional[Any]			)					->    Tuple:


					"""simple docstring"""

					__snake_case					:      Optional[Any]							      =     0
					__snake_case					:      Dict							      =     0

		def 	lowercase__     (   self					:   Union[str, Any]			)					->    int:


					"""simple docstring"""

					__snake_case					:      Optional[Any]							      =     self.inference_layers_num / self.inference_instances_num
					__snake_case					:      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(__magic_name__			)

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


					"""simple docstring"""

					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					:      Union[str, Any]							      =     input_ids.size()
					elif inputs_embeds is not None:
								__snake_case					:      List[str]							      =     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					:      int							      =     torch.ones(__magic_name__  ,      device=__magic_name__			)
					if token_type_ids is None:
								__snake_case					:      int							      =     torch.zeros(__magic_name__  ,      dtype=torch.long  ,      device=__magic_name__			)

					# 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(__magic_name__  ,      __magic_name__  ,      __magic_name__			)

					# 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					:      str							      =     encoder_hidden_states.size()
								__snake_case					:      List[Any]							      =     (encoder_batch_size, encoder_sequence_length)
								if encoder_attention_mask is None:
											__snake_case					:      List[Any]							      =     torch.ones(__magic_name__  ,      device=__magic_name__			)
								__snake_case					:      Optional[Any]							      =     self.invert_attention_mask(__magic_name__			)
					else:
								__snake_case					:      Any							      =     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(__magic_name__  ,      self.config.num_hidden_layers			)

					__snake_case					:      Optional[Any]							      =     self.embeddings(
					    input_ids=__magic_name__  ,      position_ids=__magic_name__  ,      token_type_ids=__magic_name__  ,      inputs_embeds=__magic_name__			)
					__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(
											    __magic_name__  ,      current_layer=__magic_name__  ,      attention_mask=__magic_name__  ,      head_mask=__magic_name__			)

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

											__snake_case					:      Union[str, Any]							      =     self.pooler(__magic_name__			)
											__snake_case					:      str							      =     output_layers[i](__magic_name__			)
											if regression:
														__snake_case					:      str							      =     logits.detach()
														if patient_result is not None:
																	__snake_case					:      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					:      Any							      =     0
											else:
														__snake_case					:      int							      =     logits.detach().argmax(dim=1			)
														if patient_result is not None:
																	__snake_case					:      Any							      =     patient_result.detach().argmax(dim=1			)
														if (patient_result is not None) and torch.all(labels.eq(__magic_name__			)			):
																	patient_counter += 1
														else:
																	__snake_case					:      Tuple							      =     0

											__snake_case					:      Optional[int]							      =     logits
											if patient_counter == self.patience:
														break
								__snake_case					:      Tuple							      =     [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
    the pooled output) e.g. for GLUE tasks. '''						,					__lowercase						,					)
class  _A (					__lowercase		):

		def __init__(   self					:   Union[str, Any]  ,      __magic_name__					:   str			)					->    Any:


					"""simple docstring"""

					super().__init__(__magic_name__			)
					__snake_case					:      List[str]							      =     config.num_labels

					__snake_case					:      Union[str, Any]							      =     BertModelWithPabee(__magic_name__			)
					__snake_case					:      str							      =     nn.Dropout(config.hidden_dropout_prob			)
					__snake_case					:      Union[str, Any]							      =     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(__magic_name__			)
		def 	lowercase__     (   self					:   List[str]  ,      __magic_name__					:   Dict=None  ,      __magic_name__					:   str=None  ,      __magic_name__					:   Union[str, Any]=None  ,      __magic_name__					:   str=None  ,      __magic_name__					:   Dict=None  ,      __magic_name__					:   str=None  ,      __magic_name__					:   Tuple=None  ,      )					->    Optional[Any]:


					"""simple docstring"""

					__snake_case					:      Any							      =     self.bert(
					    input_ids=__magic_name__  ,      attention_mask=__magic_name__  ,      token_type_ids=__magic_name__  ,      position_ids=__magic_name__  ,      head_mask=__magic_name__  ,      inputs_embeds=__magic_name__  ,      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					:      List[str]							      =     None
								__snake_case					:      Optional[Any]							      =     0
								for ix, logits_item in enumerate(__magic_name__			):
											if self.num_labels == 1:
														#  We are doing regression
														__snake_case					:      Optional[int]							      =     MSELoss()
														__snake_case					:      List[Any]							      =     loss_fct(logits_item.view(-1			)  ,      labels.view(-1			)			)
											else:
														__snake_case					:      Union[str, Any]							      =     CrossEntropyLoss()
														__snake_case					:      List[Any]							      =     loss_fct(logits_item.view(-1  ,      self.num_labels			)  ,      labels.view(-1			)			)
											if total_loss is None:
														__snake_case					:      List[Any]							      =     loss
											else:
														total_loss += loss * (ix + 1)
											total_weights += ix + 1
								__snake_case					:      List[Any]							      =     (total_loss / total_weights,) + outputs

					return outputs





import os
import re
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


A_			:					Tuple    			=     logging.get_logger(__name__)

A_			:					Tuple    			=     {'vocab_file': 'spiece.model'}

A_			:					Optional[Any]    			=     {
    'vocab_file': {
        'google/bigbird-roberta-base': 'https://huggingface.co/google/bigbird-roberta-base/resolve/main/spiece.model',
        'google/bigbird-roberta-large': (
            'https://huggingface.co/google/bigbird-roberta-large/resolve/main/spiece.model'
        ),
        'google/bigbird-base-trivia-itc': (
            'https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/spiece.model'
        ),
    }
}

A_			:					Tuple    			=     {
    'google/bigbird-roberta-base': 4096,
    'google/bigbird-roberta-large': 4096,
    'google/bigbird-base-trivia-itc': 4096,
}



class        _a		(__magic_name__							):


					'''simple docstring'''






					UpperCAmelCase__:  Optional[int]    				=			VOCAB_FILES_NAMES
					UpperCAmelCase__:  str    				=			PRETRAINED_VOCAB_FILES_MAP
					UpperCAmelCase__:  List[str]    				=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					UpperCAmelCase__:  str    				=			['''input_ids''', '''attention_mask''']
					UpperCAmelCase__:  List[int]    				=			[]






					def __init__(					self       ,			A__       ,			A__="<unk>"       ,			A__="<s>"       ,			A__="</s>"       ,			A__="<pad>"       ,			A__="[SEP]"       ,			A__="[MASK]"       ,			A__="[CLS]"       ,			A__ = None       ,			**A__       ,			):
						A__				:       Tuple	  =						AddedToken(A__       ,			lstrip=A__       ,			rstrip=A__     ) if isinstance(A__       ,			A__     ) else bos_token
						A__				:       int	  =						AddedToken(A__       ,			lstrip=A__       ,			rstrip=A__     ) if isinstance(A__       ,			A__     ) else eos_token
						A__				:       Tuple	  =						AddedToken(A__       ,			lstrip=A__       ,			rstrip=A__     ) if isinstance(A__       ,			A__     ) else unk_token
						A__				:       str	  =						AddedToken(A__       ,			lstrip=A__       ,			rstrip=A__     ) if isinstance(A__       ,			A__     ) else pad_token
						A__				:       Dict	  =						AddedToken(A__       ,			lstrip=A__       ,			rstrip=A__     ) if isinstance(A__       ,			A__     ) else cls_token
						A__				:       Any	  =						AddedToken(A__       ,			lstrip=A__       ,			rstrip=A__     ) if isinstance(A__       ,			A__     ) else sep_token

						# Mask token behave like a normal word, i.e. include the space before it
						A__				:       Optional[Any]	  =						AddedToken(A__       ,			lstrip=A__       ,			rstrip=A__     ) if isinstance(A__       ,			A__     ) else mask_token

						A__				:       Optional[Any]	  =						{} if sp_model_kwargs is None else sp_model_kwargs

						super().__init__(
						    bos_token=A__       ,			eos_token=A__       ,			unk_token=A__       ,			pad_token=A__       ,			sep_token=A__       ,			mask_token=A__       ,			cls_token=A__       ,			sp_model_kwargs=self.sp_model_kwargs       ,			**A__       ,			)

						A__				:       Any	  =						vocab_file

						A__				:       Optional[int]	  =						spm.SentencePieceProcessor(**self.sp_model_kwargs     )
						self.sp_model.Load(A__     )






					@property
					def     __A     (					self     ):
						return self.sp_model.get_piece_size()






					def     __A     (					self     ):
						A__				:       List[str]	  =						{self.convert_ids_to_tokens(A__     ): i for i in range(self.vocab_size     )}
						vocab.update(self.added_tokens_encoder     )
						return vocab






					def __getstate__(					self     ):
						A__				:       Optional[Any]	  =						self.__dict__.copy()
						A__				:       Optional[Any]	  =						None
						return state






					def __setstate__(					self       ,			A__     ):
						A__				:       Dict	  =						d

						# for backward compatibility
						if not hasattr(self       ,			"""sp_model_kwargs"""     ):
							A__				:       str	  =						{}

						A__				:       Tuple	  =						spm.SentencePieceProcessor(**self.sp_model_kwargs     )
						self.sp_model.Load(self.vocab_file     )






					def     __A     (					self       ,			A__     ):
						return self.sp_model.encode(A__       ,			out_type=A__     )






					def     __A     (					self       ,			A__     ):
						return self.sp_model.piece_to_id(A__     )






					def     __A     (					self       ,			A__     ):
						A__				:       Dict	  =						self.sp_model.IdToPiece(A__     )
						return token






					def     __A     (					self       ,			A__     ):
						A__				:       str	  =						[]
						A__				:       Optional[int]	  =						""""""
						A__				:       Any	  =						False
						for token in tokens:
							# make sure that special tokens are not decoded using sentencepiece model
							if token in self.all_special_tokens:
								if not prev_is_special:
									out_string += " "
								out_string += self.sp_model.decode(A__     ) + token
								A__				:       Union[str, Any]	  =						True
								A__				:       Union[str, Any]	  =						[]
							else:
								current_sub_tokens.append(A__     )
								A__				:       Any	  =						False
						out_string += self.sp_model.decode(A__     )
						return out_string.strip()






					def     __A     (					self       ,			A__       ,			A__ = False       ,			A__ = None       ,			A__ = True       ,			**A__       ,			):
						A__				:       List[str]	  =						kwargs.pop("""use_source_tokenizer"""       ,			A__     )

						A__				:       Union[str, Any]	  =						self.convert_ids_to_tokens(A__       ,			skip_special_tokens=A__     )

						# To avoid mixing byte-level and unicode for byte-level BPT
						# we need to build string separately for added tokens and byte-level tokens
						# cf. https://github.com/huggingface/transformers/issues/1133
						A__				:       Dict	  =						[]
						A__				:       List[Any]	  =						[]
						for token in filtered_tokens:
							if skip_special_tokens and token in self.all_special_ids:
								continue
							if token in self.added_tokens_encoder:
								if current_sub_text:
									sub_texts.append(self.convert_tokens_to_string(A__     )     )
									A__				:       Optional[Any]	  =						[]
								sub_texts.append(A__     )
							else:
								current_sub_text.append(A__     )
						if current_sub_text:
							sub_texts.append(self.convert_tokens_to_string(A__     )     )

						# Mimic the behavior of the Rust tokenizer:
						# No space before [MASK] and [SEP]
						if spaces_between_special_tokens:
							A__				:       List[Any]	  =						re.sub(r""" (\[(MASK|SEP)\])"""       ,			r"""\1"""       ,			""" """.join(A__     )     )
						else:
							A__				:       Optional[Any]	  =						"""""".join(A__     )

						A__				:       Union[str, Any]	  =						(
						    clean_up_tokenization_spaces
						    if clean_up_tokenization_spaces is not None
						    else self.clean_up_tokenization_spaces
						)
						if clean_up_tokenization_spaces:
							A__				:       Union[str, Any]	  =						self.clean_up_tokenization(A__     )
							return clean_text
						else:
							return text






					def     __A     (					self       ,			A__       ,			A__ = None     ):
						if not os.path.isdir(A__     ):
							logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""     )
							return
						A__				:       str	  =						os.path.join(
						    A__       ,			(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]     )

						if os.path.abspath(self.vocab_file     ) != os.path.abspath(A__     ) and os.path.isfile(self.vocab_file     ):
							copyfile(self.vocab_file       ,			A__     )
						elif not os.path.isfile(self.vocab_file     ):
							with open(A__       ,			"""wb"""     ) as fi:
								A__				:       Dict	  =						self.sp_model.serialized_model_proto()
								fi.write(A__     )

						return (out_vocab_file,)






					def     __A     (					self       ,			A__       ,			A__ = None     ):
						if token_ids_a is None:
							return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
						A__				:       List[Any]	  =						[self.cls_token_id]
						A__				:       List[Any]	  =						[self.sep_token_id]
						return cls + token_ids_a + sep + token_ids_a + sep






					def     __A     (					self       ,			A__       ,			A__ = None       ,			A__ = False     ):
						if already_has_special_tokens:
							return super().get_special_tokens_mask(
							    token_ids_a=A__       ,			token_ids_a=A__       ,			already_has_special_tokens=A__     )

						if token_ids_a is None:
							return [1] + ([0] * len(A__     )) + [1]
						return [1] + ([0] * len(A__     )) + [1] + ([0] * len(A__     )) + [1]






					def     __A     (					self       ,			A__       ,			A__ = None     ):
						A__				:       Any	  =						[self.sep_token_id]
						A__				:       List[Any]	  =						[self.cls_token_id]
						if token_ids_a is None:
							return len(cls + token_ids_a + sep     ) * [0]
						return len(cls + token_ids_a + sep     ) * [0] + len(token_ids_a + sep     ) * [1]



