Stop generating at eos token

This commit is contained in:
devonthomas35 2023-12-14 15:50:59 -08:00 committed by GitHub
parent 95a1d50318
commit 8d496ba61a

View File

@ -202,6 +202,10 @@ if __name__ == "__main__":
tokens = [] tokens = []
for token, _ in zip(generate(prompt, model), range(args.max_tokens)): for token, _ in zip(generate(prompt, model), range(args.max_tokens)):
if token == tokenizer.eos_token_id:
break
else:
tokens.append(token) tokens.append(token)
if (len(tokens) % 10) == 0: if (len(tokens) % 10) == 0: