MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT-LLM / sample

Method sample

tensorrt_llm/serve/scripts/benchmark_dataset.py:804–856  ·  view source on GitHub ↗
(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        lora_path: Optional[str] = None,
        max_loras: Optional[int] = None,
        output_len: Optional[int] = None,
        enable_multimodal_chat: bool = False,
        **kwargs,
    )

Source from the content-addressed store, hash-verified

802 random.shuffle(self.data)
803
804 def sample(
805 self,
806 tokenizer: PreTrainedTokenizerBase,
807 num_requests: int,
808 lora_path: Optional[str] = None,
809 max_loras: Optional[int] = None,
810 output_len: Optional[int] = None,
811 enable_multimodal_chat: bool = False,
812 **kwargs,
813 ) -> list:
814 if enable_multimodal_chat:
815 raise NotImplementedError
816
817 # Collect prompts and completions for batch processing
818 prompts = []
819 completions = []
820
821 for entry in self.data:
822 if len(prompts) >= num_requests:
823 break
824 prompt, completion = (
825 entry["conversations"][0]["value"],
826 entry["conversations"][1]["value"],
827 )
828 prompts.append(prompt)
829 completions.append(completion)
830
831 # Batch tokenize prompts and completions
832 prompt_lengths, _ = batch_tokenize_prompts(
833 prompts, tokenizer, progress_name="ShareGPT prompts")
834 completion_lengths, _ = batch_tokenize_prompts(
835 completions, tokenizer, progress_name="ShareGPT completions")
836
837 # Filter and create samples
838 samples: list = []
839 for prompt, completion, prompt_len, completion_len in zip(
840 prompts, completions, prompt_lengths, completion_lengths):
841 new_output_len = completion_len if output_len is None else output_len
842 if not is_valid_sequence(prompt_len,
843 new_output_len,
844 skip_min_output_len_check=output_len
845 is not None):
846 continue
847
848 samples.append(
849 SampleRequest(
850 prompt=prompt,
851 prompt_len=prompt_len,
852 expected_output_len=new_output_len,
853 ))
854
855 self.maybe_oversample_requests(samples, num_requests)
856 return samples
857
858
859# -----------------------------------------------------------------------------

Callers

nothing calls this directly

Calls 5

batch_tokenize_promptsFunction · 0.85
is_valid_sequenceFunction · 0.85
SampleRequestClass · 0.85
appendMethod · 0.45

Tested by

no test coverage detected