Get the sample requests from the specified dataset.
(args)
| 427 | |
| 428 | |
| 429 | def get_samples(args): |
| 430 | """Get the sample requests from the specified dataset.""" |
| 431 | if not hasattr(args, "request_id_prefix"): |
| 432 | args.request_id_prefix = "" |
| 433 | |
| 434 | # For datasets that follow a similar structure, use a mapping. |
| 435 | dataset_mapping = { |
| 436 | "EB": lambda: EBDataset(random_seed=args.seed, dataset_path=args.dataset_path, shuffle=args.shuffle).sample( |
| 437 | num_requests=args.num_prompts, |
| 438 | output_len=args.sharegpt_output_len, |
| 439 | ), |
| 440 | "EBChat": lambda: EBChatDataset( |
| 441 | random_seed=args.seed, dataset_path=args.dataset_path, shuffle=args.shuffle |
| 442 | ).sample( |
| 443 | num_requests=args.num_prompts, |
| 444 | output_len=args.sharegpt_output_len, |
| 445 | ), |
| 446 | } |
| 447 | |
| 448 | try: |
| 449 | input_requests = dataset_mapping[args.dataset_name]() |
| 450 | except KeyError as err: |
| 451 | raise ValueError(f"Unknown dataset: {args.dataset_name}") from err |
| 452 | |
| 453 | return input_requests |
| 454 | |
| 455 | |
| 456 | def add_dataset_parser(parser: FlexibleArgumentParser): |
no test coverage detected