code: def prepare_inputs_for_generation(
      self,
      input_ids: torch.LongTensor,
      past_key_values: Optional[torch.Tensor] = None,
      attention_mask: Optional[torch.Tensor] = None,
      inputs_embeds: Optional[torch.Tensor] = None,
      **kwargs,
  ) -> dict:
      # only last token for input_ids if past is not None
      if past_key_values:
          input_ids = input_ids[:, -1].unsqueeze(-1)
  
          # the cache may be in the stardard format (e.g. in contrastive search), convert to bloom's format if needed
          if past_key_values[0][0].shape[0] == input_ids.shape[0]:
              past_key_values = self._convert_to_bloom_cache(
                  past_key_values
              )
  
      # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
      if inputs_embeds is not None and past_key_values is None:
          model_inputs = {"inputs_embeds": inputs_embeds}
      else:
          model_inputs = {"input_ids": input_ids}
  
      model_inputs.update(
          {
              "past_key_values": past_key_values,
              "use_cache": kwargs.get("use_cache"),
              "attention_mask": attention_mask,
          }
      )
      return model_inputs
  
imports: import math
  import os
  import warnings
  from typing import Optional, Tuple, Union
  import torch
  import torch.distributed
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import LayerNorm
  from torch.nn import functional
  from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions
  from transformers import BloomConfig, PreTrainedModel
  from text_generation_server.utils.layers import TensorParallelColumnLinear, TensorParallelEmbedding, TensorParallelRowLinear, TensorParallelHead
label_desc: Use image node_id 2 for calling the BloomForCausalLM obj's underlying member method code with example usage: obj.prepare_inputs_for_generation(input_ids, past_key_values, attention_mask, inputs_embeds) and returns: model_inputs