



'''simple docstring'''







import importlib.util
import weakref
from functools import partialmethod

from .dependency_versions_check import dep_version_check
from .utils import is_accelerate_available, is_torch_available, logging


if is_torch_available():
      import torch

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





def 						snake_case__ ():
    return importlib.util.find_spec("deepspeed"	) is not None


if is_accelerate_available() and is_deepspeed_available():
      from accelerate.utils.deepspeed import HfDeepSpeedConfig as DeepSpeedConfig
else:
      # Inherits from a dummy `object` if accelerate is not available, so that python succeeds to import this file.
      # Deepspeed glue code will never inherit this dummy object as it checks if accelerate is available.
      from builtins import object as DeepSpeedConfig

class       lowerCamelCase_							(	lowerCAmelCase_  ):






   '''simple docstring'''

   def __init__(  self     ,						__UpperCamelCase ):
       # set global weakref object
       set_hf_deepspeed_config(self )
       dep_version_check("accelerate" )
       dep_version_check("deepspeed" )
       super().__init__(__UpperCamelCase )

class       lowerCamelCase_							(	lowerCAmelCase_  ):






   '''simple docstring'''

   def __init__(  self     ,						__UpperCamelCase ):
       super().__init__(__UpperCamelCase )
       UpperCamelCase__      =			None
       UpperCamelCase__      =			[]


   def 	_UpperCamelCase       (  self ):
       if self._dtype is None:
           raise ValueError("trainer_config_process() wasn't called yet to tell dtype" )
       return self._dtype


   def 	_UpperCamelCase       (  self     ,						__UpperCamelCase ):
       UpperCamelCase__      =			self.get_value(__UpperCamelCase )
       if val is None:
           return False
       else:
           return val == "auto"


   def 	_UpperCamelCase       (  self     ,						__UpperCamelCase     ,						__UpperCamelCase     ,						__UpperCamelCase=None     ,						__UpperCamelCase=True ):
       UpperCamelCase__    ,  UpperCamelCase__      =			self.find_config_node(__UpperCamelCase )
       if config is None:
           return

       if config.get(__UpperCamelCase ) == "auto":
           UpperCamelCase__      =			hf_val
           return

       if not must_match:
           return

       UpperCamelCase__      =			config.get(__UpperCamelCase )
       if ds_val is not None and ds_val != hf_val:
           self.mismatches.append(f'- ds {ds_key_long}={ds_val} vs hf {hf_key}={hf_val}' )

   A :     Any				=	partialmethod(lowerCAmelCase_					,    must_match=lowerCAmelCase_  )


   def 	_UpperCamelCase       (  self     ,						__UpperCamelCase ):
       UpperCamelCase__      =			args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps
       self.fill_match(
           "train_micro_batch_size_per_gpu"     ,						args.per_device_train_batch_size     ,						"per_device_train_batch_size" )
       self.fill_match("gradient_accumulation_steps"     ,						args.gradient_accumulation_steps     ,						"gradient_accumulation_steps" )
       self.fill_match("train_batch_size"     ,						__UpperCamelCase     ,						"train_batch_size (calculated)" )
       self.fill_match("gradient_clipping"     ,						args.max_grad_norm     ,						"max_grad_norm" )

       self.fill_match("optimizer.params.lr"     ,						args.learning_rate     ,						"learning_rate" )
       self.fill_match("optimizer.params.betas"     ,						[args.adam_betaa, args.adam_betaa]     ,						"adam_beta1+adam_beta2" )
       self.fill_match("optimizer.params.eps"     ,						args.adam_epsilon     ,						"adam_epsilon" )
       self.fill_match("optimizer.params.weight_decay"     ,						args.weight_decay     ,						"weight_decay" )

       self.fill_only("scheduler.params.warmup_min_lr"     ,						0 )  # not a trainer arg
       self.fill_match("scheduler.params.warmup_max_lr"     ,						args.learning_rate     ,						"learning_rate" )
       # total_num_steps - will get set in trainer_config_finalize

       # fp16
       if args.fpaa or args.fpaa_full_eval:
           UpperCamelCase__      =			"apex" if args.fpaa_backend == "apex" else "amp"
       else:
           UpperCamelCase__      =			None

       if args.save_on_each_node:
           # deepspeed uses shared storage by default. Let's override this setting if save_on_each_node == True
           UpperCamelCase__      =			self.config.get("checkpoint"     ,						{} )
           UpperCamelCase__      =			args.save_on_each_node

       # amp: similar to the pytorch native amp - it has a bunch of optional params but we won't set
       # any here unless the user did the work
       self.fill_match(
           "fp16.enabled"     ,						((args.fpaa or args.fpaa_full_eval) and fpaa_backend == "amp")     ,						"fp16|fp16_full_eval+fp16_backend(amp)"     ,						)

       # apex: delegates amp work to apex (which needs to be available), but it cannot be used with any
       # ZeRO features
       self.fill_match("amp.enabled"     ,						fpaa_backend == "apex"     ,						"fp16+fp16_backend(apex)" )
       self.fill_match("amp.opt_level"     ,						args.fpaa_opt_level     ,						"fp16_opt_level" )

       self.fill_match("bf16.enabled"     ,						(args.bfaa or args.bfaa_full_eval)     ,						"bf16|bf16_full_eval" )

       # deepspeed's default mode is fp16 unless there is a config that says differently
       if self.is_true("bf16.enabled" ):
           UpperCamelCase__      =			torch.bfloataa
       elif self.is_false("fp16.enabled" ):
           UpperCamelCase__      =			torch.floataa
       else:
           UpperCamelCase__      =			torch.floataa







   def 	_UpperCamelCase       (  self     ,						__UpperCamelCase     ,						__UpperCamelCase     ,						__UpperCamelCase ):
       UpperCamelCase__      =			[
           "zero_optimization.reduce_bucket_size",
           "zero_optimization.stage3_prefetch_bucket_size",
           "zero_optimization.stage3_param_persistence_threshold",
       ]
       UpperCamelCase__      =			[x for x in hidden_size_based_keys if self.is_auto(__UpperCamelCase )]

       if len(__UpperCamelCase ) > 0:
           if hasattr(model.config     ,						"hidden_size" ):
               UpperCamelCase__      =			model.config.hidden_size
           elif hasattr(model.config     ,						"hidden_sizes" ):
               # if there are many hidden sizes pick the largest one
               UpperCamelCase__      =			max(model.config.hidden_sizes )
           else:
               raise ValueError(
                   "The model's config file has neither `hidden_size` nor `hidden_sizes` entry, "
                   "therefore it's not possible to automatically fill out the following `auto` entries "
                   f'in the DeepSpeed config file: {hidden_size_auto_keys}. You can fix that by replacing '
                   "`auto` values for these keys with an integer value of your choice." )

           self.fill_only("zero_optimization.reduce_bucket_size"     ,						hidden_size * hidden_size )
           if self.is_zeroa():
               # automatically assign the optimal config values based on model config
               self.fill_only("zero_optimization.stage3_prefetch_bucket_size"     ,						0.9 * hidden_size * hidden_size )
               self.fill_only("zero_optimization.stage3_param_persistence_threshold"     ,						10 * hidden_size )

        # scheduler
       self.fill_match("scheduler.params.total_num_steps"     ,						__UpperCamelCase     ,						"num_training_steps (calculated)" )
       self.fill_match("scheduler.params.warmup_num_steps"     ,						args.get_warmup_steps(__UpperCamelCase )     ,						"warmup_steps" )

       if len(self.mismatches ) > 0:
           UpperCamelCase__      =			"\n".join(self.mismatches )
           raise ValueError(
               "Please correct the following DeepSpeed config values that mismatch TrainingArguments"
               f' values:\n{mismatches}\nThe easiest method is to set these DeepSpeed config values to \'auto\'.' )

lowercase__     :  List[str]    					=			None





def 						snake_case__ (lowerCAmelCase_	):
    # this is a special weakref global object to allow us to get to Deepspeed config from APIs
    # that don't have an easy way to get to the Deepspeed config outside of the Trainer domain.
    global _hf_deepspeed_config_weak_ref
    # will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed)
    UpperCamelCase__      =			weakref.ref(lowerCAmelCase_	)





def 						snake_case__ ():
    # useful for unit tests to ensure the global state doesn't leak - call from `tearDown` method
    global _hf_deepspeed_config_weak_ref
    UpperCamelCase__      =			None





def 						snake_case__ ():
    if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
        return _hf_deepspeed_config_weak_ref().is_zeroa()
    else:
        return False





def 						snake_case__ ():
    if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
        return _hf_deepspeed_config_weak_ref().config
    else:
        return None





def 						snake_case__ (lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_	):
    from accelerate.utils import DummyOptim, DummyScheduler

    UpperCamelCase__      =			hf_deepspeed_config.config

    # Optimizer + Scheduler
    # Currently supported combos:
    # 1. DS scheduler + DS optimizer: Yes
    # 2. HF scheduler + HF optimizer: Yes
    # 3. DS scheduler + HF optimizer: Yes
    # 4. HF scheduler + DS optimizer: No
    #
    # Unless Offload is enabled in which case it's:
    # 1. DS scheduler + DS optimizer: Yes
    # 2. HF scheduler + HF optimizer: Mostly*
    # 3. DS scheduler + HF optimizer: Mostly*
    # 4. HF scheduler + DS optimizer: No
    #
    # Mostly*: All non-native DeepSpeed optimizers that have both CPU and GPU implementation should work (except LAMB)

    UpperCamelCase__      =			None
    if "optimizer" in config:
        if args.adafactor:
            raise ValueError(
                "--adafactor was passed, but also found `optimizer` configured in the DeepSpeed config. "
                "Only one optimizer can be configured."	)
        UpperCamelCase__      =			DummyOptim(params=lowerCAmelCase_	)
    else:
        if hf_deepspeed_config.is_offload():
            logger.info(
                "Detected ZeRO Offload and non-DeepSpeed optimizers: This combination should work as long as the"
                " custom optimizer has both CPU and GPU implementation (except LAMB)"	)

        # ds supports Adam, OneBitAdam, and Lamb optimizers and can import other optimizers from torch.
        # But trainer uses AdamW by default.
        UpperCamelCase__      =			trainer.create_optimizer()
        # To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
        UpperCamelCase__      =			True

    UpperCamelCase__      =			None
    if "scheduler" in config:
        UpperCamelCase__      =			DummyScheduler(lowerCAmelCase_	)
    else:
        if isinstance(lowerCAmelCase_					,  lowerCAmelCase_	):
            raise ValueError(
                "Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. "
                "Please configure a scheduler in the DeepSpeed config."	)
        UpperCamelCase__      =			trainer.create_scheduler(num_training_steps=lowerCAmelCase_					,  optimizer=lowerCAmelCase_	)

    return optimizer, lr_scheduler





