| 191 | |
| 192 | |
| 193 | def get_requests(args, tokenizer): |
| 194 | # Common parameters for all dataset types. |
| 195 | common_kwargs = { |
| 196 | "dataset_path": args.dataset_path, |
| 197 | "random_seed": args.seed, |
| 198 | } |
| 199 | sample_kwargs = { |
| 200 | # "tokenizer": tokenizer, |
| 201 | "lora_path": args.lora_path, |
| 202 | # "max_loras": args.max_loras, |
| 203 | "num_requests": args.num_prompts, |
| 204 | "input_len": args.input_len, |
| 205 | "output_len": args.output_len, |
| 206 | } |
| 207 | if args.dataset_path is None or args.dataset_name == "random": |
| 208 | sample_kwargs["range_ratio"] = args.random_range_ratio |
| 209 | sample_kwargs["prefix_len"] = args.prefix_len |
| 210 | sample_kwargs["tokenizer"] = tokenizer |
| 211 | dataset_cls = RandomDataset |
| 212 | elif args.dataset_name == "EB": |
| 213 | dataset_cls = EBDataset |
| 214 | elif args.dataset_name == "EBChat": |
| 215 | dataset_cls = EBChatDataset |
| 216 | else: |
| 217 | raise ValueError(f"Unknown dataset name: {args.dataset_name}") |
| 218 | # Remove None values |
| 219 | sample_kwargs = {k: v for k, v in sample_kwargs.items() if v is not None} |
| 220 | return dataset_cls(**common_kwargs).sample(**sample_kwargs) |
| 221 | |
| 222 | |
| 223 | def validate_args(args): |