docstr: If labels are provided, they must be 0 for masked out tokens (as specified in the attention
  mask).
  Outputs:
      if `labels` and `next_sentence_label` are not `None`:
          Outputs the total_loss which is the sum of the masked language modeling loss and the next
          sentence classification loss.
      if `labels` or `next_sentence_label` is `None`:
          Outputs a tuple comprising
          - the masked language modeling logits of shape [batch_size, sequence_length, vocab_size], and
          - the next sentence classification logits of shape [batch_size, 2].
docstr_tok: ["If","labels","are","provided",",","they","must","be","0","for","masked","out","tokens","(","as","specified","in","the","attention","mask",")",".","Outputs",":","if","`","labels","`","and","`","next_sentence_label","`","are","not","`","None","`",":","Outputs","the","total_loss","which","is","the","sum","of","the","masked","language","modeling","loss","and","the","next","sentence","classification","loss",".","if","`","labels","`","or","`","next_sentence_label","`","is","`","None","`",":","Outputs","a","tuple","comprising","-","the","masked","language","modeling","logits","of","shape","[","batch_size",",","sequence_length",",","vocab_size","]",",","and","-","the","next","sentence","classification","logits","of","shape","[","batch_size",",","2","]","."]
code: def forward(
      self,
      input_ids,
      position_ids=None,
      token_type_ids=None,
      attention_mask=None,
      labels=None,
      next_sentence_label=None,
  ):
      """
      If labels are provided, they must be 0 for masked out tokens (as specified in the attention
      mask).
      Outputs:
          if `labels` and `next_sentence_label` are not `None`:
              Outputs the total_loss which is the sum of the masked language modeling loss and the next
              sentence classification loss.
          if `labels` or `next_sentence_label` is `None`:
              Outputs a tuple comprising
              - the masked language modeling logits of shape [batch_size, sequence_length, vocab_size], and
              - the next sentence classification logits of shape [batch_size, 2].
  
      """
      masked_tokens_mask = (
          labels > 0
          if (self.last_layer_subset and labels is not None)
          else None
      )
      outputs = self.bert(
          input_ids,
          position_ids=position_ids,
          token_type_ids=token_type_ids,
          attention_mask=attention_mask.bool()
          if attention_mask is not None
          else None,
          masked_tokens_mask=masked_tokens_mask,
      )
      sequence_output, pooled_output = (
          outputs.last_hidden_state,
          outputs.pooler_output,
      )
      if self.dense_seq_output and labels is not None:
          masked_token_idx = torch.nonzero(
              labels.flatten() > 0, as_tuple=False
          ).flatten()
          if not self.last_layer_subset:
              sequence_output = index_first_axis(
                  rearrange(sequence_output, "b s d -> (b s) d"),
                  masked_token_idx,
              )
      prediction_scores, seq_relationship_score = self.cls(
          sequence_output, pooled_output
      )
  
      total_loss = None
      if labels is not None and next_sentence_label is not None:
          if (
              self.dense_seq_output and labels is not None
          ):  # prediction_scores are already flattened
              masked_lm_loss = self.mlm_loss(
                  prediction_scores, labels.flatten()[masked_token_idx]
              )
          else:
              masked_lm_loss = self.mlm_loss(
                  rearrange(prediction_scores, "... v -> (...) v"),
                  rearrange(labels, "... -> (...)"),
              )
          next_sentence_loss = self.nsp_loss(
              rearrange(seq_relationship_score, "... t -> (...) t"),
              rearrange(next_sentence_label, "... -> (...)"),
          )
          total_loss = (
              masked_lm_loss.float() + next_sentence_loss.float()
          )
  
      return BertForPreTrainingOutput(
          loss=total_loss,
          prediction_logits=prediction_scores,
          seq_relationship_logits=seq_relationship_score,
      )
  
imports: import logging
  import re
  from collections import OrderedDict
  from collections.abc import Sequence
  from functools import partial
  from typing import Any, Mapping
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange
  from transformers import BertConfig, PretrainedConfig
  from transformers.models.bert.modeling_bert import BaseModelOutputWithPoolingAndCrossAttentions, BertForPreTrainingOutput
  from flash_attn.bert_padding import index_first_axis, index_first_axis_residual, pad_input, unpad_input
  from flash_attn.modules.block import Block
  from flash_attn.modules.embedding import BertEmbeddings
  from flash_attn.modules.mha import MHA
  from flash_attn.modules.mlp import FusedMLP, Mlp
  from flash_attn.utils.pretrained import state_dict_from_pretrained
label_desc: Use image node_id 3 for calling the BertForPreTraining obj's underlying member method code with example usage: obj.forward(input_ids, position_ids, token_type_ids, attention_mask, labels, next_sentence_label) and returns: BertForPreTrainingOutput