zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-29 18:26:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Stop generating at eos token

Browse Source
This commit is contained in:
devonthomas35 2023-12-14 15:50:59 -08:00 committed by GitHub
parent 95a1d50318
commit 8d496ba61a
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
phi2/phi2.py
View File

@ -202,7 +202,11 @@ 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)):
tokens.append(token)
if token == tokenizer.eos_token_id:
break
else:
tokens.append(token)
if (len(tokens) % 10) == 0: if (len(tokens) % 10) == 0:
mx.eval(tokens) mx.eval(tokens)