@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 207 lines 5.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Convert a plain-text file (local path or URL) into a JSONL dataset
4 compatible with ``CustomDataset`` (``--dataset-name custom``), by
5 randomly slicing the tokenized text into prompts.
6
7 Each line of the output JSONL contains a ``prompt`` (decoded from a random
8 slice of the tokenized source text) and an ``output_tokens`` count.
9
10 Usage
11 -----
12 ::
13
14 python -m vllm.benchmarks.datasets.create_txt_slices_dataset \\
15 --input sonnet.txt \\
16 --output sonnet_dataset.jsonl \\
17 --tokenizer gpt2 \\
18 --num-prompts 1000 \\
19 --input-len 1024 \\
20 --output-len 128
21
22 The resulting JSONL file can then be used with the serving benchmark::
23
24 python -m vllm.benchmarks.serve \\
25 --dataset-name custom \\
26 --dataset-path sonnet_dataset.jsonl \\
27 ...
28 """
29
30 from __future__ import annotations
31
32 import json
33 import logging
34 import random
35 import urllib.request
36
37 import numpy as np
38 from transformers import AutoTokenizer
39
40 from vllm.benchmarks.datasets.utils import RangeRatio, get_sampling_params
41 from vllm.utils.argparse_utils import FlexibleArgumentParser
42
43 logger = logging.getLogger(__name__)
44
45
46 def load_text(path: str) -> str:
47 """Load text from a local file or URL."""
48 if path.startswith(("http://", "https://")):
49 with urllib.request.urlopen(path) as response:
50 return response.read().decode("utf-8")
51 with open(path, encoding="utf-8") as f:
52 return f.read()
53
54
55 def create_txt_slices_jsonl(
56 *,
57 input_path: str,
58 output_path: str,
59 tokenizer_name: str,
60 num_prompts: int,
61 input_len: int,
62 output_len: int,
63 range_ratio: RangeRatio = 0.0,
64 seed: int = 0,
65 trust_remote_code: bool = False,
66 ) -> None:
67 """Read *input_path*, slice it into prompts, and write JSONL to
68 *output_path*."""
69 tokenizer = AutoTokenizer.from_pretrained(
70 tokenizer_name, trust_remote_code=trust_remote_code
71 )
72
73 text = load_text(input_path)
74 if not text:
75 raise ValueError("The text file is empty and cannot be sampled from.")
76
77 token_ids = tokenizer(text, add_special_tokens=False).input_ids
78 if not token_ids:
79 raise ValueError("Tokenizing the text produced zero tokens; cannot sample.")
80
81 rng_np = np.random.default_rng(seed)
82 rng_py = random.Random(seed)
83
84 input_lens, output_lens, _ = get_sampling_params(
85 rng_np,
86 num_prompts,
87 range_ratio,
88 input_len,
89 output_len,
90 tokenizer,
91 )
92
93 num_available_tokens = len(token_ids)
94
95 records: list[dict[str, object]] = []
96 for i in range(num_prompts):
97 req_input_len = int(input_lens[i])
98 req_output_len = int(output_lens[i])
99
100 # Randomly select a start position and slice with cycling
101 start_pos = rng_py.randint(0, num_available_tokens - 1)
102 prompt_token_ids = [
103 token_ids[(start_pos + j) % num_available_tokens]
104 for j in range(req_input_len)
105 ]
106 prompt = tokenizer.decode(prompt_token_ids, skip_special_tokens=False)
107
108 records.append({"prompt": prompt, "output_tokens": req_output_len})
109
110 with open(output_path, "w", encoding="utf-8") as f:
111 for record in records:
112 f.write(json.dumps(record, ensure_ascii=False) + "\n")
113
114 logger.info(
115 "Wrote %d prompts to %s",
116 len(records),
117 output_path,
118 )
119
120
121 def main(argv: list[str] | None = None) -> None:
122 parser = FlexibleArgumentParser(
123 description="Convert a plain-text file into a JSONL dataset "
124 "for CustomDataset (--dataset-name custom).",
125 )
126 parser.add_argument(
127 "--input",
128 required=True,
129 help="Path or URL to the source text file.",
130 )
131 parser.add_argument(
132 "--output",
133 required=True,
134 help="Path for the output JSONL file.",
135 )
136 parser.add_argument(
137 "--tokenizer",
138 required=True,
139 help="HuggingFace tokenizer name or path.",
140 )
141 parser.add_argument(
142 "--num-prompts",
143 type=int,
144 default=1000,
145 help="Number of prompt samples to generate (default: 1000).",
146 )
147 parser.add_argument(
148 "--input-len",
149 type=int,
150 default=1024,
151 help="Target number of input tokens per prompt (default: 1024).",
152 )
153 parser.add_argument(
154 "--output-len",
155 type=int,
156 default=128,
157 help="Target number of output tokens per prompt (default: 128).",
158 )
159 parser.add_argument(
160 "--range-ratio",
161 type=str,
162 default="0.0",
163 help="Range ratio for input/output length sampling (default: 0.0). "
164 "A single float applies to both ISL and OSL. "
165 'A JSON dict like \'{"input": 0.3, "output": 0.5}\' sets them '
166 "independently. Values must be in [0, 1).",
167 )
168 parser.add_argument(
169 "--seed",
170 type=int,
171 default=0,
172 help="Random seed for reproducibility (default: 0).",
173 )
174 parser.add_argument(
175 "--trust-remote-code",
176 action="store_true",
177 help="Trust remote code from HuggingFace.",
178 )
179
180 args = parser.parse_args(argv)
181
182 logging.basicConfig(level=logging.INFO)
183
184 # Parse --range-ratio: try float first, then JSON dict.
185 range_ratio: RangeRatio
186 try:
187 range_ratio = float(args.range_ratio)
188 except ValueError:
189 import json as _json
190
191 range_ratio = _json.loads(args.range_ratio)
192
193 create_txt_slices_jsonl(
194 input_path=args.input,
195 output_path=args.output,
196 tokenizer_name=args.tokenizer,
197 num_prompts=args.num_prompts,
198 input_len=args.input_len,
199 output_len=args.output_len,
200 range_ratio=range_ratio,
201 seed=args.seed,
202 trust_remote_code=args.trust_remote_code,
203 )
204
205
206 if __name__ == "__main__":
207 main()