def 						snake_case__ (lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_=False	):
    from deepspeed.utils import logger as ds_logger

    UpperCamelCase__      =			trainer.model
    UpperCamelCase__      =			trainer.args

    UpperCamelCase__      =			trainer.accelerator.state.deepspeed_plugin.hf_ds_config

    # resume config update - some bits like `model` and `num_training_steps` only become available during train
    hf_deepspeed_config.trainer_config_finalize(lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_	)

    # set the Deepspeed log level consistent with the Trainer
    ds_logger.setLevel(args.get_process_log_level()	)

    if inference:
        # only Z3 makes sense for the inference
        if not hf_deepspeed_config.is_zeroa():
            raise ValueError("ZeRO inference only makes sense with ZeRO Stage 3 - please adjust your config"	)

        # in case the training config is re-used for inference
        hf_deepspeed_config.del_config_sub_tree("optimizer"	)
        hf_deepspeed_config.del_config_sub_tree("lr_scheduler"	)
        UpperCamelCase__    ,  UpperCamelCase__      =			None, None
        UpperCamelCase__      =			None
    else:
        UpperCamelCase__      =			None  # important for when deepspeed_init is used as re-init
        UpperCamelCase__      =			list(filter(lambda lowerCAmelCase_	: p.requires_grad					,  model.parameters()	)	)
        UpperCamelCase__    ,  UpperCamelCase__      =			deepspeed_optim_sched(
            lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_	)

    # keep for quick debug:
    # from pprint import pprint; pprint(config)

    return optimizer, lr_scheduler





def 						snake_case__ (lowerCAmelCase_					,  lowerCAmelCase_	):
    # it's possible that the user is trying to resume from model_path, which doesn't necessarily
    # contain a deepspeed checkpoint. e.g. examples just check if the dir exists and assume it's
    # a resume from a checkpoint and not just a local pretrained weight. So we check here if the
    # path contains what looks like a deepspeed checkpoint
    import glob

    UpperCamelCase__      =			sorted(glob.glob(F'{checkpoint_path}/global_step*'	)	)

    if len(lowerCAmelCase_	) > 0:
        logger.info(F'Attempting to resume from {checkpoint_path}'	)
        # this magically updates self.optimizer and self.lr_scheduler
        UpperCamelCase__    ,  UpperCamelCase__      =			deepspeed_engine.load_checkpoint(
            lowerCAmelCase_					,  load_optimizer_states=lowerCAmelCase_					,  load_lr_scheduler_states=lowerCAmelCase_	)
        if load_path is None:
            raise ValueError(F'[deepspeed] failed to resume from checkpoint {checkpoint_path}'	)
    else:
        raise ValueError(F'Can\'t find a valid checkpoint at {checkpoint_path}'	)









