fix: max_tokens handling in generate.py (#70)

Author: Mirza-Samad-Ahmed-Baig

gpt_oss/generate.py

@@ -28,7 +28,8 @@ def main(args):
 
     tokenizer = get_tokenizer()
     tokens = tokenizer.encode(args.prompt)
-    for token, logprob in generator.generate(tokens, stop_tokens=[tokenizer.eot_token], temperature=args.temperature, max_tokens=args.limit, return_logprobs=True):
+    max_tokens = None if args.limit == 0 else args.limit
+    for token, logprob in generator.generate(tokens, stop_tokens=[tokenizer.eot_token], temperature=args.temperature, max_tokens=max_tokens, return_logprobs=True):
         tokens.append(token)
         decoded_token = tokenizer.decode([token])
         print(