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