class_docstr: This class serves as a micro library for forward function substitution of OPT models
  under pipeline setting.
class_docstr_tok: ["This","class","serves","as","a","micro","library","for","forward","function","substitution","of","OPT","models","under","pipeline","setting","."]
docstr: This function is modified on the basis of transformers.models.opt.modeling_opt.OPTForQuestionAnswering.forward.
  Please refer to original code of transformers for more details.
docstr_tok: ["This","function","is","modified","on","the","basis","of","transformers.models.opt.modeling_opt.OPTForQuestionAnswering.forward",".","Please","refer","to","original","code","of","transformers","for","more","details","."]
code: def opt_for_question_answering_forward(
      self: OPTForQuestionAnswering,
      input_ids: Optional[torch.LongTensor] = None,
      attention_mask: Optional[torch.FloatTensor] = None,
      head_mask: Optional[torch.FloatTensor] = None,
      past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
      inputs_embeds: Optional[torch.FloatTensor] = None,
      start_positions: Optional[torch.LongTensor] = None,
      end_positions: Optional[torch.LongTensor] = None,
      use_cache: Optional[bool] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
      stage_manager: Optional[PipelineStageManager] = None,
      hidden_states: Optional[torch.FloatTensor] = None,
      stage_index: Optional[List[int]] = None,
  ) -> Union[Tuple, QuestionAnsweringModelOutput]:
      r"""
      This function is modified on the basis of transformers.models.opt.modeling_opt.OPTForQuestionAnswering.forward.
      Please refer to original code of transformers for more details.
      """
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
  
      transformer_outputs = OPTPipelineForwards.opt_model_forward(
          self.model,
          input_ids,
          past_key_values=past_key_values,
          attention_mask=attention_mask,
          head_mask=head_mask,
          inputs_embeds=inputs_embeds,
          use_cache=use_cache,
          output_attentions=output_attentions,
          output_hidden_states=output_hidden_states,
          return_dict=return_dict,
          stage_manager=stage_manager,
          hidden_states=hidden_states,
          stage_index=stage_index,
      )
      if stage_manager.is_last_stage():
          hidden_states = transformer_outputs[0]
  
          logits = self.qa_outputs(hidden_states)
          start_logits, end_logits = logits.split(1, dim=-1)
          start_logits = start_logits.squeeze(-1).contiguous()
          end_logits = end_logits.squeeze(-1).contiguous()
  
          total_loss = None
          if start_positions is not None and end_positions is not None:
              # If we are on multi-GPU, split add a dimension
              if len(start_positions.size()) > 1:
                  start_positions = start_positions.squeeze(-1)
              if len(end_positions.size()) > 1:
                  end_positions = end_positions.squeeze(-1)
              # sometimes the start/end positions are outside our model inputs, we ignore these terms
              ignored_index = start_logits.size(1)
              start_positions = start_positions.clamp(0, ignored_index)
              end_positions = end_positions.clamp(0, ignored_index)
  
              loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
              start_loss = loss_fct(start_logits, start_positions)
              end_loss = loss_fct(end_logits, end_positions)
              total_loss = (start_loss + end_loss) / 2
  
          if not return_dict:
              output = (start_logits, end_logits) + transformer_outputs[
                  2:
              ]
              return (
                  ((total_loss,) + output)
                  if total_loss is not None
                  else output
              )
  
          return QuestionAnsweringModelOutput(
              loss=total_loss,
              start_logits=start_logits,
              end_logits=end_logits,
              hidden_states=transformer_outputs.hidden_states,
              attentions=transformer_outputs.attentions,
          )
      else:
          hidden_states = transformer_outputs.get("hidden_states")
          return {"hidden_states": hidden_states}
  
file_path: colossalai/colossalai/shardformer/modeling/opt.py
imports: import random
  from typing import List, Optional, Tuple, Union
  import torch
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast
  from transformers.models.opt.modeling_opt import OPTForCausalLM, OPTForQuestionAnswering, OPTForSequenceClassification, OPTModel
  from transformers.utils import logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
label_desc: Use image node_id 6 for calling the OPTPipelineForwards obj's underlying member method code with example usage: obj.opt_for_question_answering_forward(input_ids, attention_mask, head_mask, past_key_values, inputs_embeds, start_positions, end_positions, use_cache, output_attentions, output_hidden_states, return_dict, stage_manager, hidden_states, stage_index) and returns: QuestionAnsweringModelOutput, dict, unknown