code: def _prune_hidden_states(
      hidden_states: torch.Tensor,
      input_metadata: InputMetadata,
  ) -> torch.Tensor:
      last_token_indices = {t: [] for t in SamplingType}
      start_idx = 0
      for i, seq_group in enumerate(input_metadata.seq_groups):
          seq_ids, sampling_params = seq_group
          sampling_type = sampling_params.sampling_type
          if i < input_metadata.num_prompts:
              prompt_len = input_metadata.prompt_lens[i]
              last_token_indices[sampling_type].append(
                  start_idx + prompt_len - 1
              )
              start_idx += prompt_len
          else:
              num_seqs = len(seq_ids)
              last_token_indices[sampling_type].extend(
                  range(start_idx, start_idx + num_seqs)
              )
              start_idx += num_seqs
  
      all_last_token_indices = []
      for sampling_type in SamplingType:
          all_last_token_indices.extend(
              last_token_indices[sampling_type]
          )
      all_last_token_indices = torch.tensor(
          all_last_token_indices,
          dtype=torch.long,
          device=hidden_states.device,
      )
      return hidden_states.index_select(0, all_last_token_indices)
  
imports: from typing import Dict, List, Optional, Tuple
  import torch
  import torch.nn
  from bigdl.llm.vllm.model_executor.input_metadata import InputMetadata
  from bigdl.llm.vllm.sampling_params import SamplingParams, SamplingType
  from bigdl.llm.vllm.sequence import SamplerOutput, SequenceGroupMetadata, SequenceData, SequenceOutputs
  import time
label_desc: Use image node_id 1 for calling a global function with example usage: _prune_hidden_states(hidden_states, input_metadata) and returns: hidden_states