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

Function create_dataset_from_stream

tensorrt_llm/bench/utils/data.py:35–177  ·  view source on GitHub ↗

Generate metadata and a list of requests to drive benchmarking. Args: tokenizer (PreTrainedTokenizer): HuggingFace tokenizer. stream (TextIO): Stream of input requests. max_input_length (int, optional): Maximum input length to cap prompts to. Defaults to 0. max_o

(
    tokenizer: PreTrainedTokenizer,
    stream: TextIO,
    max_input_length: int = 0,
    max_output_length: int = 0,
    num_requests: int = 0,
    model_dir: str = None,
    model_type: str = None,
    modality: str = None,
    image_data_format: str = "pt",
    data_device: str = "cpu",
    max_input_seq_len_for_multimodal: int = 4096,
)

Source from the content-addressed store, hash-verified

33
34
35def create_dataset_from_stream(
36 tokenizer: PreTrainedTokenizer,
37 stream: TextIO,
38 max_input_length: int = 0,
39 max_output_length: int = 0,
40 num_requests: int = 0,
41 model_dir: str = None,
42 model_type: str = None,
43 modality: str = None,
44 image_data_format: str = "pt",
45 data_device: str = "cpu",
46 max_input_seq_len_for_multimodal: int = 4096,
47) -> Tuple[DatasetMetadata, List[InferenceRequest]]:
48 """Generate metadata and a list of requests to drive benchmarking.
49
50 Args:
51 tokenizer (PreTrainedTokenizer): HuggingFace tokenizer.
52 stream (TextIO): Stream of input requests.
53 max_input_length (int, optional): Maximum input length to cap prompts to. Defaults to 0.
54 max_output_length (int, optional): Maximum output length to cap prompts to.. Defaults to 0.
55 num_requests (int, optional): Number of requests to limit to. Defaults to 0.
56
57 Returns:
58 Tuple[DatasetMetadata, List[InferenceRequest]]: A tuple containing a dataclass of dataset
59 statistics and a list of inference requests for benchmarking.
60 """
61 # Initialize dataset list, and metadata tracking variables.
62 dataset = []
63 max_requests = num_requests if num_requests > 0 else float("inf")
64
65 # If we're limiting the input length to a certain size, then set up
66 # a partial to truncate the data down to size. Otherwise, just use the
67 # unmodified tokenizer callable.
68 tokenize = (partial(
69 tokenizer,
70 padding="max_length",
71 max_length=max_input_length,
72 truncation=True,
73 ) if max_input_length > 0 else tokenizer)
74
75 # If we need to limit the output length, fill in a partial callable
76 # for max, otherwise a lambda that just returns x with no bounds.
77 output_limiter = (partial(max, max_output_length)
78 if max_output_length > 0 else lambda x: x)
79
80 # For each line in the standard input, parse out the JSON string we expect
81 # to see.
82 # Note the := walrus -- we're assigning and checking the condition.
83 all_osl = []
84 prompts = []
85 media_paths = []
86 all_logits = []
87 task_ids = []
88 lora_requests = []
89 while (line := stream.readline()) and len(task_ids) < max_requests:
90 # We expect the data to come in as a JSON string.
91 # For example:
92 # {"task_id": 1, "prompt": "Generate an infinite response to the following:

Callers 2

throughput_commandFunction · 0.90
latency_commandFunction · 0.90

Calls 7

LoRARequestClass · 0.90
InferenceRequestClass · 0.90
DatasetMetadataClass · 0.90
from_iterableMethod · 0.80
appendMethod · 0.45
getMethod · 0.45

Tested by

no test coverage detected