docstr: Concatenate multiple batches together by padding internal torch tensors
docstr_tok: ["Concatenate","multiple","batches","together","by","padding","internal","torch","tensors"]
code: def concatenate(
      cls, batches: List["Seq2SeqLMBatch"]
  ) -> "Seq2SeqLMBatch":
      """Concatenate multiple batches together by padding internal torch tensors"""
  
      # Used for padding
      total_batch_size = 0
      max_input_length = 0
      max_decoder_input_length = 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_decoder_input_length = max(
              max_decoder_input_length, batch.max_decoder_input_length
          )
          padding_right_offset = max(
              padding_right_offset, batch.padding_right_offset
          )
  
      # Batch attributes
      requests = []
      requests_idx_mapping = {}
      all_decoder_input_ids = []
      input_lengths = []
      decoder_input_lengths = []
      prefix_offsets = []
      read_offsets = []
      next_token_choosers = []
      stopping_criterias = []
      top_n_tokens = []
      max_tokens = 0
  
      # Batch tensors
      attention_mask = None
      decoder_input_ids = None
      decoder_attention_mask = None
      encoder_last_hidden_state = None
      top_n_tokens_tensor = 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):
          # Extend all list attributes
          requests.extend(batch.requests)
          all_decoder_input_ids.extend(batch.all_decoder_input_ids)
          input_lengths.extend(batch.input_lengths)
          decoder_input_lengths.extend(batch.decoder_input_lengths)
          prefix_offsets.extend(batch.prefix_offsets)
          read_offsets.extend(batch.read_offsets)
          next_token_choosers.extend(batch.next_token_choosers)
          stopping_criterias.extend(batch.stopping_criterias)
          top_n_tokens.extend(batch.top_n_tokens)
  
          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.encoder_last_hidden_state is None:
              raise ValueError(
                  "Batch encoder_last_hidden_state cannot be None"
              )
  
          # Create padded tensor
          if attention_mask is None:
              attention_mask = batch.attention_mask.new_zeros(
                  (total_batch_size, max_input_length),
              )
          # Copy to correct indices
          attention_mask[
              start_index:end_index, -batch.max_input_length :
          ] = batch.attention_mask[:, -batch.max_input_length :]
  
          # Create padded tensor
          if decoder_input_ids is None:
              decoder_input_ids = batch.decoder_input_ids.new_zeros(
                  (total_batch_size, 1),
              )
          # Copy to correct indices
          decoder_input_ids[
              start_index:end_index
          ] = batch.decoder_input_ids
  
          # Create padded tensor
          if decoder_attention_mask is None:
              # As decoder_attention_mask might not exist, we use `batch.attention_mask` for device here
              decoder_attention_mask = batch.attention_mask.new_zeros(
                  (
                      total_batch_size,
                      max_decoder_input_length + padding_right_offset,
                  ),
              )
          # If the decoder mask does not exist yet, all generations started at the same time and we never concatenated
          # this batch. All generations are of length `batch.max_decoder_input_length`.
          left_offset = (
              max_decoder_input_length - batch.max_decoder_input_length
          )
          if batch.decoder_attention_mask is None:
              decoder_attention_mask[
                  start_index:end_index,
                  left_offset:-padding_right_offset,
              ] = 1
          # If it exists, we need to index
          else:
              batch_left_offset = (
                  batch.decoder_attention_mask.shape[1]
                  - batch.max_decoder_input_length
                  - batch.padding_right_offset
              )
              decoder_attention_mask[
                  start_index:end_index,
                  left_offset:-padding_right_offset,
              ] = batch.decoder_attention_mask[
                  :,
                  batch_left_offset : -batch.padding_right_offset,
              ]
  
          # Create padded tensor
          if encoder_last_hidden_state is None:
              encoder_last_hidden_state = (
                  batch.encoder_last_hidden_state.new_zeros(
                      (
                          total_batch_size,
                          max_input_length,
                          batch.encoder_last_hidden_state.shape[-1],
                      ),
                  )
              )
  
          if top_n_tokens_tensor is None:
              top_n_tokens_tensor = batches[
                  0
              ].top_n_tokens_tensor.new_zeros(
                  total_batch_size,
              )
          top_n_tokens_tensor[
              start_index:end_index
          ] = batch.top_n_tokens_tensor
  
          # Copy to correct indices
          encoder_last_hidden_state[
              start_index:end_index, -batch.max_input_length :, :
          ] = batch.encoder_last_hidden_state[
              :, -batch.max_input_length :, :
          ]
          batch.encoder_last_hidden_state = None
  
          # Ensure that we can update tensors in-place
          if type(batch.past_key_values[0]) == tuple:
              batch.past_key_values = [
                  [t for t in layer] for layer in batch.past_key_values
              ]
  
          # Add eventual padding tokens that were added while concatenating
          max_tokens += batch.max_tokens + (
              max_input_length
              - batch.max_input_length
              + max_decoder_input_length
              - batch.max_decoder_input_length
          ) * len(batch)
  
          start_index = end_index
  
      # Determine shapes for new past kv tensors
      first_past_kvs = batches[0].past_key_values
      _, num_heads, _, head_dim = first_past_kvs[0][0].shape
  
      padded_dec_t_shape = (
          total_batch_size,
          num_heads,
          (max_decoder_input_length - 1),
          head_dim,
      )
  
      padded_enc_t_shape = (
          total_batch_size,
          num_heads,
          max_input_length,
          head_dim,
      )
  
      # Iterate over attention layers
      for j in range(len(first_past_kvs)):
          past_key_values.append([])
  
          # Decoder past
          for k in range(0, 2):
              # Initialize tensors
              padded_past_values = first_past_kvs[j][k].new_zeros(
                  padded_dec_t_shape
              )
              past_key_values[j].append(padded_past_values)
  
              start_index = 0
              for batch in batches:
                  t = batch.past_key_values[j][k]
                  # Clear reference to the original tensor
                  batch.past_key_values[j][k] = None
                  # Slicing end index for this batch
                  end_index = start_index + len(batch)
                  # We slice the past keys and values to remove the padding from previous batches
                  past_seq_len = batch.max_decoder_input_length - 1
                  padded_past_values[
                      start_index:end_index, :, -past_seq_len:, :
                  ] = t[:, :, -past_seq_len:, :]
                  del t
  
                  start_index = end_index
  
          # Encoder past
          for k in range(2, 4):
              # Initialize tensors
              padded_past_values = first_past_kvs[j][k].new_zeros(
                  padded_enc_t_shape
              )
              past_key_values[j].append(padded_past_values)
  
              start_index = 0
              for batch in batches:
                  t = batch.past_key_values[j][k]
                  # Clear reference to the original tensor
                  batch.past_key_values[j][k] = None
                  # Slicing end index for this batch
                  end_index = start_index + len(batch)
                  # We slice the past keys and values to remove the padding from previous batches
                  padded_past_values[
                      start_index:end_index,
                      :,
                      -batch.max_input_length :,
                      :,
                  ] = t[:, :, -batch.max_input_length :, :]
                  del t
  
                  start_index = end_index
  
      return cls(
          batch_id=batches[0].batch_id,
          requests=requests,
          requests_idx_mapping=requests_idx_mapping,
          input_ids=None,
          attention_mask=attention_mask,
          decoder_input_ids=decoder_input_ids,
          all_decoder_input_ids=all_decoder_input_ids,
          decoder_attention_mask=decoder_attention_mask,
          encoder_last_hidden_state=encoder_last_hidden_state,
          past_key_values=past_key_values,
          input_lengths=input_lengths,
          decoder_input_lengths=decoder_input_lengths,
          prefix_offsets=prefix_offsets,
          read_offsets=read_offsets,
          next_token_choosers=next_token_choosers,
          stopping_criterias=stopping_criterias,
          top_n_tokens=top_n_tokens,
          top_n_tokens_tensor=top_n_tokens_tensor,
          max_input_length=max_input_length,
          max_decoder_input_length=max_decoder_input_length,
          padding_right_offset=padding_right_offset,
          max_tokens=max_tokens,
      )
  
imports: import torch
  import time
  from dataclasses import dataclass
  from opentelemetry import trace
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, PreTrainedTokenizerBase
  from typing import Optional, Tuple, List, Type, Dict
  from text_generation_server.utils.tokens import batch_top_tokens
  from text_generation_server.models import Model
  from text_generation_server.models.types import GeneratedText, Batch, Generation, Tokens
  from text_generation_server.pb import generate_pb2
  from text_generation_server.utils import NextTokenChooser, StoppingCriteria, Sampling
label_desc: Use image node_id 4 for calling the Seq2SeqLMBatch obj's underlying member method code with example usage: obj.concatenate(cls, batches) and returns: cls