'''simple docstring'''







import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from ....modeling_outputs import BaseModelOutputWithPooling, SequenceClassifierOutput
from ....modeling_utils import ModuleUtilsMixin
from ....utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings


lowercase__     :  int    					=			logging.get_logger(__name__)

lowercase__     :  Any    					=			'MMBTConfig'

class       lowerCamelCase_							(	nn.Module  ):






   '''simple docstring'''

   def __init__(  self     ,						__UpperCamelCase     ,						__UpperCamelCase     ,						__UpperCamelCase ):
       super().__init__()
       UpperCamelCase__      =			config
       UpperCamelCase__      =			encoder
       UpperCamelCase__      =			nn.Linear(config.modal_hidden_size     ,						config.hidden_size )
       UpperCamelCase__      =			embeddings.position_embeddings
       UpperCamelCase__      =			embeddings.token_type_embeddings
       UpperCamelCase__      =			embeddings.word_embeddings
       UpperCamelCase__      =			embeddings.LayerNorm
       UpperCamelCase__      =			nn.Dropout(p=config.hidden_dropout_prob )







   def 	_UpperCamelCase       (  self     ,						__UpperCamelCase     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None ):
       UpperCamelCase__      =			self.proj_embeddings(self.encoder(__UpperCamelCase ) )
       UpperCamelCase__      =			token_embeddings.size(1 )

       if start_token is not None:
           UpperCamelCase__      =			self.word_embeddings(__UpperCamelCase )
           seq_length += 1
           UpperCamelCase__      =			torch.cat([start_token_embeds.unsqueeze(1 ), token_embeddings]     ,						dim=1 )

       if end_token is not None:
           UpperCamelCase__      =			self.word_embeddings(__UpperCamelCase )
           seq_length += 1
           UpperCamelCase__      =			torch.cat([token_embeddings, end_token_embeds.unsqueeze(1 )]     ,						dim=1 )

       if position_ids is None:
           UpperCamelCase__      =			torch.arange(__UpperCamelCase     ,						dtype=torch.long     ,						device=input_modal.device )
           UpperCamelCase__      =			position_ids.unsqueeze(0 ).expand(input_modal.size(0 )     ,						__UpperCamelCase )

       if token_type_ids is None:
           UpperCamelCase__      =			torch.zeros(
               (input_modal.size(0 ), seq_length)     ,						dtype=torch.long     ,						device=input_modal.device )

       UpperCamelCase__      =			self.position_embeddings(__UpperCamelCase )
       UpperCamelCase__      =			self.token_type_embeddings(__UpperCamelCase )
       UpperCamelCase__      =			token_embeddings + position_embeddings + token_type_embeddings
       UpperCamelCase__      =			self.LayerNorm(__UpperCamelCase )
       UpperCamelCase__      =			self.dropout(__UpperCamelCase )
       return embeddings

lowercase__     :  Union[str, Any]    					=			r'\n    MMBT model was proposed in [Supervised Multimodal Bitransformers for Classifying Images and\n    Text](https://github.com/facebookresearch/mmbt) by Douwe Kiela, Suvrat Bhooshan, Hamed Firooz, Davide Testuggine.\n    It\'s a supervised multimodal bitransformer model that fuses information from text and other image encoders, and\n    obtain state-of-the-art performance on various multimodal classification benchmark tasks.\n\n    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the\n    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads\n    etc.)\n\n    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.\n    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage\n    and behavior.\n\n    Parameters:\n        config ([`MMBTConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration.\n        transformer (`nn.Module`): A text transformer that is used by MMBT.\n            It should have embeddings, encoder, and pooler attributes.\n        encoder (`nn.Module`): Encoder for the second modality.\n            It should take in a batch of modal inputs and return k, n dimension embeddings.\n'

