





'''simple docstring'''
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,
    BertEmbeddings,
    BertLayer,
    BertPooler,
    BertPreTrainedModel,
)
def 				lowercase__   (   __lowercase						:  Optional[Any]     ) ->      List[Any]:



       """simple docstring"""






       __UpperCamelCase       =  torch.exp(__lowercase     )
       __UpperCamelCase       =  torch.sum(__lowercase				,  dim=1     )  # sum of exp(x_i)
       __UpperCamelCase       =  torch.sum(x * exp_x				,  dim=1     )  # sum of x_i * exp(x_i)
       return torch.log(__lowercase     ) - B / A

class      snake_case     (    nn.Module       ):




 """simple docstring"""



 def __init__(       self       :							Any							,	__A       :							Optional[int]	):
        super().__init__()
        __UpperCamelCase       =  config.output_attentions
        __UpperCamelCase       =  config.output_hidden_states
        __UpperCamelCase       =  nn.ModuleList([BertLayer(__A	) for _ in range(config.num_hidden_layers	)]	)
        __UpperCamelCase       =  nn.ModuleList([BertHighway(__A	) for _ in range(config.num_hidden_layers	)]	)

        __UpperCamelCase       =  [-1 for _ in range(config.num_hidden_layers	)]






 def        _lowerCamelCase					(       self       :							Optional[Any]							,	__A       :							List[Any]	):
        if (type(__A	) is float) or (type(__A	) is int):
               for i in range(len(self.early_exit_entropy	)	):
                      __UpperCamelCase       =  x
        else:
               __UpperCamelCase       =  x






 def        _lowerCamelCase					(       self       :							str							,	__A       :							Tuple	):
        __UpperCamelCase       =  pooler.state_dict()
        for highway in self.highway:
               for name, param in highway.pooler.state_dict().items():
                      param.copy_(loaded_model[name]	)






 def        _lowerCamelCase					(       self       :							List[str]							,	__A       :							List[str]							,	__A       :							Optional[Any]=None							,	__A       :							Optional[Any]=None							,	__A       :							Union[str, Any]=None							,	__A       :							str=None							,	):
        __UpperCamelCase       =  ()
        __UpperCamelCase       =  ()
        __UpperCamelCase       =  ()
        for i, layer_module in enumerate(self.layer	):
               if self.output_hidden_states:
                      __UpperCamelCase       =  all_hidden_states + (hidden_states,)

               __UpperCamelCase       =  layer_module(
                   __A							,	__A							,	head_mask[i]							,	__A							,	__A	)
               __UpperCamelCase       =  layer_outputs[0]

               if self.output_attentions:
                      __UpperCamelCase       =  all_attentions + (layer_outputs[1],)

               __UpperCamelCase       =  (hidden_states,)
               if self.output_hidden_states:
                      __UpperCamelCase       =  current_outputs + (all_hidden_states,)
               if self.output_attentions:
                      __UpperCamelCase       =  current_outputs + (all_attentions,)

               __UpperCamelCase       =  self.highway[i](__A	)
               # logits, pooled_output

               if not self.training:
                      __UpperCamelCase       =  highway_exit[0]
                      __UpperCamelCase       =  entropy(__A	)
                      __UpperCamelCase       =  highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
                      __UpperCamelCase       =  all_highway_exits + (highway_exit,)

                      if highway_entropy < self.early_exit_entropy[i]:
                             __UpperCamelCase       =  (highway_logits,) + current_outputs[1:] + (all_highway_exits,)
                             raise HighwayException(__A							,	i + 1	)
               else:
                      __UpperCamelCase       =  all_highway_exits + (highway_exit,)

        # Add last layer
        if self.output_hidden_states:
               __UpperCamelCase       =  all_hidden_states + (hidden_states,)

        __UpperCamelCase       =  (hidden_states,)
        if self.output_hidden_states:
               __UpperCamelCase       =  outputs + (all_hidden_states,)
        if self.output_attentions:
               __UpperCamelCase       =  outputs + (all_attentions,)

        __UpperCamelCase       =  outputs + (all_highway_exits,)
        return outputs  # last-layer hidden state, (all hidden states), (all attentions), all highway exits






@add_start_docstrings(
    "The Bert Model transformer with early exiting (DeeBERT). "			,       __lowerCamelCase			,       )
