Skip to content

Commit e936b07

Browse files
committed
Set priority in the random inputs case
Signed-off-by: Jeremy Arnold <[email protected]>
1 parent 6f0ba0f commit e936b07

File tree

1 file changed

+5
-3
lines changed

1 file changed

+5
-3
lines changed

benchmarks/benchmark_prioritization.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ def sample_requests(
1717
num_requests: int,
1818
tokenizer: PreTrainedTokenizerBase,
1919
fixed_output_len: Optional[int],
20-
) -> List[Tuple[str, int, int]]:
20+
) -> List[Tuple[str, int, int, int]]:
2121
if fixed_output_len is not None and fixed_output_len < 4:
2222
raise ValueError("output_len too small")
2323

@@ -104,8 +104,10 @@ def main(args: argparse.Namespace):
104104
if args.dataset is None:
105105
# Synthesize a prompt with the given input length.
106106
prompt = "hi" * (args.input_len - 1)
107-
requests = [(prompt, args.input_len, args.output_len)
108-
for _ in range(args.num_prompts)]
107+
requests = [(prompt, args.input_len, args.output_len,
108+
# Select a equi-probable random priority
109+
0 if random.random() < 0.5 else 1
110+
)]
109111
else:
110112
requests = sample_requests(args.dataset, args.num_prompts, tokenizer,
111113
args.output_len)

0 commit comments

Comments
 (0)