code: def from_pb(
      cls,
      pb: generate_pb2.Batch,
      tokenizer: PreTrainedTokenizerBase,
      dtype: torch.dtype,
      device: torch.device,
  ) -> "FlashCausalLMBatch":
      batch_inputs = []
      max_truncation = 0
      for r in pb.requests:
          batch_inputs.append(r.inputs)
          max_truncation = max(max_truncation, r.truncate)
  
      batch_tokenized_inputs = tokenizer(
          batch_inputs, truncation=True, max_length=max_truncation
      )["input_ids"]
  
      position_ids = []
      speculative_ids = []
      cu_seqlen_prefill = [0]
      needed_blocks_slots = []
      start_slots = []
      slot_indices = []
  
      input_lengths = []
      prefix_offsets = []
      read_offsets = []
      all_input_ids = []
      requests_idx_mapping = {}
  
      all_prefill_logprobs = True
      no_prefill_logprobs = True
      prefill_head_indices = []
      prefill_next_token_indices = []
      prefill_cu_outlens = [0]
  
      next_token_chooser_parameters = []
      stopping_criterias = []
      top_n_tokens = []
  
      # Cumulative length
      cumulative_length = 0
      cumulative_max_length = 0
      prefill_out_cumulative_length = 0
  
      blocks = 0
      max_seqlen = 0
      max_length = 0
      max_blocks = 0
  
      # Parse batch
      for i, (r, tokenized_input) in enumerate(
          zip(pb.requests, batch_tokenized_inputs)
      ):
          # request id -> idx in list mapping
          requests_idx_mapping[r.id] = i
  
          tokenized_input = tokenized_input[-r.truncate :]
  
          input_length = len(tokenized_input)
          input_lengths.append(input_length)
  
          prefix_offsets.append(input_length - 5)
          read_offsets.append(input_length)
  
          all_input_ids.append(tokenized_input)
  
          # Position ids
          request_position_ids = torch.arange(
              0, input_length, dtype=torch.int32
          )
          position_ids.append(request_position_ids)
  
          # Add cumulative lengths of all previous inputs
          cu_seqlen_prefill.append(cumulative_length + input_length)
  
          next_token_chooser_parameters.append(r.parameters)
  
          stopping_criteria = StoppingCriteria.from_pb(
              r.stopping_parameters, tokenizer
          )
          max_new_tokens = stopping_criteria.max_new_tokens
          stopping_criterias.append(stopping_criteria)
          top_n_tokens.append(r.top_n_tokens)
  
          # Paged attention
          # Remove one as the first token des not have a past
          speculative_length = get_speculate()
          total_tokens = (
              input_length + max_new_tokens - 1 + speculative_length
          )
          needed_blocks = math.ceil(total_tokens / BLOCK_SIZE)
          blocks += needed_blocks
          needed_blocks_slots.append((needed_blocks, total_tokens))
          start_slots.append(cumulative_max_length)
  
          request_slot_indices = torch.arange(
              cumulative_max_length,
              cumulative_max_length + input_length,
              dtype=torch.int64,
          )
          slot_indices.append(request_slot_indices)
  
          all_prefill_logprobs = (
              all_prefill_logprobs and r.prefill_logprobs
          )
          no_prefill_logprobs = (
              no_prefill_logprobs and not r.prefill_logprobs
          )
  
          if r.prefill_logprobs:
              prefill_head_indices.append(
                  request_position_ids + cumulative_length
              )
              prefill_next_token_indices.append(
                  prefill_out_cumulative_length + input_length - 1
              )
              prefill_cu_outlens.append(
                  prefill_out_cumulative_length + input_length
              )
              prefill_out_cumulative_length += input_length
          else:
              prefill_head_indices.append(
                  torch.tensor(
                      [cumulative_length + input_length - 1],
                      dtype=torch.int32,
                  )
              )
              prefill_next_token_indices.append(
                  prefill_out_cumulative_length
              )
              prefill_cu_outlens.append(
                  prefill_out_cumulative_length + 1
              )
              prefill_out_cumulative_length += 1
  
          # Update
          cumulative_length += input_length
          cumulative_max_length += total_tokens
          max_seqlen = max(max_seqlen, input_length)
          max_blocks = max(max_blocks, needed_blocks)
          max_length = max(
              max_length,
              input_length + max_new_tokens + speculative_length,
          )
  
      next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
          next_token_chooser_parameters, dtype, device
      )
      start_slots = torch.tensor(start_slots, dtype=torch.int64)
  
      # Padded all_input_ids_tensor
      all_input_ids_tensor = np.zeros(
          (len(all_input_ids), max_length), dtype=np.int64
      )
      for i, input_ids in enumerate(all_input_ids):
          all_input_ids_tensor[i, : len(input_ids)] = input_ids
  
      # Create tensors on device
      all_input_ids_tensor = torch.tensor(
          all_input_ids_tensor, dtype=torch.int64, device=device
      )
  
      if len(pb.requests) > 1:
          input_ids = np.concatenate(all_input_ids, dtype=np.int64)
          position_ids = torch.cat(position_ids)
          slot_indices = torch.cat(slot_indices)
      else:
          input_ids = all_input_ids[0]
          position_ids = position_ids[0]
          slot_indices = slot_indices[0]
  
      cu_seqlen_prefill = torch.tensor(
          cu_seqlen_prefill, device=device, dtype=torch.int32
      )
      position_ids = position_ids.to(device)
      slot_indices = slot_indices.to(device)
      input_ids = torch.tensor(
          input_ids, dtype=torch.int64, device=device
      )
      input_lengths_tensor = torch.tensor(
          input_lengths, dtype=torch.int32, device=device
      )
  
      if all_prefill_logprobs:
          prefill_head_indices = None
          prefill_next_token_indices = cu_seqlen_prefill[1:] - 1
      elif no_prefill_logprobs:
          prefill_head_indices = cu_seqlen_prefill[1:] - 1
          prefill_next_token_indices = None
      else:
          prefill_head_indices = torch.tensor(
              torch.cat(prefill_head_indices),
              dtype=torch.int64,
              device=device,
          )
          prefill_next_token_indices = torch.tensor(
              prefill_next_token_indices,
              dtype=torch.int64,
              device=device,
          )
      top_n_tokens_tensor = torch.tensor(
          top_n_tokens, device=device, dtype=torch.int64
      )
  
      return cls(
          batch_id=pb.id,
          requests=pb.requests,
          requests_idx_mapping=requests_idx_mapping,
          input_ids=input_ids,
          position_ids=position_ids,
          cu_seqlen_prefill=cu_seqlen_prefill,
          start_slots=start_slots,
          slot_indices=slot_indices,
          needed_blocks_slots=needed_blocks_slots,
          block_tables=None,
          block_tables_tensor=None,
          slots=None,
          max_seqlen=max_seqlen,
          prefill_head_indices=prefill_head_indices,
          prefill_next_token_indices=prefill_next_token_indices,
          prefill_cu_outlens=prefill_cu_outlens,
          input_lengths=input_lengths,
          input_lengths_tensor=input_lengths_tensor,
          prefix_offsets=prefix_offsets,
          read_offsets=read_offsets,
          all_input_ids=all_input_ids,
          all_input_ids_tensor=all_input_ids_tensor,
          next_token_chooser=next_token_chooser,
          stopping_criterias=stopping_criterias,
          top_n_tokens=top_n_tokens,
          top_n_tokens_tensor=top_n_tokens_tensor,
          blocks=blocks,
          max_blocks=max_blocks,
          speculative_ids=None,
      )
  
imports: import math
  import time
  import itertools
  import torch
  import torch.distributed
  import numpy
  from dataclasses import dataclass
  from opentelemetry import trace
  from transformers import PreTrainedTokenizerBase
  from typing import Optional, Tuple, List, Type, Dict
  from text_generation_server.models import Model
  from text_generation_server.utils.tokens import batch_top_tokens
  from text_generation_server.utils.speculate import get_speculate
  from text_generation_server.models.types import Batch, Tokens, Generation, GeneratedText
  from text_generation_server.models.cache_manager import get_cache_manager, set_cache_manager, BLOCK_SIZE
  from text_generation_server.pb import generate_pb2
  from text_generation_server.utils import StoppingCriteria, HeterogeneousNextTokenChooser
  from text_generation_server.utils.dist import MEMORY_FRACTION
label_desc: Use image node_id 2 for calling the FlashCausalLMBatch obj's underlying member method code with example usage: obj.from_pb(cls, pb, tokenizer, dtype, device) and returns: cls