class      snake_case     (    __lowerCamelCase       ):




 """simple docstring"""



 def __init__(       self       :							int							,	__A       :							int	):
        super().__init__(__A	)
        __UpperCamelCase       =  config

        __UpperCamelCase       =  BertEmbeddings(__A	)
        __UpperCamelCase       =  DeeBertEncoder(__A	)
        __UpperCamelCase       =  BertPooler(__A	)

        self.init_weights()






 def        _lowerCamelCase					(       self       :							Union[str, Any]	):
        self.encoder.init_highway_pooler(self.pooler	)






 def        _lowerCamelCase					(       self       :							List[str]	):
        return self.embeddings.word_embeddings






 def        _lowerCamelCase					(       self       :							Optional[Any]							,	__A       :							Union[str, Any]	):
        __UpperCamelCase       =  value






 def        _lowerCamelCase					(       self       :							Dict							,	__A       :							List[Any]	):
        for layer, heads in heads_to_prune.items():
               self.encoder.layer[layer].attention.prune_heads(__A	)






 @add_start_docstrings_to_model_forward(__A	)
 def        _lowerCamelCase					(       self       :							List[Any]							,	__A       :							Any=None							,	__A       :							Any=None							,	__A       :							Union[str, Any]=None							,	__A       :							Any=None							,	__A       :							Optional[int]=None							,	__A       :							Tuple=None							,	__A       :							Any=None							,	__A       :							List[str]=None							,	):
        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(__A							,	device=__A	)
        if encoder_attention_mask is None:
               __UpperCamelCase       =  torch.ones(__A							,	device=__A	)
        if token_type_ids is None:
               __UpperCamelCase       =  torch.zeros(__A							,	dtype=torch.long							,	device=__A	)

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

        # 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 encoder_attention_mask.dim() == 3:
               __UpperCamelCase       =  encoder_attention_mask[:, None, :, :]
        if encoder_attention_mask.dim() == 2:
               __UpperCamelCase       =  encoder_attention_mask[:, None, None, :]

        __UpperCamelCase       =  encoder_extended_attention_mask.to(
            dtype=next(self.parameters()	).dtype	)  # fp16 compatibility
        __UpperCamelCase       =  (1.0 - encoder_extended_attention_mask) * -1_0000.0

        # 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(__A							,	self.config.num_hidden_layers	)

        __UpperCamelCase       =  self.embeddings(
            input_ids=__A							,	position_ids=__A							,	token_type_ids=__A							,	inputs_embeds=__A	)
        __UpperCamelCase       =  self.encoder(
            __A							,	attention_mask=__A							,	head_mask=__A							,	encoder_hidden_states=__A							,	encoder_attention_mask=__A							,	)
        __UpperCamelCase       =  encoder_outputs[0]
        __UpperCamelCase       =  self.pooler(__A	)

        __UpperCamelCase       =  (
            sequence_output,
            pooled_output,
        ) + encoder_outputs[
            1:
        ]  # add hidden_states and attentions if they are here
        return outputs  # sequence_output, pooled_output, (hidden_states), (attentions), highway exits






class      snake_case     (    __lowerCamelCase       ):




 """simple docstring"""



 def __init__(       self       :							Optional[Any]							,	__A       :							Optional[Any]							,	__A       :							Optional[Any]	):
        __UpperCamelCase       =  message
        __UpperCamelCase       =  exit_layer  # start from 1!






class      snake_case     (    nn.Module       ):




 """simple docstring"""



 def __init__(       self       :							str							,	__A       :							List[str]	):
        super().__init__()
        __UpperCamelCase       =  BertPooler(__A	)
        __UpperCamelCase       =  nn.Dropout(config.hidden_dropout_prob	)
        __UpperCamelCase       =  nn.Linear(config.hidden_size							,	config.num_labels	)






 def        _lowerCamelCase					(       self       :							Dict							,	__A       :							Optional[int]	):
        # Pooler
        __UpperCamelCase       =  encoder_outputs[0]
        __UpperCamelCase       =  self.pooler(__A	)
        # "return" pooler_output

        # BertModel
        __UpperCamelCase       =  (pooler_input, pooler_output) + encoder_outputs[1:]
        # "return" bmodel_output

        # Dropout and classification
        __UpperCamelCase       =  bmodel_output[1]

        __UpperCamelCase       =  self.dropout(__A	)
        __UpperCamelCase       =  self.classifier(__A	)

        return logits, pooled_output






