code: def _sample(
      probs: torch.Tensor,
      logprobs: torch.Tensor,
      input_metadata: InputMetadata,
      st_timestamp: Optional[float] = None,
  ) -> SamplerOutput:
      categorized_seq_group_ids = {t: [] for t in SamplingType}
      category_num_tokens = {t: 0 for t in SamplingType}
      for i, seq_group in enumerate(input_metadata.seq_groups):
          seq_ids, sampling_params = seq_group
          sampling_type = sampling_params.sampling_type
          categorized_seq_group_ids[sampling_type].append(i)
          num_seqs = len(seq_ids)
          category_num_tokens[sampling_type] += num_seqs
  
      seq_outputs_dict: Dict[int, List[SequenceOutputs]] = {}
      category_start_idx = 0
      for sampling_type in SamplingType:
          seq_group_ids = categorized_seq_group_ids[sampling_type]
          seq_groups = [
              input_metadata.seq_groups[i] for i in seq_group_ids
          ]
          is_prompts = [
              i < input_metadata.num_prompts for i in seq_group_ids
          ]
          num_tokens = category_num_tokens[sampling_type]
          if num_tokens == 0:
              continue
          category_logprobs = logprobs[
              category_start_idx : category_start_idx + num_tokens
          ]
          category_probs = probs[
              category_start_idx : category_start_idx + num_tokens
          ]
          if sampling_type == SamplingType.GREEDY:
              sample_results = _greedy_sample(
                  seq_groups, category_logprobs
              )
          elif sampling_type == SamplingType.RANDOM:
              sample_results = _random_sample(
                  seq_groups, is_prompts, category_probs
              )
          elif sampling_type == SamplingType.BEAM:
              sample_results = _beam_search_sample(
                  seq_groups,
                  is_prompts,
                  input_metadata.seq_data,
                  category_logprobs,
              )
  
          # Batched query for logprobs of selected token
          batched_logprobs_query_seq_indices: List[int] = []
          batched_logprobs_query_token_indices: List[int] = []
          sample_idx = 0
          for seq_group_id, seq_group, sample_result in zip(
              seq_group_ids, seq_groups, sample_results
          ):
              seq_ids, sampling_params = seq_group
              next_token_ids, parent_ids = sample_result
              num_parent_seqs = len(seq_ids)
              batched_logprobs_query_seq_indices.extend(
                  [sample_idx + parent_id for parent_id in parent_ids]
              )
              batched_logprobs_query_token_indices.extend(
                  next_token_ids
              )
              sample_idx += num_parent_seqs
          batched_logprobs_query_result = category_logprobs[
              [
                  batched_logprobs_query_seq_indices,
                  batched_logprobs_query_token_indices,
              ]
          ].tolist()
  
          # Build the sequence outputs.
          sample_idx = 0
          result_idx = 0
          for seq_group_id, seq_group, sample_result in zip(
              seq_group_ids, seq_groups, sample_results
          ):
              seq_ids, sampling_params = seq_group
              next_token_ids, parent_ids = sample_result
              num_results = len(next_token_ids)
              num_parent_seqs = len(seq_ids)
              parent_logprobs = category_logprobs[
                  sample_idx : sample_idx + num_parent_seqs
              ]
              selected_token_logprobs = batched_logprobs_query_result[
                  result_idx : result_idx + num_results
              ]
              seq_output = _build_sequence_outputs(
                  parent_ids,
                  next_token_ids,
                  selected_token_logprobs,
                  seq_ids,
                  parent_logprobs,
                  sampling_params.logprobs,
                  st_timestamp,
              )
              seq_outputs_dict[seq_group_id] = seq_output
              sample_idx += num_parent_seqs
              result_idx += num_results
          category_start_idx += num_tokens
  
      return [
          seq_outputs_dict[i]
          for i in range(len(input_metadata.seq_groups))
      ]
  
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 13 for calling a global function with example usage: _sample(probs, logprobs, input_metadata, st_timestamp) and returns: unknown