From e5760bc40aaace96c41b2564017c3b99b3f7bb03 Mon Sep 17 00:00:00 2001
From: Lu Changqi <58518876+zeroorhero@users.noreply.github.com>
Date: Wed, 5 Mar 2025 17:21:37 +0800
Subject: [PATCH] bench: add dataset param for bench_multiturn (#3990)

---
 benchmark/hicache/bench_multiturn.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/benchmark/hicache/bench_multiturn.py b/benchmark/hicache/bench_multiturn.py
index 1fb58e024..6bd0bd99e 100644
--- a/benchmark/hicache/bench_multiturn.py
+++ b/benchmark/hicache/bench_multiturn.py
@@ -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]