fix

app.py

@@ -13,9 +13,11 @@ from transformers import TextIteratorStreamer
 import queue
 
 class RichTextStreamer(TextIteratorStreamer):
-    def __init__(self, tokenizer, **kwargs):
+    def __init__(self, tokenizer, prompt_len=0, **kwargs):
         super().__init__(tokenizer, **kwargs)
         self.token_queue = queue.Queue()
+        self.prompt_len = prompt_len
+        self.count = 0
 
     def put(self, value):
         if isinstance(value, torch.Tensor):
@@ -26,6 +28,9 @@ class RichTextStreamer(TextIteratorStreamer):
             token_ids = [value]
 
         for token_id in token_ids:
+            self.count += 1
+            if self.count <= self.prompt_len:
+                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({
@@ -60,25 +65,29 @@ def chat_with_model(messages):
     device = torch.device("cuda")
     current_model.to(device).half()
 
-    inputs = current_tokenizer(prompt, return_tensors="pt")
-    inputs = {k: v.to(device) for k, v in inputs.items()}
+    # 1. Tokenize prompt
+    prompt = "Your input here"
+    inputs = current_tokenizer(prompt, return_tensors="pt").to(device)
+    prompt_len = inputs["input_ids"].shape[-1]
 
-    streamer = RichTextStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)
-
-    max_new_tokens = 256
-    generated_tokens = 0
-    output_text = ""
-    in_think = False
+    # 2. Init streamer with prompt_len
+    streamer = RichTextStreamer(
+        tokenizer=current_tokenizer,
+        prompt_len=prompt_len,
+        skip_special_tokens=False
+    )
 
+    # 3. Build generation kwargs
     generation_kwargs = dict(
         **inputs,
-        max_new_tokens=max_new_tokens,
+        max_new_tokens=256,
         do_sample=True,
         streamer=streamer,
         eos_token_id=eos_id,
         pad_token_id=pad_id
     )
 
+    # 4. Launch generation in a thread
     thread = threading.Thread(target=current_model.generate, kwargs=generation_kwargs)
     thread.start()