From 7da89adef8ea60817322827bd0cbdf48c62fe4ab Mon Sep 17 00:00:00 2001
From: Leon Ericsson <leon.ericsson@icloud.com>
Date: Mon, 1 Jan 2024 12:53:11 +0100
Subject: [PATCH] nist

---
 llms/lookup_decoding/decoder.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/llms/lookup_decoding/decoder.py b/llms/lookup_decoding/decoder.py
index d82e0681..afcba4bc 100644
--- a/llms/lookup_decoding/decoder.py
+++ b/llms/lookup_decoding/decoder.py
@@ -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",