Unverified Commit e5760bc4 authored by Lu Changqi's avatar Lu Changqi Committed by GitHub
Browse files

bench: add dataset param for bench_multiturn (#3990)

parent 56a724eb
......@@ -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]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment