



"""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.albert.modeling_albert import (
    ALBERT_INPUTS_DOCSTRING,
    ALBERT_START_DOCSTRING,
    AlbertModel,
    AlbertPreTrainedModel,
    AlbertTransformer,
)


_lowerCAmelCase					:Union[str, Any]  =      logging.getLogger(__name__)
class 							_UpperCAmelCase							(				a   ):





     '''simple docstring'''



     def 				__lowerCAmelCase			(      self ,			A ,			A ,			A=None ,			A=None	)							->      Tuple:
         if current_layer == 0:
             _UpperCAmelCase							:						Optional[Any]       							=					self.embedding_hidden_mapping_in(A	)
         else:
             _UpperCAmelCase							:						Union[str, Any]       							=					hidden_states[0]

         _UpperCAmelCase							:						Dict       							=					int(self.config.num_hidden_layers / self.config.num_hidden_groups	)

         # Index of the hidden group
         _UpperCAmelCase							:						Dict       							=					int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)	)

         _UpperCAmelCase							:						List[Any]       							=					self.albert_layer_groups[group_idx](
             A ,			A ,			head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group] ,			)
         _UpperCAmelCase							:						str       							=					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 							_UpperCAmelCase							(				a   ):





     '''simple docstring'''



     def __init__(      self ,			A	)							->      str:
         super().__init__(A	)

         _UpperCAmelCase							:						Dict       							=					AlbertTransformerWithPabee(A	)

         self.init_weights()
         _UpperCAmelCase							:						str       							=					0
         _UpperCAmelCase							:						Union[str, Any]       							=					0
         _UpperCAmelCase							:						Dict       							=					0

         _UpperCAmelCase							:						str       							=					0





     def 				__lowerCAmelCase			(      self ,			A	)							->      Optional[int]:
         _UpperCAmelCase							:						List[Any]       							=					threshold





     def 				__lowerCAmelCase			(      self ,			A	)							->      Optional[Any]:
         _UpperCAmelCase							:						int       							=					patience





     def 				__lowerCAmelCase			(      self	)							->      Any:
         _UpperCAmelCase							:						Dict       							=					0
         _UpperCAmelCase							:						List[str]       							=					0





     def 				__lowerCAmelCase			(      self	)							->      List[Any]:
         _UpperCAmelCase							:						Union[str, Any]       							=					self.inference_layers_num / self.inference_instances_num
         _UpperCAmelCase							:						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(A	)





     @add_start_docstrings_to_model_forward(A	)
     def 				__lowerCAmelCase			(      self ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=False ,			)							->      int:

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

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

         if attention_mask is None:
             _UpperCAmelCase							:						Union[str, Any]       							=					torch.ones(A ,			device=A	)
         if token_type_ids is None:
             _UpperCAmelCase							:						Tuple       							=					torch.zeros(A ,			dtype=torch.long ,			device=A	)

         _UpperCAmelCase							:						Optional[Any]       							=					attention_mask.unsqueeze(1	).unsqueeze(2	)
         _UpperCAmelCase							:						Union[str, Any]       							=					extended_attention_mask.to(dtype=self.dtype	)  # fp16 compatibility
         _UpperCAmelCase							:						Dict       							=					(1.0 - extended_attention_mask) * -10_000.0
         _UpperCAmelCase							:						Any       							=					self.get_head_mask(A ,			self.config.num_hidden_layers	)

         _UpperCAmelCase							:						int       							=					self.embeddings(
             A ,			position_ids=A ,			token_type_ids=A ,			inputs_embeds=A	)
         _UpperCAmelCase							:						Any       							=					embedding_output

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

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

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

                 _UpperCAmelCase							:						Any       							=					logits
                 if patient_counter == self.patience:
                     break
             _UpperCAmelCase							:						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 							_UpperCAmelCase							(				a   ):





     '''simple docstring'''



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

         _UpperCAmelCase							:						Optional[int]       							=					AlbertModelWithPabee(A	)
         _UpperCAmelCase							:						List[str]       							=					nn.Dropout(config.classifier_dropout_prob	)
         _UpperCAmelCase							:						Optional[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(A	)
     def 				__lowerCAmelCase			(      self ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			)							->      List[str]:
         _UpperCAmelCase							:						Optional[Any]       							=					self.albert(
             input_ids=A ,			attention_mask=A ,			token_type_ids=A ,			position_ids=A ,			head_mask=A ,			inputs_embeds=A ,			output_dropout=self.dropout ,			output_layers=self.classifiers ,			regression=self.num_labels == 1 ,			)

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

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

         return outputs







"""simple docstring"""
import os
from typing import List, Union

import tensorflow as tf
from tensorflow_text import BertTokenizer as BertTokenizerLayer
from tensorflow_text import FastBertTokenizer, ShrinkLongestTrimmer, case_fold_utfa, combine_segments, pad_model_inputs

from .tokenization_bert import BertTokenizer
class 							_UpperCAmelCase							(				tf.keras.layers.Layer   ):





     '''simple docstring'''



     def __init__(      self ,			A ,			A ,			A = None ,			A = None ,			A = None ,			A = "longest" ,			A = True ,			A = 5_1_2 ,			A = None ,			A = True ,			A = True ,			A = True ,			**A ,			)							->      Tuple:
         super().__init__()
         if use_fast_bert_tokenizer:
             _UpperCAmelCase							:						Union[str, Any]       							=					FastBertTokenizer(
                 A ,			token_out_type=tf.intaa ,			lower_case_nfd_strip_accents=A ,			**A	)
         else:
             _UpperCAmelCase							:						List[Any]       							=					tf.lookup.StaticVocabularyTable(
                 tf.lookup.KeyValueTensorInitializer(
                     keys=A ,			key_dtype=tf.string ,			values=tf.range(tf.size(A ,			out_type=tf.intaa	) ,			dtype=tf.intaa	) ,			value_dtype=tf.intaa ,			) ,			num_oov_buckets=1 ,			)
             _UpperCAmelCase							:						str       							=					BertTokenizerLayer(
                 A ,			token_out_type=tf.intaa ,			lower_case=A ,			**A	)

         _UpperCAmelCase							:						int       							=					vocab_list
         _UpperCAmelCase							:						Dict       							=					do_lower_case
         _UpperCAmelCase							:						Optional[int]       							=					cls_token_id or vocab_list.index('''[CLS]'''	)
         _UpperCAmelCase							:						Any       							=					sep_token_id or vocab_list.index('''[SEP]'''	)
         _UpperCAmelCase							:						List[Any]       							=					pad_token_id or vocab_list.index('''[PAD]'''	)
         _UpperCAmelCase							:						List[str]       							=					ShrinkLongestTrimmer(max_length - 3 ,			axis=1	)  # Allow room for special tokens
         _UpperCAmelCase							:						Tuple       							=					max_length
         _UpperCAmelCase							:						List[Any]       							=					padding
         _UpperCAmelCase							:						List[Any]       							=					truncation
         _UpperCAmelCase							:						Union[str, Any]       							=					pad_to_multiple_of
         _UpperCAmelCase							:						Union[str, Any]       							=					return_token_type_ids
         _UpperCAmelCase							:						Union[str, Any]       							=					return_attention_mask





     @classmethod
     def 				__lowerCAmelCase			(      cls ,			A ,			**A	)							->      Dict:  # noqa: F821
         _UpperCAmelCase							:						Optional[int]       							=					kwargs.pop('''do_lower_case''' ,			A	)
         _UpperCAmelCase							:						int       							=					tokenizer.do_lower_case if do_lower_case is None else do_lower_case
         _UpperCAmelCase							:						Optional[Any]       							=					kwargs.pop('''cls_token_id''' ,			A	)
         _UpperCAmelCase							:						str       							=					tokenizer.cls_token_id if cls_token_id is None else cls_token_id
         _UpperCAmelCase							:						str       							=					kwargs.pop('''sep_token_id''' ,			A	)
         _UpperCAmelCase							:						int       							=					tokenizer.sep_token_id if sep_token_id is None else sep_token_id
         _UpperCAmelCase							:						List[Any]       							=					kwargs.pop('''pad_token_id''' ,			A	)
         _UpperCAmelCase							:						Optional[Any]       							=					tokenizer.pad_token_id if pad_token_id is None else pad_token_id

         _UpperCAmelCase							:						int       							=					tokenizer.get_vocab()
         _UpperCAmelCase							:						List[str]       							=					sorted([(wordpiece, idx) for wordpiece, idx in vocab.items()] ,			key=lambda A	: x[1]	)
         _UpperCAmelCase							:						Dict       							=					[entry[0] for entry in vocab]
         return cls(
             vocab_list=A ,			do_lower_case=A ,			cls_token_id=A ,			sep_token_id=A ,			pad_token_id=A ,			**A ,			)





     @classmethod
     def 				__lowerCAmelCase			(      cls ,			A ,			*A ,			**A	)							->      int:
         try:
             _UpperCAmelCase							:						int       							=					BertTokenizer.from_pretrained(A ,			*A ,			**A	)
         except:  # noqa: E722
             from .tokenization_bert_fast import BertTokenizerFast

             _UpperCAmelCase							:						Dict       							=					BertTokenizerFast.from_pretrained(A ,			*A ,			**A	)
         return cls.from_tokenizer(A ,			**A	)





     def 				__lowerCAmelCase			(      self ,			A	)							->      Any:
         if self.do_lower_case:
             _UpperCAmelCase							:						int       							=					case_fold_utfa(A	)
         _UpperCAmelCase							:						str       							=					self.tf_tokenizer.tokenize(A	)
         return tokens.merge_dims(1 ,			-1	)





     def 				__lowerCAmelCase			(      self ,			A ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			)							->      Optional[Any]:
         if padding is None:
             _UpperCAmelCase							:						str       							=					self.padding
         if padding not in ("longest", "max_length"):
             raise ValueError('''Padding must be either \'longest\' or \'max_length\'!'''	)
         if max_length is not None and text_pair is not None:
             # Because we have to instantiate a Trimmer to do it properly
             raise ValueError('''max_length cannot be overridden at call time when truncating paired texts!'''	)
         if max_length is None:
             _UpperCAmelCase							:						Union[str, Any]       							=					self.max_length
         if truncation is None:
             _UpperCAmelCase							:						List[Any]       							=					self.truncation
         if pad_to_multiple_of is None:
             _UpperCAmelCase							:						List[Any]       							=					self.pad_to_multiple_of
         if return_token_type_ids is None:
             _UpperCAmelCase							:						List[str]       							=					self.return_token_type_ids
         if return_attention_mask is None:
             _UpperCAmelCase							:						Any       							=					self.return_attention_mask
         if not isinstance(A ,			tf.Tensor	):
             _UpperCAmelCase							:						Dict       							=					tf.convert_to_tensor(A	)
         if text_pair is not None and not isinstance(A ,			tf.Tensor	):
             _UpperCAmelCase							:						str       							=					tf.convert_to_tensor(A	)
         if text_pair is not None:
             if text.shape.rank > 1:
                 raise ValueError('''text argument should not be multidimensional when a text pair is supplied!'''	)
             if text_pair.shape.rank > 1:
                 raise ValueError('''text_pair should not be multidimensional!'''	)
         if text.shape.rank == 2:
             _UpperCAmelCase							,       _UpperCAmelCase							:						str       							=					text[:, 0], text[:, 1]
         _UpperCAmelCase							:						Optional[Any]       							=					self.unpaired_tokenize(A	)
         if text_pair is None:  # Unpaired text
             if truncation:
                 _UpperCAmelCase							:						Optional[int]       							=					text[:, : max_length - 2]  # Allow room for special tokens
             _UpperCAmelCase							,       _UpperCAmelCase							:						Any       							=					combine_segments(
                 (text,) ,			start_of_sequence_id=self.cls_token_id ,			end_of_segment_id=self.sep_token_id	)
         else:  # Paired text
             _UpperCAmelCase							:						Tuple       							=					self.unpaired_tokenize(A	)
             if truncation:
                 _UpperCAmelCase							,       _UpperCAmelCase							:						Optional[int]       							=					self.paired_trimmer.trim([text, text_pair]	)
             _UpperCAmelCase							,       _UpperCAmelCase							:						Union[str, Any]       							=					combine_segments(
                 (text, text_pair) ,			start_of_sequence_id=self.cls_token_id ,			end_of_segment_id=self.sep_token_id	)
         if padding == "longest":
             _UpperCAmelCase							:						List[Any]       							=					input_ids.bounding_shape(axis=1	)
             if pad_to_multiple_of is not None:
                 # No ceiling division in tensorflow, so we negate floordiv instead
                 _UpperCAmelCase							:						Tuple       							=					pad_to_multiple_of * (-tf.math.floordiv(-pad_length ,			A	))
         else:
             _UpperCAmelCase							:						int       							=					max_length

         _UpperCAmelCase							,       _UpperCAmelCase							:						Any       							=					pad_model_inputs(A ,			max_seq_length=A ,			pad_value=self.pad_token_id	)
         _UpperCAmelCase							:						Dict       							=					{'''input_ids''': input_ids}
         if return_attention_mask:
             _UpperCAmelCase							:						Optional[int]       							=					attention_mask
         if return_token_type_ids:
             _UpperCAmelCase							,       _UpperCAmelCase							:						Tuple       							=					pad_model_inputs(
                 A ,			max_seq_length=A ,			pad_value=self.pad_token_id	)
             _UpperCAmelCase							:						int       							=					token_type_ids
         return output





     def 				__lowerCAmelCase			(      self	)							->      List[str]:
         return {
             "vocab_list": self.vocab_list,
             "do_lower_case": self.do_lower_case,
             "cls_token_id": self.cls_token_id,
             "sep_token_id": self.sep_token_id,
             "pad_token_id": self.pad_token_id,
         }


