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

Merge pull request #108 from devonthomas35/phi2_eos

Browse Source 
Phi-2: Stop generating at eos token
This commit is contained in:
Awni Hannun 2023-12-15 07:34:11 -08:00 committed by GitHub
commit 71b40c474c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
phi2/phi2.py
View File

@ -206,9 +206,16 @@ if __name__ == "__main__":
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])