diff --git a/python/sglang/test/send_one.py b/python/sglang/test/send_one.py index dfde64f4f..46838a44a 100644 --- a/python/sglang/test/send_one.py +++ b/python/sglang/test/send_one.py @@ -5,6 +5,7 @@ Usage: python3 -m sglang.test.send_one python3 -m sglang.test.send_one --profile --profile-steps 5 python3 -m sglang.test.send_one --profile --profile-by-stage +python3 -m sglang.test.send_one --stop "<|separator|>" "<|eos|>" --max-new-tokens 2048 """ import argparse @@ -36,6 +37,7 @@ class BenchArgs: ) image: bool = False many_images: bool = False + stop: Optional[list] = None stream: bool = False profile: bool = False profile_steps: int = 3 @@ -66,6 +68,7 @@ class BenchArgs: parser.add_argument("--json", action="store_true") parser.add_argument("--return-logprob", action="store_true") parser.add_argument("--prompt", type=str, default=BenchArgs.prompt) + parser.add_argument("--stop", type=str, nargs="*", default=None) parser.add_argument("--image", action="store_true") parser.add_argument("--many-images", action="store_true") parser.add_argument("--stream", action="store_true") @@ -135,7 +138,7 @@ def send_one_prompt(args: BenchArgs): "frequency_penalty": args.frequency_penalty, "presence_penalty": args.presence_penalty, "json_schema": json_schema, - "stop": ["Question", "Assistant:", "<|separator|>", "<|eos|>"], + "stop": args.stop, }, "return_logprob": args.return_logprob, "stream": args.stream,