@add_start_docstrings(
    "Bert Model (with early exiting - DeeBERT) with a classifier on top,\n    also takes care of multi-layer training. "			,       __lowerCamelCase			,       )
class      snake_case     (    __lowerCamelCase       ):




 """simple docstring"""



 def __init__(       self       :							List[str]							,	__A       :							Any	):
        super().__init__(__A	)
        __UpperCamelCase       =  config.num_labels
        __UpperCamelCase       =  config.num_hidden_layers

        __UpperCamelCase       =  DeeBertModel(__A	)
        __UpperCamelCase       =  nn.Dropout(config.hidden_dropout_prob	)
        __UpperCamelCase       =  nn.Linear(config.hidden_size							,	self.config.num_labels	)

        self.init_weights()






 @add_start_docstrings_to_model_forward(__A	)
 def        _lowerCamelCase					(       self       :							Optional[int]							,	__A       :							Optional[int]=None							,	__A       :							Optional[int]=None							,	__A       :							List[Any]=None							,	__A       :							List[str]=None							,	__A       :							int=None							,	__A       :							Optional[Any]=None							,	__A       :							Any=None							,	__A       :							str=-1							,	__A       :							int=False							,	):
        __UpperCamelCase       =  self.num_layers
        try:
               __UpperCamelCase       =  self.bert(
                   __A							,	attention_mask=__A							,	token_type_ids=__A							,	position_ids=__A							,	head_mask=__A							,	inputs_embeds=__A							,	)
               # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

               __UpperCamelCase       =  outputs[1]

               __UpperCamelCase       =  self.dropout(__A	)
               __UpperCamelCase       =  self.classifier(__A	)
               __UpperCamelCase       =  (logits,) + outputs[2:]  # add hidden states and attention if they are here
        except HighwayException as e:
               __UpperCamelCase       =  e.message
               __UpperCamelCase       =  e.exit_layer
               __UpperCamelCase       =  outputs[0]

        if not self.training:
               __UpperCamelCase       =  entropy(__A	)
               __UpperCamelCase       =  []
               __UpperCamelCase       =  []
        if labels is not None:
               if self.num_labels == 1:
                      #  We are doing regression
                      __UpperCamelCase       =  MSELoss()
                      __UpperCamelCase       =  loss_fct(logits.view(-1	)							,	labels.view(-1	)	)
               else:
                      __UpperCamelCase       =  CrossEntropyLoss()
                      __UpperCamelCase       =  loss_fct(logits.view(-1							,	self.num_labels	)							,	labels.view(-1	)	)

               # work with highway exits
               __UpperCamelCase       =  []
               for highway_exit in outputs[-1]:
                      __UpperCamelCase       =  highway_exit[0]
                      if not self.training:
                             highway_logits_all.append(__A	)
                             highway_entropy.append(highway_exit[2]	)
                      if self.num_labels == 1:
                             #  We are doing regression
                             __UpperCamelCase       =  MSELoss()
                             __UpperCamelCase       =  loss_fct(highway_logits.view(-1	)							,	labels.view(-1	)	)
                      else:
                             __UpperCamelCase       =  CrossEntropyLoss()
                             __UpperCamelCase       =  loss_fct(highway_logits.view(-1							,	self.num_labels	)							,	labels.view(-1	)	)
                      highway_losses.append(__A	)

               if train_highway:
                      __UpperCamelCase       =  (sum(highway_losses[:-1]	),) + outputs
                      # exclude the final highway, of course
               else:
                      __UpperCamelCase       =  (loss,) + outputs
        if not self.training:
               __UpperCamelCase       =  outputs + ((original_entropy, highway_entropy), exit_layer)
               if output_layer >= 0:
                      __UpperCamelCase       =  (
                          (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
                      )  # use the highway of the last layer

        return outputs  # (loss), logits, (hidden_states), (attentions), (highway_exits)







'''simple docstring'''
import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


a__							:							List[Any] =logging.get_logger(__name__)

a__							:							List[Any] ={
    '''BAAI/AltCLIP''': '''https://huggingface.co/BAAI/AltCLIP/resolve/main/config.json''',
    # See all AltCLIP models at https://huggingface.co/models?filter=altclip
}

class      snake_case     (    __lowerCamelCase       ):




 """simple docstring"""



 SCREAMING_SNAKE_CASE_	:					Tuple					="altclip_text_model"






 def __init__(       self       :							str							,	__A       :							List[Any]=2_5_0_0_0_2							,	__A       :							Any=1_0_2_4							,	__A       :							int=2_4							,	__A       :							Dict=1_6							,	__A       :							Optional[Any]=4_0_9_6							,	__A       :							Union[str, Any]="gelu"							,	__A       :							Dict=0.1							,	__A       :							Dict=0.1							,	__A       :							List[str]=5_1_4							,	__A       :							Optional[int]=1							,	__A       :							int=0.02							,	__A       :							Optional[Any]=0.02							,	__A       :							Optional[Any]=1e-05							,	__A       :							Dict=1							,	__A       :							List[Any]=0							,	__A       :							int=2							,	__A       :							Tuple="absolute"							,	__A       :							Optional[Any]=True							,	__A       :							Optional[int]=7_6_8							,	**__A       :							List[str]							,	):
        super().__init__(pad_token_id=__A							,	bos_token_id=__A							,	eos_token_id=__A							,	**__A	)

        __UpperCamelCase       =  vocab_size
        __UpperCamelCase       =  hidden_size
        __UpperCamelCase       =  num_hidden_layers
        __UpperCamelCase       =  num_attention_heads
        __UpperCamelCase       =  hidden_act
        __UpperCamelCase       =  intermediate_size
        __UpperCamelCase       =  hidden_dropout_prob
        __UpperCamelCase       =  attention_probs_dropout_prob
        __UpperCamelCase       =  max_position_embeddings
        __UpperCamelCase       =  type_vocab_size
        __UpperCamelCase       =  initializer_range
        __UpperCamelCase       =  initializer_factor
        __UpperCamelCase       =  layer_norm_eps
        __UpperCamelCase       =  position_embedding_type
        __UpperCamelCase       =  use_cache
        __UpperCamelCase       =  project_dim






class      snake_case     (    __lowerCamelCase       ):




 """simple docstring"""



 SCREAMING_SNAKE_CASE_	:					Tuple					="altclip_vision_model"






 def __init__(       self       :							List[Any]							,	__A       :							Union[str, Any]=7_6_8							,	__A       :							Optional[int]=3_0_7_2							,	__A       :							Optional[Any]=5_1_2							,	__A       :							Tuple=1_2							,	__A       :							Union[str, Any]=1_2							,	__A       :							Optional[int]=3							,	__A       :							Dict=2_2_4							,	__A       :							Tuple=3_2							,	__A       :							str="quick_gelu"							,	__A       :							Dict=1e-5							,	__A       :							Optional[int]=0.0							,	__A       :							List[Any]=0.02							,	__A       :							int=1.0							,	**__A       :							Optional[int]							,	):
        super().__init__(**__A	)

        __UpperCamelCase       =  hidden_size
        __UpperCamelCase       =  intermediate_size
        __UpperCamelCase       =  projection_dim
        __UpperCamelCase       =  num_hidden_layers
        __UpperCamelCase       =  num_attention_heads
        __UpperCamelCase       =  num_channels
        __UpperCamelCase       =  patch_size
        __UpperCamelCase       =  image_size
        __UpperCamelCase       =  initializer_range
        __UpperCamelCase       =  initializer_factor
        __UpperCamelCase       =  attention_dropout
        __UpperCamelCase       =  layer_norm_eps
        __UpperCamelCase       =  hidden_act






 @classmethod
 def        _lowerCamelCase					(       cls       :							Optional[Any]							,	__A       :							Union[str, os.PathLike]							,	**__A       :							Optional[Any]	):
        cls._set_token_in_kwargs(__A	)

        __UpperCamelCase       ,						__UpperCamelCase       =  cls.get_config_dict(__A							,	**__A	)

        # get the vision config dict if we are loading from AltCLIPConfig
        if config_dict.get('model_type'	) == "altclip":
               __UpperCamelCase       =  config_dict['vision_config']

        if "model_type" in config_dict and hasattr(cls							,	'model_type'	) and config_dict["model_type"] != cls.model_type:
               logger.warning(
                   f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
                   f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''	)

        return cls.from_dict(__A							,	**__A	)






class      snake_case     (    __lowerCamelCase       ):




 """simple docstring"""



 SCREAMING_SNAKE_CASE_	:					List[str]					="altclip"
 SCREAMING_SNAKE_CASE_	:					Optional[int]					=True






 def __init__(       self       :							Any							,	__A       :							List[str]=None							,	__A       :							List[Any]=None							,	__A       :							List[str]=7_6_8							,	__A       :							List[str]=2.6592							,	**__A       :							Dict	):
        # If `_config_dict` exist, we use them for the backward compatibility.
        # We pop out these 2 attributes before calling `super().__init__` to avoid them being saved (which causes a lot
        # of confusion!).
        __UpperCamelCase       =  kwargs.pop('text_config_dict'							,	__A	)
        __UpperCamelCase       =  kwargs.pop('vision_config_dict'							,	__A	)

        super().__init__(**__A	)

        # Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in
        # `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most
        # cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`.
        if text_config_dict is not None:
               if text_config is None:
                      __UpperCamelCase       =  {}

               # This is the complete result when using `text_config_dict`.
               __UpperCamelCase       =  AltCLIPTextConfig(**__A	).to_dict()

               # Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
               for key, value in _text_config_dict.items():
                      if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
                             # If specified in `text_config_dict`
                             if key in text_config_dict:
                                    __UpperCamelCase       =  (
                                        f'''`{key}` is found in both `text_config_dict` and `text_config` but with different values. '''
                                        f'''The value `text_config_dict["{key}"]` will be used instead.'''
                                    )
                             # If inferred from default argument values (just to be super careful)
                             else:
                                    __UpperCamelCase       =  (
                                        f'''`text_config_dict` is provided which will be used to initialize `AltCLIPTextConfig`. The '''
                                        f'''value `text_config["{key}"]` will be overriden.'''
                                    )
                             logger.warning(__A	)

            # Update all values in `text_config` with the ones in `_text_config_dict`.
               text_config.update(_text_config_dict	)

        if vision_config_dict is not None:
               if vision_config is None:
                      __UpperCamelCase       =  {}

               # This is the complete result when using `vision_config_dict`.
               __UpperCamelCase       =  AltCLIPVisionConfig(**__A	).to_dict()
               # convert keys to string instead of integer
               if "id2label" in _vision_config_dict:
                      __UpperCamelCase       =  {
                          str(__A	): value for key, value in _vision_config_dict['id2label'].items()
                      }

               # Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
               for key, value in _vision_config_dict.items():
                      if key in vision_config and value != vision_config[key] and key not in ["transformers_version"]:
                             # If specified in `vision_config_dict`
                             if key in vision_config_dict:
                                    __UpperCamelCase       =  (
                                        f'''`{key}` is found in both `vision_config_dict` and `vision_config` but with different '''
                                        f'''values. The value `vision_config_dict["{key}"]` will be used instead.'''
                                    )
                             # If inferred from default argument values (just to be super careful)
                             else:
                                    __UpperCamelCase       =  (
                                        f'''`vision_config_dict` is provided which will be used to initialize `AltCLIPVisionConfig`. '''
                                        f'''The value `vision_config["{key}"]` will be overriden.'''
                                    )
                             logger.warning(__A	)

            # Update all values in `vision_config` with the ones in `_vision_config_dict`.
               vision_config.update(_vision_config_dict	)

        if text_config is None:
               __UpperCamelCase       =  {}
               logger.info('`text_config` is `None`. Initializing the `AltCLIPTextConfig` with default values.'	)

        if vision_config is None:
               __UpperCamelCase       =  {}
               logger.info('`vision_config` is `None`. initializing the `AltCLIPVisionConfig` with default values.'	)

        __UpperCamelCase       =  AltCLIPTextConfig(**__A	)
        __UpperCamelCase       =  AltCLIPVisionConfig(**__A	)

        __UpperCamelCase       =  projection_dim
        __UpperCamelCase       =  logit_scale_init_value
        __UpperCamelCase       =  1.0






 @classmethod
 def        _lowerCamelCase					(       cls       :							Union[str, Any]							,	__A       :							AltCLIPTextConfig							,	__A       :							AltCLIPVisionConfig							,	**__A       :							Optional[Any]	):
        return cls(text_config=text_config.to_dict()							,	vision_config=vision_config.to_dict()							,	**__A	)






 def        _lowerCamelCase					(       self       :							List[Any]	):
        __UpperCamelCase       =  copy.deepcopy(self.__dict__	)
        __UpperCamelCase       =  self.text_config.to_dict()
        __UpperCamelCase       =  self.vision_config.to_dict()
        __UpperCamelCase       =  self.__class__.model_type
        return output
