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

Method sample

tensorrt_llm/serve/scripts/benchmark_dataset.py:888–938  ·  view source on GitHub ↗
(
        self,
        tokenizer,
        num_requests: int,
        prefix_len: int = DEFAULT_PREFIX_LEN,
        input_len: int = DEFAULT_INPUT_LEN,
        output_len: int = DEFAULT_OUTPUT_LEN,
        return_prompt_formatted: bool = False,
        **kwargs,
    )

Source from the content-addressed store, hash-verified

886 self.data = f.readlines()
887
888 def sample(
889 self,
890 tokenizer,
891 num_requests: int,
892 prefix_len: int = DEFAULT_PREFIX_LEN,
893 input_len: int = DEFAULT_INPUT_LEN,
894 output_len: int = DEFAULT_OUTPUT_LEN,
895 return_prompt_formatted: bool = False,
896 **kwargs,
897 ) -> list:
898 # Calculate average token length for poem lines using batch tokenization
899 line_lengths, _ = batch_tokenize_prompts(self.data,
900 tokenizer,
901 progress_name="sonnet lines")
902 avg_len = sum(line_lengths) / len(line_lengths)
903
904 # Build the base prompt.
905 base_prompt = "Pick as many lines as you can from these poem lines:\n"
906 base_msg = [{"role": "user", "content": base_prompt}]
907 base_fmt = tokenizer.apply_chat_template(base_msg,
908 add_generation_prompt=True,
909 tokenize=False)
910 base_offset = len(tokenizer(base_fmt).input_ids)
911 if input_len <= base_offset:
912 raise ValueError(
913 f"'input_len' must be higher than the base prompt length "
914 f"({base_offset}).")
915
916 # Determine how many poem lines to use.
917 num_input_lines = round((input_len - base_offset) / avg_len)
918 num_prefix_lines = max(round((prefix_len - base_offset) / avg_len), 0)
919 prefix_lines = self.data[:num_prefix_lines]
920
921 samples = []
922 while len(samples) < num_requests:
923 extra_lines = random.choices(self.data,
924 k=num_input_lines - num_prefix_lines)
925 prompt = f"{base_prompt}{''.join(prefix_lines + extra_lines)}"
926 msg = [{"role": "user", "content": prompt}]
927 prompt_formatted = tokenizer.apply_chat_template(
928 msg, add_generation_prompt=True, tokenize=False)
929 prompt_len = len(tokenizer(prompt_formatted).input_ids)
930 if prompt_len <= input_len:
931 samples.append(
932 SampleRequest(
933 prompt=prompt_formatted
934 if return_prompt_formatted else prompt,
935 prompt_len=prompt_len,
936 expected_output_len=output_len,
937 ))
938 return samples
939
940
941# -----------------------------------------------------------------------------

Callers 1

mainFunction · 0.95

Calls 7

batch_tokenize_promptsFunction · 0.85
sumFunction · 0.85
maxFunction · 0.85
SampleRequestClass · 0.85
choicesMethod · 0.80
apply_chat_templateMethod · 0.45
appendMethod · 0.45

Tested by

no test coverage detected