fix please

app.py

@@ -16,7 +16,6 @@ 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):
         if isinstance(value, torch.Tensor):
@@ -27,9 +26,6 @@ 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({
@@ -37,7 +33,6 @@ class RichTextStreamer(TextIteratorStreamer):
                 "token": token_str,
                 "is_special": is_special
             })
-            self.prompt_shown = True
 
     def __iter__(self):
         while True:
@@ -92,6 +87,8 @@ def chat_with_model(messages):
 
     print(f'Step 1: {messages}')
 
+    prompt_text = current_tokenizer.decode(inputs["input_ids"][0], skip_special_tokens=True)
+
     for token_info in streamer:
         token_str = token_info["token"]
         token_id = token_info["token_id"]
@@ -119,6 +116,12 @@ def chat_with_model(messages):
             output_text = output_text.split("\nUser:")[0].rstrip()
             break
 
+        # Strip prompt from start of generated output
+        if output_text.startswith(prompt_text):
+            stripped_output = output_text[len(prompt_text):]
+        else:
+            stripped_output = output_text
+
         generated_tokens += 1
         if generated_tokens >= max_new_tokens:
             break