






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


lowerCamelCase_				: Any    	=							logging.getLogger(__name__)
class    _SCREAMING_SNAKE_CASE							(		UpperCamelCase_	):


			'''simple docstring'''





			def 				A      (   self							:						Optional[int] ,      lowercase							:						str ,      lowercase							:						Any ,      lowercase							:						List[Any]=None ,      lowercase							:						Tuple=None			)							->    Tuple:

									'''simple docstring'''







									UpperCamelCase__         = self.layer[current_layer](lowercase ,      lowercase ,      head_mask[current_layer]			)

									UpperCamelCase__         = 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."       ,UpperCamelCase_       ,)
class    _SCREAMING_SNAKE_CASE							(		UpperCamelCase_	):


			'''simple docstring'''





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

									'''simple docstring'''







									super().__init__(lowercase			)

									UpperCamelCase__         = BertEncoderWithPabee(lowercase			)

									self.init_weights()
									UpperCamelCase__         = 0
									UpperCamelCase__         = 0
									UpperCamelCase__         = 0

									UpperCamelCase__         = 0






			def 				A      (   self							:						List[str] ,      lowercase							:						Any			)							->    Tuple:

									'''simple docstring'''







									UpperCamelCase__         = threshold






			def 				A      (   self							:						Any ,      lowercase							:						Any			)							->    Dict:

									'''simple docstring'''







									UpperCamelCase__         = patience






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

									'''simple docstring'''







									UpperCamelCase__         = 0
									UpperCamelCase__         = 0






			def 				A      (   self							:						List[str]			)							->    List[str]:

									'''simple docstring'''







									UpperCamelCase__         = self.inference_layers_num / self.inference_instances_num
									UpperCamelCase__         = (
									    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(lowercase			)







			@add_start_docstrings_to_model_forward(lowercase			)
			def 				A      (   self							:						Optional[int] ,      lowercase							:						int=None ,      lowercase							:						Optional[int]=None ,      lowercase							:						int=None ,      lowercase							:						Any=None ,      lowercase							:						Dict=None ,      lowercase							:						Any=None ,      lowercase							:						Optional[Any]=None ,      lowercase							:						int=None ,      lowercase							:						Union[str, Any]=None ,      lowercase							:						Tuple=None ,      lowercase							:						Tuple=False ,      )							->    Union[str, Any]:

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

									UpperCamelCase__         = input_ids.device if input_ids is not None else inputs_embeds.device

									if attention_mask is None:
															UpperCamelCase__         = torch.ones(lowercase ,      device=lowercase			)
									if token_type_ids is None:
															UpperCamelCase__         = torch.zeros(lowercase ,      dtype=torch.long ,      device=lowercase			)

									# 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.
									UpperCamelCase__         = self.get_extended_attention_mask(lowercase ,      lowercase ,      lowercase			)

									# 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:
															UpperCamelCase__         = encoder_hidden_states.size()
															UpperCamelCase__         = (encoder_batch_size, encoder_sequence_length)
															if encoder_attention_mask is None:
																					UpperCamelCase__         = torch.ones(lowercase ,      device=lowercase			)
															UpperCamelCase__         = self.invert_attention_mask(lowercase			)
									else:
															UpperCamelCase__         = 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]
									UpperCamelCase__         = self.get_head_mask(lowercase ,      self.config.num_hidden_layers			)

									UpperCamelCase__         = self.embeddings(
									    input_ids=lowercase ,      position_ids=lowercase ,      token_type_ids=lowercase ,      inputs_embeds=lowercase			)
									UpperCamelCase__         = embedding_output

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

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

																					UpperCamelCase__         = self.pooler(lowercase			)
																					UpperCamelCase__         = output_layers[i](lowercase			)
																					if regression:
																											UpperCamelCase__         = logits.detach()
																											if patient_result is not None:
																																	UpperCamelCase__         = patient_result.detach()
																											if (patient_result is not None) and torch.abs(patient_result - labels			) < self.regression_threshold:
																																	patient_counter += 1
																											else:
																																	UpperCamelCase__         = 0
																					else:
																											UpperCamelCase__         = logits.detach().argmax(dim=1			)
																											if patient_result is not None:
																																	UpperCamelCase__         = patient_result.detach().argmax(dim=1			)
																											if (patient_result is not None) and torch.all(labels.eq(lowercase			)			):
																																	patient_counter += 1
																											else:
																																	UpperCamelCase__         = 0

																					UpperCamelCase__         = logits
																					if patient_counter == self.patience:
																											break
															UpperCamelCase__         = [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. "       ,UpperCamelCase_       ,)
class    _SCREAMING_SNAKE_CASE							(		UpperCamelCase_	):


			'''simple docstring'''





			def __init__(   self							:						int ,      lowercase							:						Optional[int]			)							->    List[str]:

									'''simple docstring'''







									super().__init__(lowercase			)
									UpperCamelCase__         = config.num_labels

									UpperCamelCase__         = BertModelWithPabee(lowercase			)
									UpperCamelCase__         = nn.Dropout(config.hidden_dropout_prob			)
									UpperCamelCase__         = 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(lowercase			)
			def 				A      (   self							:						Any ,      lowercase							:						Dict=None ,      lowercase							:						Optional[int]=None ,      lowercase							:						Optional[Any]=None ,      lowercase							:						Dict=None ,      lowercase							:						Union[str, Any]=None ,      lowercase							:						Optional[Any]=None ,      lowercase							:						Tuple=None ,      )							->    Dict:

									'''simple docstring'''







									UpperCamelCase__         = self.bert(
									    input_ids=lowercase ,      attention_mask=lowercase ,      token_type_ids=lowercase ,      position_ids=lowercase ,      head_mask=lowercase ,      inputs_embeds=lowercase ,      output_dropout=self.dropout ,      output_layers=self.classifiers ,      regression=self.num_labels == 1 ,      )

									UpperCamelCase__         = (logits[-1],)

									if labels is not None:
															UpperCamelCase__         = None
															UpperCamelCase__         = 0
															for ix, logits_item in enumerate(lowercase			):
																					if self.num_labels == 1:
																											#  We are doing regression
																											UpperCamelCase__         = MSELoss()
																											UpperCamelCase__         = loss_fct(logits_item.view(-1			) ,      labels.view(-1			)			)
																					else:
																											UpperCamelCase__         = CrossEntropyLoss()
																											UpperCamelCase__         = loss_fct(logits_item.view(-1 ,      self.num_labels			) ,      labels.view(-1			)			)
																					if total_loss is None:
																											UpperCamelCase__         = loss
																					else:
																											total_loss += loss * (ix + 1)
																					total_weights += ix + 1
															UpperCamelCase__         = (total_loss / total_weights,) + outputs

									return outputs








'''simple docstring'''


from __future__ import annotations

from collections.abc import Callable
from typing import Generic, TypeVar

lowerCamelCase_				: List[str]    	=							TypeVar('''T''')
lowerCamelCase_				: Optional[int]    	=							TypeVar('''U''')
class    _SCREAMING_SNAKE_CASE							(		Generic[T, U]	):


			'''simple docstring'''





			def __init__(   self							:						Dict ,      lowercase							:						T | None ,      lowercase							:						U | None			)							->    List[Any]:

									'''simple docstring'''







									UpperCamelCase__         = key
									UpperCamelCase__         = val
									UpperCamelCase__         = None
									UpperCamelCase__         = None







			def __repr__(   self							:						List[Any]			)							->    str:

									'''simple docstring'''







									return (
									    f"Node: key: {self.key}, val: {self.val}, "
									    f"has next: {bool(self.next			)}, has prev: {bool(self.prev			)}"
									)
class    _SCREAMING_SNAKE_CASE							(		Generic[T, U]	):


			'''simple docstring'''





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

									'''simple docstring'''







									UpperCamelCase__         = DoubleLinkedListNode(lowercase ,      lowercase			)
									UpperCamelCase__         = DoubleLinkedListNode(lowercase ,      lowercase			)
									UpperCamelCase__    ,   UpperCamelCase__         = self.rear, self.head






			def __repr__(   self							:						int			)							->    str:

									'''simple docstring'''







									UpperCamelCase__         = ["""DoubleLinkedList"""]
									UpperCamelCase__         = self.head
									while node.next is not None:
															rep.append(str(lowercase			)			)
															UpperCamelCase__         = node.next
									rep.append(str(self.rear			)			)
									return ",\n    ".join(lowercase			)






			def 				A      (   self							:						str ,      lowercase							:						DoubleLinkedListNode[T, U]			)							->    None:

									'''simple docstring'''







									UpperCamelCase__         = self.rear.prev

									# All nodes other than self.head are guaranteed to have non-None previous
									assert previous is not None

									UpperCamelCase__         = node
									UpperCamelCase__         = previous
									UpperCamelCase__         = node
									UpperCamelCase__         = self.rear







			def 				A      (   self							:						Any ,      lowercase							:						DoubleLinkedListNode[T, U]			)							->    DoubleLinkedListNode[T, U] | None:

									'''simple docstring'''







									if node.prev is None or node.next is None:
															return None

									UpperCamelCase__         = node.next
									UpperCamelCase__         = node.prev
									UpperCamelCase__         = None
									UpperCamelCase__         = None
									return node
class    _SCREAMING_SNAKE_CASE							(		Generic[T, U]	):


			'''simple docstring'''





			__a							:							dict[Callable[[T], U], LRUCache[T, U]]					     =					{}






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

									'''simple docstring'''







									UpperCamelCase__         = DoubleLinkedList()
									UpperCamelCase__         = capacity
									UpperCamelCase__         = 0
									UpperCamelCase__         = 0
									UpperCamelCase__         = 0
									UpperCamelCase__         = {}






			def __repr__(   self							:						Any			)							->    str:

									'''simple docstring'''







									return (
									    f"CacheInfo(hits={self.hits}, misses={self.miss}, "
									    f"capacity={self.capacity}, current size={self.num_keys})"
									)






			def __contains__(   self							:						Any ,      lowercase							:						T			)							->    bool:

									'''simple docstring'''







									return key in self.cache






			def 				A      (   self							:						Tuple ,      lowercase							:						T			)							->    U | None:

									'''simple docstring'''







									if key in self.cache:
															self.hits += 1
															UpperCamelCase__         = self.cache[key]
															UpperCamelCase__         = self.list.remove(self.cache[key]			)
															assert node == value_node

															# node is guaranteed not None because it is in self.cache
															assert node is not None
															self.list.add(lowercase			)
															return node.val
									self.miss += 1
									return None






			def 				A      (   self							:						Dict ,      lowercase							:						T ,      lowercase							:						U			)							->    None:

									'''simple docstring'''







									if key not in self.cache:
															if self.num_keys >= self.capacity:
																					# delete first node (oldest) when over capacity
																					UpperCamelCase__         = self.list.head.next

																					# guaranteed to have a non-None first node when num_keys > 0
																					# explain to type checker via assertions
																					assert first_node is not None
																					assert first_node.key is not None
																					assert (
																					    self.list.remove(lowercase			) is not None
																					)  # node guaranteed to be in list assert node.key is not None

																					del self.cache[first_node.key]
																					self.num_keys -= 1
															UpperCamelCase__         = DoubleLinkedListNode(lowercase ,      lowercase			)
															self.list.add(self.cache[key]			)
															self.num_keys += 1

									else:
															# bump node to the end of the list, update value
															UpperCamelCase__         = self.list.remove(self.cache[key]			)
															assert node is not None  # node guaranteed to be in list
															UpperCamelCase__         = value
															self.list.add(lowercase			)







			@classmethod
			def 				A      (   cls							:						Optional[int] ,      lowercase							:						int = 1_2_8			)							->    Callable[[Callable[[T], U]], Callable[..., U]]:

									'''simple docstring'''







									def cache_decorator_inner(lowercase							:						Callable[[T], U]			) -> Callable[..., U]:
															def cache_decorator_wrapper(*lowercase							:						T			) -> U:
																					if func not in cls.decorator_function_to_instance_map:
																											UpperCamelCase__         = LRUCache(lowercase			)

																					UpperCamelCase__         = cls.decorator_function_to_instance_map[func].get(args[0]			)
																					if result is None:
																											UpperCamelCase__         = func(*lowercase			)
																											cls.decorator_function_to_instance_map[func].put(args[0] ,      lowercase			)
																					return result

															def cache_info() -> LRUCache[T, U]:
																					return cls.decorator_function_to_instance_map[func]

															setattr(lowercase ,      """cache_info""" ,      lowercase			)  # noqa: B010

															return cache_decorator_wrapper

									return cache_decorator_inner


if __name__ == "__main__":
					import doctest

					doctest.testmod()
