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

Refactor EOS check

Browse Source
This commit is contained in:
devonthomas35 2023-12-14 21:11:23 -08:00 committed by GitHub
parent d7d7aabded
commit d74d9453dd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
phi2/phi2.py
View File

@ -202,16 +202,20 @@ 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:
mx.eval(tokens) mx.eval(tokens)
eos_index = next((i for i, t in enumerate(tokens) if t.item() == tokenizer.eos_token_id), None)
if eos_index is not None:
tokens = tokens[:eos_index]
s = tokenizer.decode([t.item() for t in tokens]) s = tokenizer.decode([t.item() for t in tokens])
print(s, end="", flush=True) print(s, end="", flush=True)
tokens = [] tokens = []
if eos_index is not None:
break
mx.eval(tokens) mx.eval(tokens)
s = tokenizer.decode([t.item() for t in tokens]) s = tokenizer.decode([t.item() for t in tokens])