code: def concatenate(
      cls, batches: List["IdeficsCausalLMBatch"]
  ) -> "IdeficsCausalLMBatch":
      # It adds new requests to the batch
      # Used for padding
      total_batch_size = 0
      max_input_length = 0
      max_num_images = 0
      padding_right_offset = 0
      for batch in batches:
          total_batch_size += len(batch)
          max_input_length = max(
              max_input_length, batch.max_input_length
          )
          max_num_images = max(
              max_num_images, batch.pixel_values.size(1)
          )
          padding_right_offset = max(
              padding_right_offset, batch.padding_right_offset
          )
  
      # Batch attributes
      requests = []
      requests_idx_mapping = {}
      input_lengths = []
      prefix_offsets = []
      read_offsets = []
      all_input_ids = []
      next_token_choosers = []
      stopping_criterias = []
      max_tokens = 0
  
      # Batch tensors
      input_ids = None
      attention_mask = None
      position_ids = None
      pixel_values = None
      image_hidden_states = None
      image_attention_mask = None
      past_key_values = []
  
      # Used for slicing correctly inside the tensors
      # Equivalent to a cumsum on batch sizes
      start_index = 0
      for i, batch in enumerate(batches):
          requests.extend(batch.requests)
          input_lengths.extend(batch.input_lengths)
          prefix_offsets.extend(batch.prefix_offsets)
          read_offsets.extend(batch.read_offsets)
          all_input_ids.extend(batch.all_input_ids)
          next_token_choosers.extend(batch.next_token_choosers)
          stopping_criterias.extend(batch.stopping_criterias)
  
          if i == 0:
              requests_idx_mapping = batch.requests_idx_mapping
          else:
              # We need to offset the mapping for each batch by the cumulative batch size
              for k, v in batch.requests_idx_mapping.items():
                  requests_idx_mapping[k] = v + start_index
  
          # Slicing end index for this batch
          end_index = start_index + len(batch)
  
          # We only concatenate batches that did at least one step
          if batch.past_key_values is None:
              raise ValueError("only concatenate prefilled batches")
  
          # Create empty tensor
          # input_ids is always of shape [batch_size, 1]
          # We do not need to pad it
          if input_ids is None:
              input_ids = batch.input_ids.new_empty(
                  (total_batch_size, 1)
              )
          # Copy to correct indices
          input_ids[start_index:end_index] = batch.input_ids
  
          # Create padded tensor
          if attention_mask is None:
              attention_mask = batch.attention_mask.new_zeros(
                  (
                      total_batch_size,
                      max_input_length + padding_right_offset,
                  ),
              )
  
          curr_batch_max_num_images = batch.pixel_values.size(1)
          if pixel_values is None:
              pixel_values = batch.pixel_values.new_zeros(
                  (total_batch_size, max_num_images, 3, 224, 224)
              )
          pixel_values[
              start_index:end_index, :curr_batch_max_num_images
          ] = batch.pixel_values
  
          if image_attention_mask is None:
              image_attention_mask = (
                  batch.image_attention_mask.new_zeros(
                      (
                          total_batch_size,
                          max_input_length + padding_right_offset,
                          max_num_images,
                      )
                  )
              )
  
          # We need to slice the attention mask to remove padding from previous steps
          # and to remove unused allocated space
          left_offset = max_input_length - batch.max_input_length
          batch_left_offset = (
              batch.attention_mask.shape[1]
              - batch.max_input_length
              - batch.padding_right_offset
          )
          attention_mask[
              start_index:end_index,
              left_offset:-padding_right_offset,
          ] = batch.attention_mask[
              :,
              batch_left_offset : -batch.padding_right_offset,
          ]
          image_attention_mask[
              start_index:end_index,
              left_offset:-padding_right_offset,
              :curr_batch_max_num_images,
          ] = batch.image_attention_mask[
              :, batch_left_offset : -batch.padding_right_offset, :
          ]
  
          # Create empty tensor
          # position_ids is always of shape [batch_size, 1]
          if position_ids is None:
              position_ids = batch.position_ids.new_empty(
                  (total_batch_size, 1)
              )
          position_ids[start_index:end_index] = batch.position_ids
  
          # Shenanigans to get dimensions because BLOOM outputs a past with a different shape
          # BLOOM Keys:   [batch_size * num_heads, head_dim, seq_length]
          # BLOOM Values: [batch_size * num_heads, seq_length, head_dim]
          # And ensure that we can update tensors in-place
          if type(batch.past_key_values[0]) == tuple:
              batch.past_key_values = [
                  [t.view(len(batch), -1, *t.shape[-2:]) for t in layer]
                  for layer in batch.past_key_values
              ]
          elif len(batch.past_key_values[0][0].shape) == 3:
              for layer in batch.past_key_values:
                  for k, t in enumerate(layer):
                      layer[k] = t.view(len(batch), -1, *t.shape[-2:])
  
          # Add eventual padding tokens that were added while concatenating
          max_tokens += batch.max_tokens + (
              max_input_length - batch.max_input_length
          ) * len(batch)
  
          start_index = end_index
  
      first_past_kvs = batches[0].past_key_values
      _, num_heads, padded_sequence_length, head_dim = first_past_kvs[
          0
      ][1].shape
  
      padded_past_values_shape = (
          total_batch_size,
          num_heads,
          max_input_length - 1,
          head_dim,
      )
  
      if batches[0].keys_head_dim_last:
          padded_past_keys_shape = padded_past_values_shape
      else:
          # seq_length is last for BLOOM
          padded_past_keys_shape = (
              total_batch_size,
              num_heads,
              head_dim,
              max_input_length - 1,
          )
  
      # Iterate over attention layers
      # Concatenate past key values layer by layer to allow incremental garbage collection
      for j in range(len(first_past_kvs)):
          padded_past_keys = first_past_kvs[j][0].new_zeros(
              padded_past_keys_shape
          )
          start_index = 0
          for batch in batches:
              past_keys = batch.past_key_values[j][0]
              # Clear reference to the original tensor
              batch.past_key_values[j][0] = None
  
              # Slicing end index for this batch
              end_index = start_index + len(batch)
              # We slice the keys to remove the padding from previous batches
              past_seq_len = batch.max_input_length - 1
              if batch.keys_head_dim_last:
                  padded_past_keys[
                      start_index:end_index, :, -past_seq_len:, :
                  ] = past_keys[:, :, -past_seq_len:, :]
              else:
                  # BLOOM case
                  padded_past_keys[
                      start_index:end_index, :, :, -past_seq_len:
                  ] = past_keys[:, :, :, -past_seq_len:]
              del past_keys
  
              start_index = end_index
  
          padded_past_values = first_past_kvs[j][1].new_zeros(
              padded_past_values_shape
          )
          start_index = 0
          for batch in batches:
              past_values = batch.past_key_values[j][1]
              # Clear reference to the original tensor
              batch.past_key_values[j][1] = None
  
              # Slicing end index for this batch
              end_index = start_index + len(batch)
              # We slice the past values to remove the padding from previous batches
              past_seq_len = batch.max_input_length - 1
              padded_past_values[
                  start_index:end_index, :, -past_seq_len:, :
              ] = past_values[:, :, -past_seq_len:, :]
              del past_values
  
              # Update values
              start_index = end_index
  
          past_key_values.append([padded_past_keys, padded_past_values])
  
      return cls(
          batch_id=batches[0].batch_id,
          requests=requests,
          requests_idx_mapping=requests_idx_mapping,
          input_ids=input_ids,
          attention_mask=attention_mask,
          position_ids=position_ids,
          pixel_values=pixel_values,
          image_hidden_states=image_hidden_states,
          image_attention_mask=image_attention_mask,
          past_key_values=past_key_values,
          all_input_ids=all_input_ids,
          input_lengths=input_lengths,
          prefix_offsets=prefix_offsets,
          read_offsets=read_offsets,
          next_token_choosers=next_token_choosers,
          stopping_criterias=stopping_criterias,
          max_input_length=max_input_length,
          padding_right_offset=padding_right_offset,
          keys_head_dim_last=batches[0].keys_head_dim_last,
          max_tokens=max_tokens,
      )
  
imports: import torch
  import time
  from dataclasses import dataclass
  from opentelemetry import trace
  from transformers import AutoProcessor, AutoTokenizer, PreTrainedTokenizerBase, ProcessorMixin
  from typing import Optional, Tuple, List, Type, Dict
  from text_generation_server.models import Model
  from text_generation_server.models.types import Batch, Tokens, Generation, GeneratedText
  from text_generation_server.pb import generate_pb2
  from text_generation_server.utils import NextTokenChooser, StoppingCriteria, Sampling
  import re
label_desc: Use image node_id 4 for calling the IdeficsCausalLMBatch obj's underlying member method code with example usage: obj.concatenate(cls, batches) and returns: cls