Fix prompt showing (now skipping)

app.py

@@ -16,9 +16,9 @@ class RichTextStreamer(TextIteratorStreamer):
     def __init__(self, tokenizer, **kwargs):
         super().__init__(tokenizer, **kwargs)
         self.token_queue = queue.Queue()
+        self.prompt_shown = not self.skip_prompt 
 
     def put(self, value):
-        # Convert incoming tensor or list to flat list of token IDs
         if isinstance(value, torch.Tensor):
             token_ids = value.view(-1).tolist()
         elif isinstance(value, list):
@@ -27,6 +27,9 @@ class RichTextStreamer(TextIteratorStreamer):
             token_ids = [value]
 
         for token_id in token_ids:
+            if self.skip_prompt and not self.prompt_shown:
+                continue  # skip prompt tokens
+
             token_str = self.tokenizer.decode([token_id], **self.decode_kwargs)
             is_special = token_id in self.tokenizer.all_special_ids
             self.token_queue.put({
@@ -34,6 +37,7 @@ class RichTextStreamer(TextIteratorStreamer):
                 "token": token_str,
                 "is_special": is_special
             })
+            self.prompt_shown = True
 
     def __iter__(self):
         while True: