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

nist

Browse Source
This commit is contained in:
Leon Ericsson 2024-01-01 12:53:11 +01:00
parent 3f750759d3
commit 7da89adef8
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
llms/lookup_decoding/decoder.py
View File

@ -32,7 +32,7 @@ class PromptLookupDecoder:
x = sample(logits)
yield x
# Generate without prompt lookup decoding (for testing)
# Normal decoding w/o prompt lookup (for testing)
def generate(
self,
prompt,
@ -140,10 +140,11 @@ class PromptLookupDecoder:
n_decoding_steps += 1
# Check stop decodig criteria:
# Check stop decodig criteria and print accepted draft tokens.
for t in new_tokens.tolist()[:-1]:
if t == self.tokenizer.eos_id:
break
if color:
print(
"\033[34m" + self.tokenizer.decode([t]) + "\033[30m",