Method sample

tensorrt_llm/serve/scripts/benchmark_dataset.py:1235–1275

(self,
               tokenizer: PreTrainedTokenizerBase,
               num_requests: int,
               output_len: Optional[int] = None,
               enable_multimodal_chat: bool = False,
               **kwargs)

    }

    def sample(self,
               tokenizer: PreTrainedTokenizerBase,
               num_requests: int,
               output_len: Optional[int] = None,
               enable_multimodal_chat: bool = False,
               **kwargs) -> list:
        output_len = (output_len
                      if output_len is not None else self.DEFAULT_OUTPUT_LEN)

        # Collect prompts for batch processing
        prompts = []
        for item in self.data:
            if len(prompts) >= num_requests:
                break
            raw_prompt = item['turns'][0]

            # apply template
            formatted_prompt = tokenizer.apply_chat_template(
                [{
                    "role": "user",
                    "content": raw_prompt
                }],
                add_generation_prompt=True,
                tokenize=False)
            prompts.append(formatted_prompt)

        # Batch tokenize prompts
        prompt_lengths, _ = batch_tokenize_prompts(
            prompts, tokenizer, progress_name="MT-Bench prompts")

        # Create samples
        sampled_requests = []
        for prompt, prompt_len in zip(prompts, prompt_lengths):
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                ))
        self.maybe_oversample_requests(sampled_requests, num_requests)
        return sampled_requests


# -----------------------------------------------------------------------------