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

Mixtral: Stop at EOS token (#183)

Browse Source 
* Stop at EOS token

* Precommit format files

* Fix precommit hooks

* Fix precommit hooks
This commit is contained in:
devonthomas35 2023-12-23 21:25:42 -08:00 committed by GitHub
parent 0371d90ccb
commit 939086e6a3
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
llms/mixtral/mixtral.py
View File

@ -329,9 +329,16 @@ if __name__ == "__main__":
if (len(tokens) % 10) == 0:
mx.eval(tokens)
eos_index = next(
(i for i, t in enumerate(tokens) if t.item() == tokenizer.eos_id), None
)
if eos_index is not None:
tokens = tokens[:eos_index]
s = tokenizer.decode([t.item() for t in tokens])
print(s, end="", flush=True)
tokens = []
if eos_index is not None:
break
mx.eval(tokens)
s = tokenizer.decode([t.item() for t in tokens])