



"""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       :      int							=		logging.getLogger(__name__)






class 			UpperCAmelCase_						(       _UpperCamelCase		):



      def 		snake_case_      (  self  :			int  ,			A  :			Union[str, Any]  ,			A  :			int  ,			A  :			Tuple=None  ,			A  :			Any=None		):
           if current_layer == 0:
                _UpperCAmelCase    :			str              =					self.embedding_hidden_mapping_in(A		)
           else:
                _UpperCAmelCase    :			str              =					hidden_states[0]

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

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

           _UpperCAmelCase    :			Optional[int]              =					self.albert_layer_groups[group_idx](
               A  ,			A  ,			head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]  ,			)
           _UpperCAmelCase    :			List[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.' ,    _UpperCamelCase ,    )
class 			UpperCAmelCase_						(       _UpperCamelCase		):



      def __init__(  self  :			Optional[Any]  ,			A  :			Optional[int]		):
           super().__init__(A		)

           _UpperCAmelCase    :			int              =					AlbertTransformerWithPabee(A		)

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

           _UpperCAmelCase    :			Dict              =					0



      def 		snake_case_      (  self  :			int  ,			A  :			Optional[Any]		):
           _UpperCAmelCase    :			int              =					threshold



      def 		snake_case_      (  self  :			List[str]  ,			A  :			List[Any]		):
           _UpperCAmelCase    :			List[str]              =					patience



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



      def 		snake_case_      (  self  :			str		):
           _UpperCAmelCase    :			Any              =					self.inference_layers_num / self.inference_instances_num
           _UpperCAmelCase    :			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(A		)



      @add_start_docstrings_to_model_forward(A		)
      def 		snake_case_      (  self  :			Dict  ,			A  :			Optional[int]=None  ,			A  :			List[str]=None  ,			A  :			Any=None  ,			A  :			Any=None  ,			A  :			List[str]=None  ,			A  :			int=None  ,			A  :			str=None  ,			A  :			List[Any]=None  ,			A  :			str=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:
                _UpperCAmelCase    :			Tuple              =					input_ids.size()
           elif inputs_embeds is not None:
                _UpperCAmelCase    :			int              =					inputs_embeds.size()[:-1]
           else:
                raise ValueError("You have to specify either input_ids or inputs_embeds"		)

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

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

           _UpperCAmelCase    :			List[str]              =					attention_mask.unsqueeze(1		).unsqueeze(2		)
           _UpperCAmelCase    :			Optional[int]              =					extended_attention_mask.to(dtype=self.dtype		)  # fp16 compatibility
           _UpperCAmelCase    :			Optional[int]              =					(1.0 - extended_attention_mask) * -1_0_0_0_0.0
           _UpperCAmelCase    :			int              =					self.get_head_mask(A  ,			self.config.num_hidden_layers		)

           _UpperCAmelCase    :			Union[str, Any]              =					self.embeddings(
               A  ,			position_ids=A  ,			token_type_ids=A  ,			inputs_embeds=A		)
           _UpperCAmelCase    :			str              =					embedding_output

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

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

                     _UpperCAmelCase    :			Dict              =					self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]		)		)
                     _UpperCAmelCase    :			List[Any]              =					output_layers[i](A		)
                     if regression:
                          _UpperCAmelCase    :			Optional[int]              =					logits.detach()
                          if patient_result is not None:
                               _UpperCAmelCase    :			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:
                               _UpperCAmelCase    :			Dict              =					0
                     else:
                          _UpperCAmelCase    :			List[Any]              =					logits.detach().argmax(dim=1		)
                          if patient_result is not None:
                               _UpperCAmelCase    :			List[str]              =					patient_result.detach().argmax(dim=1		)
                          if (patient_result is not None) and torch.all(labels.eq(A		)		):
                               patient_counter += 1
                          else:
                               _UpperCAmelCase    :			Dict              =					0

                     _UpperCAmelCase    :			Tuple              =					logits
                     if patient_counter == self.patience:
                          break
                _UpperCAmelCase    :			Tuple              =					[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\n    the pooled output) e.g. for GLUE tasks. ' ,    _UpperCamelCase ,    )
class 			UpperCAmelCase_						(       _UpperCamelCase		):



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

           _UpperCAmelCase    :			Union[str, Any]              =					AlbertModelWithPabee(A		)
           _UpperCAmelCase    :			Any              =					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 		snake_case_      (  self  :			int  ,			A  :			Tuple=None  ,			A  :			Union[str, Any]=None  ,			A  :			Optional[int]=None  ,			A  :			Any=None  ,			A  :			Dict=None  ,			A  :			int=None  ,			A  :			Optional[Any]=None  ,			):
           _UpperCAmelCase    :			Optional[int]              =					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    :			List[str]              =					(logits[-1],)

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

           return outputs





"""simple docstring"""




import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
 from .tokenization_xglm import XGLMTokenizer
else:
 _lowerCAmelCase       :      Tuple							=							None


_lowerCAmelCase       :      List[str]							=							logging.get_logger(__name__)

_lowerCAmelCase       :      str							=							{"vocab_file": "sentencepiece.bpe.model", "tokenizer_file": "tokenizer.json"}

_lowerCAmelCase       :      Optional[Any]							=							{
    "vocab_file": {
        "facebook/xglm-564M": "https://huggingface.co/facebook/xglm-564M/resolve/main/sentencepiece.bpe.model",
    },
    "tokenizer_file": {
        "facebook/xglm-564M": "https://huggingface.co/facebook/xglm-564M/resolve/main/tokenizer.json",
    },
}

_lowerCAmelCase       :      Union[str, Any]							=							{
    "facebook/xglm-564M": 20_48,
}






class 			UpperCAmelCase_						(       _UpperCamelCase		):
      __SCREAMING_SNAKE_CASE		:     str				  =				VOCAB_FILES_NAMES
      __SCREAMING_SNAKE_CASE		:     List[str]				  =				PRETRAINED_VOCAB_FILES_MAP
      __SCREAMING_SNAKE_CASE		:     str				  =				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __SCREAMING_SNAKE_CASE		:     int				  =				['input_ids', 'attention_mask']
      __SCREAMING_SNAKE_CASE		:     str				  =				XGLMTokenizer



      def __init__(  self  :			List[Any]  ,			A  :			Union[str, Any]=None  ,			A  :			Optional[Any]=None  ,			A  :			Optional[Any]="<s>"  ,			A  :			Optional[int]="</s>"  ,			A  :			str="</s>"  ,			A  :			int="<s>"  ,			A  :			int="<unk>"  ,			A  :			List[str]="<pad>"  ,			**A  :			Tuple  ,			):
           # Compatibility with the original tokenizer
           _UpperCAmelCase    :			Union[str, Any]              =					7
           _UpperCAmelCase    :			Optional[int]              =					[f'<madeupword{i}>' for i in range(self.num_madeup_words		)]

           _UpperCAmelCase    :			Any              =					kwargs.get("additional_special_tokens"  ,			[]		)
           kwargs["additional_special_tokens"] += [
               word for word in madeup_words if word not in kwargs["additional_special_tokens"]
           ]

           super().__init__(
               A  ,			tokenizer_file=A  ,			bos_token=A  ,			eos_token=A  ,			sep_token=A  ,			cls_token=A  ,			unk_token=A  ,			pad_token=A  ,			**A  ,			)

           _UpperCAmelCase    :			Any              =					vocab_file
           _UpperCAmelCase    :			Union[str, Any]              =					False if not self.vocab_file else True



      def 		snake_case_      (  self  :			List[str]  ,			A  :			List[int]  ,			A  :			Optional[List[int]] = None		):

           if token_ids_a is None:
                return [self.sep_token_id] + token_ids_a
           _UpperCAmelCase    :			Tuple              =					[self.sep_token_id]
           return sep + token_ids_a + sep + sep + token_ids_a



      def 		snake_case_      (  self  :			str  ,			A  :			List[int]  ,			A  :			Optional[List[int]] = None		):
           _UpperCAmelCase    :			Optional[int]              =					[self.sep_token_id]

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



      def 		snake_case_      (  self  :			Dict  ,			A  :			str  ,			A  :			Optional[str] = None		):
           if not self.can_save_slow_tokenizer:
                raise ValueError(
                    "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
                    "tokenizer."		)

           if not os.path.isdir(A		):
                logger.error(f'Vocabulary path ({save_directory}) should be a directory.'		)
                return
           _UpperCAmelCase    :			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		):
                copyfile(self.vocab_file  ,			A		)

           return (out_vocab_file,)
