From 7110f781f7f29420ca582a85400baef28f212711 Mon Sep 17 00:00:00 2001 From: lijiahao Date: Sat, 25 Nov 2023 00:51:26 +0800 Subject: [PATCH] Support in cli_demo --- examples/cli_demo.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/examples/cli_demo.py b/examples/cli_demo.py index 1a07efa..c134b7d 100644 --- a/examples/cli_demo.py +++ b/examples/cli_demo.py @@ -41,6 +41,12 @@ def main() -> None: parser.add_argument( "-l", "--max_length", default=2048, type=int, help="max total length including prompt and output" ) + parser.add_argument( + "--max_new_tokens", + default=-1, + type=int, + help="max number of tokens to generate, ignoring the number of prompt tokens", + ) parser.add_argument("-c", "--max_context_length", default=512, type=int, help="max context length") parser.add_argument("--top_k", default=0, type=int, help="top-k sampling") parser.add_argument("--top_p", default=0.7, type=float, help="top-p sampling") @@ -65,6 +71,7 @@ def main() -> None: generation_kwargs = dict( max_length=args.max_length, + max_new_tokens=args.max_new_tokens, max_context_length=args.max_context_length, do_sample=args.temp > 0, top_k=args.top_k,