zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-10-31 10:58:07 +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
1 changed files with 7 additions and 0 deletions
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])