class_docstr: ONNX model with a causal language modeling head for ONNX Runtime inference. This class officially supports bloom, codegen, falcon, gpt2, gpt_bigcode, gpt_neo, gpt_neox, gptj, llama.
class_docstr_tok: ["ONNX","model","with","a","causal","language","modeling","head","for","ONNX","Runtime","inference",".","This","class","officially","supports","bloom",",","codegen",",","falcon",",","gpt2",",","gpt_bigcode",",","gpt_neo",",","gpt_neox",",","gptj",",","llama","."]
code: def forward(
      self,
      input_ids: torch.LongTensor,
      attention_mask: Optional[torch.FloatTensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
      labels: Optional[torch.LongTensor] = None,
      use_cache_branch: bool = None,
      **kwargs,
  ) -> CausalLMOutputWithPast:
      # adding use_cache_branch in the signature here is just a hack for IO Binding
      use_torch = isinstance(input_ids, torch.Tensor)
      self.raise_on_numpy_input_io_binding(use_torch)
  
      inputs = {}
      known_output_shapes = {}
      use_cache_branch = None
      loss = None
      if self.use_cache:
          if past_key_values is not None:
              # Flatten the past_key_values (gpt_bigcode has fused key/value cache, so no need to flatten it)
              if self.model_type != "gpt_bigcode":
                  past_key_values = tuple(
                      past_key_value
                      for pkv_per_layer in past_key_values
                      for past_key_value in pkv_per_layer
                  )
  
          # Create dummy past_key_values for decoder first generation step if none given
          (
              use_cache_branch,
              past_key_values,
              known_output_shapes,
          ) = self.prepare_past_key_values(
              input_ids, past_key_values, use_torch
          )
  
      if self.use_io_binding:
          # TODO: fix transformers generate to have contiguous input_ids here already
          # For an unknown reason, calling `contiguous()` here is necessary to not have errors
          # on CPU EP with batch size > 1, despite it being also called in _prepare_io_binding.
          # I suspect the reason is the contiguous python list that messes something up?
          model_inputs = [input_ids.contiguous()]
  
          if "attention_mask" in self.inputs_names:
              model_inputs.append(attention_mask)
  
          if "position_ids" in self.inputs_names:
              if position_ids is None:
                  raise ValueError(
                      "position_ids was not passed but is a required input for this ONNX model."
                  )
              model_inputs.append(position_ids.contiguous())
  
          if past_key_values is not None:
              model_inputs += past_key_values
  
          if use_cache_branch is not None:
              model_inputs.append(use_cache_branch)
  
          if "labels" in self.inputs_names:
              model_inputs.append(labels)
              known_output_shapes.update({"loss": []})
  
          (
              io_binding,
              output_shapes,
              output_buffers,
          ) = self._prepare_io_binding(
              self.model,
              *model_inputs,
              known_output_shapes=known_output_shapes,
              ordered_input_names=self._ordered_input_names,
          )
  
          if self.device.type == "cpu":
              self.model.run_with_iobinding(io_binding)
          else:
              io_binding.synchronize_inputs()
              self.model.run_with_iobinding(io_binding)
              io_binding.synchronize_outputs()
  
          if self.use_cache:
              # Tuple of length equal to : number of layer * number of past_key_value per decoder layer(2)
              past_key_values = ()
              for name in self.key_value_output_names:
                  past_key_values += (
                      output_buffers[name].view(output_shapes[name]),
                  )
  
          logits = output_buffers["logits"].view(
              output_shapes["logits"]
          )
  
          if "loss" in self.output_names:
              loss = output_buffers["loss"].view(output_shapes["loss"])
      else:
          inputs["input_ids"] = (
              input_ids.cpu().detach().numpy()
              if use_torch
              else input_ids
          )
  
          if "attention_mask" in self.inputs_names:
              inputs["attention_mask"] = (
                  attention_mask.cpu().detach().numpy()
                  if use_torch
                  else attention_mask
              )
  
          if "labels" in self.inputs_names:
              inputs["labels"] = (
                  labels.cpu().detach().numpy() if use_torch else labels
              )
  
          if "position_ids" in self.inputs_names:
              if position_ids is None:
                  raise ValueError(
                      "position_ids was not passed but is a required input for this ONNX model."
                  )
              inputs["position_ids"] = (
                  position_ids.cpu().detach().numpy()
                  if use_torch
                  else position_ids
              )
  
          # Add the past_key_values to the decoder inputs
          if past_key_values is not None:
              for input_name, past_key_value in zip(
                  self.key_value_input_names, past_key_values
              ):
                  inputs[input_name] = (
                      past_key_value.cpu().detach().numpy()
                      if use_torch
                      else past_key_value
                  )
  
          if use_cache_branch is not None:
              inputs["use_cache_branch"] = (
                  use_cache_branch.cpu().detach().numpy()
                  if use_torch
                  else use_cache_branch
              )
  
          outputs = self.model.run(None, inputs)
  
          if self.use_cache:
              # Tuple of length equal to : number of layer * number of past_key_value per decoder layer (2 for the self-attention)
              past_key_values = tuple(
                  torch.from_numpy(outputs[self.output_names[key]]).to(
                      self.device
                  )
                  for key in self.key_value_output_names
              )
  
          logits = torch.from_numpy(
              outputs[self.output_names["logits"]]
          ).to(self.device)
          if "loss" in self.output_names:
              loss = torch.from_numpy(
                  outputs[self.output_names["loss"]]
              ).to(self.device)
  
      if self.use_cache and self.model_type != "gpt_bigcode":
          # Tuple of tuple of length `n_layers`, with each tuple of length equal to the number of self-attention and
          # per decoder layer
          past_key_values = tuple(
              past_key_values[i : i + self.num_pkv]
              for i in range(0, len(past_key_values), self.num_pkv)
          )
  
      return CausalLMOutputWithPast(
          loss=loss, logits=logits, past_key_values=past_key_values
      )
  
imports: import logging
  from pathlib import Path
  from tempfile import TemporaryDirectory
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
  import numpy
  import onnx
  import torch
  from onnx.tools import update_model_dims
  from transformers import AutoModelForCausalLM, GenerationConfig
  from transformers.file_utils import add_end_docstrings, add_start_docstrings_to_model_forward
  from transformers.modeling_outputs import CausalLMOutputWithPast
  import onnxruntime
  from ..exporters.onnx import MODEL_TYPES_REQUIRING_POSITION_IDS, main_export
  from ..onnx.utils import check_model_uses_external_data
  from ..utils import NormalizedConfigManager, check_if_transformers_greater
  from ..utils.modeling_utils import MODEL_TO_PATCH_FOR_PAST
  from ..utils.save_utils import maybe_save_preprocessors
  from .constants import DECODER_MERGED_ONNX_FILE_PATTERN, DECODER_ONNX_FILE_PATTERN, DECODER_WITH_PAST_ONNX_FILE_PATTERN
  from .modeling_ort import ONNX_MODEL_END_DOCSTRING, ORTModel
  from .models.bloom import bloom_convert_to_bloom_cache, bloom_convert_to_standard_cache
  from .utils import ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME, ONNX_WEIGHTS_NAME
label_desc: Use image node_id 2 for calling the ORTModelForCausalLM obj's underlying member method code with example usage: obj.forward(input_ids, attention_mask, position_ids, past_key_values, labels, use_cache_branch) and returns: CausalLMOutputWithPast