lowercase__     :  List[str]    					=			r'\n    Args:\n        input_modal (`torch.FloatTensor` of shape `(batch_size, ***)`):\n            The other modality data. It will be the shape that the encoder for that type expects. e.g. With an Image\n            Encoder, the shape would be (batch_size, channels, height, width)\n        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary. It does not expect [CLS] token to be added as it\'s\n            appended to the end of other modality embeddings. Indices can be obtained using [`AutoTokenizer`]. See\n            [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        modal_start_tokens (`torch.LongTensor` of shape `(batch_size,)`, *optional*):\n            Optional start token to be added to Other Modality Embedding. [CLS] Most commonly used for classification\n            tasks.\n        modal_end_tokens (`torch.LongTensor` of shape `(batch_size,)`, *optional*):\n            Optional end token to be added to Other Modality Embedding. [SEP] Most commonly used.\n        attention_mask (*optional*) `torch.FloatTensor` of shape `(batch_size, sequence_length)`:\n            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:\n\n            - 1 for tokens that are **not masked**,\n            - 0 for tokens that are **masked**.\n\n            [What are attention masks?](../glossary#attention-mask)\n        token_type_ids (*optional*) `torch.LongTensor` of shape `(batch_size, sequence_length)`:\n            Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,\n            1]`:\n\n            - 0 corresponds to a *sentence A* token,\n            - 1 corresponds to a *sentence B* token.\n\n            [What are token type IDs?](../glossary#token-type-ids)\n        modal_token_type_ids (*optional*) `torch.LongTensor` of shape `(batch_size, modal_sequence_length)`:\n            Segment token indices to indicate different portions of the non-text modality. The embeddings from these\n            tokens will be summed with the respective token embeddings for the non-text modality.\n        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):\n            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,\n            config.max_position_embeddings - 1]`.\n\n            [What are position IDs?](../glossary#position-ids)\n        modal_position_ids (`torch.LongTensor` of shape `(batch_size, modal_sequence_length)`, *optional*):\n            Indices of positions of each input sequence tokens in the position embeddings for the non-text modality.\n            Selected in the range `[0, config.max_position_embeddings - 1]`.\n\n            [What are position IDs?](../glossary#position-ids)\n        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):\n            Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:\n\n            - 1 indicates the head is **not masked**,\n            - 0 indicates the head is **masked**.\n\n        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, embedding_dim)`, *optional*):\n            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This\n            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the\n            model\'s internal embedding lookup matrix.\n        encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):\n            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if\n            the model is configured as a decoder.\n        encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):\n            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in\n            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:\n\n            - 1 for tokens that are **not masked**,\n            - 0 for tokens that are **masked**.\n\n        output_attentions (`bool`, *optional*):\n            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned\n            tensors for more detail.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'

@add_start_docstrings(
    "The bare MMBT Model outputting raw hidden-states without any specific head on top."					,    lowerCAmelCase_					,    )
