bench: add dataset param for bench_multiturn (#3990)

This commit is contained in:
Lu Changqi
2025-03-05 17:21:37 +08:00
committed by GitHub
parent 56a724eba3
commit e5760bc40a

View File

@@ -85,6 +85,12 @@ def parse_args():
default="meta-llama/Llama-3.1-8B-Instruct",
help="model path compatible with Hugging Face Transformers",
)
parser.add_argument(
"--dataset-path",
type=str,
default="",
help="local dataset to sample tokens from",
)
parser.add_argument(
"--log-file",
type=str,
@@ -231,7 +237,7 @@ class WorkloadGenerator:
num_prompts=args.num_clients * args.num_rounds,
range_ratio=1.0,
tokenizer=self.tokenizer,
dataset_path="",
dataset_path=args.dataset_path,
)
self.candidate_inputs = [i[0] for i in self.candidate_inputs]