From 62713e2a833d9270f06c50b5f21258fb72ac5cee Mon Sep 17 00:00:00 2001 From: wangzy Date: Mon, 5 Feb 2024 16:24:03 +0800 Subject: [PATCH] change `max_tokens` to `max_new_tokens` --- agent/streaming_inference.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/agent/streaming_inference.py b/agent/streaming_inference.py index d21dd41..37ccf4f 100644 --- a/agent/streaming_inference.py +++ b/agent/streaming_inference.py @@ -98,7 +98,7 @@ def parse_args(): default=['<|action_end|>', '<|im_end|>'], action='append', help='Stop words') - parser.add_argument('--max_tokens', + parser.add_argument('--max_new_tokens', type=int, default=512, help='Number of maximum generated tokens.') @@ -557,7 +557,7 @@ def predict(args): stop_words=args.stop_words, top_p=args.top_p, top_k=args.top_k, - max_tokens=args.max_tokens, + max_new_tokens=args.max_new_tokens, ) with jsonlines.open(args.output_path, 'w') as f: for item in tqdm(