class       lowerCamelCase_							(	nn.Module					,    lowerCAmelCase_  ):






   '''simple docstring'''

   def __init__(  self     ,						__UpperCamelCase     ,						__UpperCamelCase     ,						__UpperCamelCase ):
       super().__init__()
       UpperCamelCase__      =			config
       UpperCamelCase__      =			transformer
       UpperCamelCase__      =			ModalEmbeddings(__UpperCamelCase     ,						__UpperCamelCase     ,						transformer.embeddings )


   @add_start_docstrings_to_model_forward(__UpperCamelCase )
   @replace_return_docstrings(output_type=__UpperCamelCase     ,						config_class=_CONFIG_FOR_DOC )
   def 	_UpperCamelCase       (  self     ,						__UpperCamelCase     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						):
       UpperCamelCase__      =			output_attentions if output_attentions is not None else self.config.output_attentions
       UpperCamelCase__      =			(
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )
       UpperCamelCase__      =			return_dict if return_dict is not None else self.config.use_return_dict

       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

       UpperCamelCase__      =			self.modal_encoder(
           __UpperCamelCase     ,						start_token=__UpperCamelCase     ,						end_token=__UpperCamelCase     ,						position_ids=__UpperCamelCase     ,						token_type_ids=__UpperCamelCase     ,						)

       UpperCamelCase__      =			modal_embeddings.size()[:-1]

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

       UpperCamelCase__      =			self.transformer.embeddings(
           input_ids=__UpperCamelCase     ,						position_ids=__UpperCamelCase     ,						token_type_ids=__UpperCamelCase     ,						inputs_embeds=__UpperCamelCase )

       UpperCamelCase__      =			torch.cat([modal_embeddings, txt_embeddings]     ,						1 )

       UpperCamelCase__      =			embedding_output.size()[:-1]

       if attention_mask is None:
           UpperCamelCase__      =			torch.ones(__UpperCamelCase     ,						device=__UpperCamelCase )
       else:
           UpperCamelCase__      =			torch.cat(
               [torch.ones(__UpperCamelCase     ,						device=__UpperCamelCase     ,						dtype=torch.long ), attention_mask]     ,						dim=1 )
       if encoder_attention_mask is None:
           UpperCamelCase__      =			torch.ones(__UpperCamelCase     ,						device=__UpperCamelCase )
       else:
           UpperCamelCase__      =			torch.cat(
               [torch.ones(__UpperCamelCase     ,						device=__UpperCamelCase ), encoder_attention_mask]     ,						dim=1 )

       UpperCamelCase__      =			self.get_extended_attention_mask(__UpperCamelCase     ,						__UpperCamelCase )
       UpperCamelCase__      =			self.invert_attention_mask(__UpperCamelCase )
       UpperCamelCase__      =			self.get_head_mask(__UpperCamelCase     ,						self.config.num_hidden_layers )

       UpperCamelCase__      =			self.transformer.encoder(
           __UpperCamelCase     ,						attention_mask=__UpperCamelCase     ,						head_mask=__UpperCamelCase     ,						encoder_hidden_states=__UpperCamelCase     ,						encoder_attention_mask=__UpperCamelCase     ,						output_attentions=__UpperCamelCase     ,						output_hidden_states=__UpperCamelCase     ,						return_dict=__UpperCamelCase     ,						)

       UpperCamelCase__      =			encoder_outputs[0]
       UpperCamelCase__      =			self.transformer.pooler(__UpperCamelCase )

       if not return_dict:
           return (sequence_output, pooled_output) + encoder_outputs[1:]

       return BaseModelOutputWithPooling(
           last_hidden_state=__UpperCamelCase     ,						pooler_output=__UpperCamelCase     ,						hidden_states=encoder_outputs.hidden_states     ,						attentions=encoder_outputs.attentions     ,						)


   def 	_UpperCamelCase       (  self ):
       return self.embeddings.word_embeddings







   def 	_UpperCamelCase       (  self     ,						__UpperCamelCase ):
       UpperCamelCase__      =			value

@add_start_docstrings(
    "\n    MMBT Model with a sequence classification/regression head on top (a linear layer on top of the pooled output)\n    "					,    lowerCAmelCase_					,    lowerCAmelCase_					,    )
class       lowerCamelCase_							(	nn.Module  ):






   '''simple docstring'''

   def __init__(  self     ,						__UpperCamelCase     ,						__UpperCamelCase     ,						__UpperCamelCase ):
       super().__init__()
       UpperCamelCase__      =			config.num_labels

       UpperCamelCase__      =			MMBTModel(__UpperCamelCase     ,						__UpperCamelCase     ,						__UpperCamelCase )
       UpperCamelCase__      =			nn.Dropout(config.hidden_dropout_prob )
       UpperCamelCase__      =			nn.Linear(config.hidden_size     ,						config.num_labels )







   def 	_UpperCamelCase       (  self     ,						__UpperCamelCase     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None     ,						):
       UpperCamelCase__      =			return_dict if return_dict is not None else self.config.use_return_dict

       UpperCamelCase__      =			self.mmbt(
           input_modal=__UpperCamelCase     ,						input_ids=__UpperCamelCase     ,						modal_start_tokens=__UpperCamelCase     ,						modal_end_tokens=__UpperCamelCase     ,						attention_mask=__UpperCamelCase     ,						token_type_ids=__UpperCamelCase     ,						modal_token_type_ids=__UpperCamelCase     ,						position_ids=__UpperCamelCase     ,						modal_position_ids=__UpperCamelCase     ,						head_mask=__UpperCamelCase     ,						inputs_embeds=__UpperCamelCase     ,						return_dict=__UpperCamelCase     ,						)

       UpperCamelCase__      =			outputs[1]

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

       UpperCamelCase__      =			None
       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 ) )

       if not return_dict:
           UpperCamelCase__      =			(logits,) + outputs[2:]
           return ((loss,) + output) if loss is not None else output

       return SequenceClassifierOutput(
           loss=__UpperCamelCase     ,						logits=__UpperCamelCase     ,						hidden_states=outputs.hidden_states     ,						attentions=outputs.attentions     ,						)




