[Misc] Add tokenizer_mode param to benchmark_serving.py (#11174)

Signed-off-by: Alexander Matveev <alexm@neuralmagic.com>
This commit is contained in:
Alexander Matveev 2024-12-13 11:19:10 -05:00 committed by GitHub
parent 5b0ed8391d
commit 238c0d93b4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -781,6 +781,7 @@ def main(args: argparse.Namespace):
backend = args.backend backend = args.backend
model_id = args.model model_id = args.model
tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model
tokenizer_mode = args.tokenizer_mode
if args.base_url is not None: if args.base_url is not None:
api_url = f"{args.base_url}{args.endpoint}" api_url = f"{args.base_url}{args.endpoint}"
@ -790,6 +791,7 @@ def main(args: argparse.Namespace):
base_url = f"http://{args.host}:{args.port}" base_url = f"http://{args.host}:{args.port}"
tokenizer = get_tokenizer(tokenizer_id, tokenizer = get_tokenizer(tokenizer_id,
tokenizer_mode=tokenizer_mode,
trust_remote_code=args.trust_remote_code) trust_remote_code=args.trust_remote_code)
if args.dataset is not None: if args.dataset is not None:
@ -1210,5 +1212,15 @@ if __name__ == "__main__":
"from the sampled HF dataset.", "from the sampled HF dataset.",
) )
parser.add_argument(
'--tokenizer-mode',
type=str,
default="auto",
choices=['auto', 'slow', 'mistral'],
help='The tokenizer mode.\n\n* "auto" will use the '
'fast tokenizer if available.\n* "slow" will '
'always use the slow tokenizer. \n* '
'"mistral" will always use the `mistral_common` tokenizer.')
args = parser.parse_args() args = parser.parse_args()
main(args) main(args)