Fix bugs

app.py

@@ -12,7 +12,6 @@ import threading
 from transformers import TextIteratorStreamer
 import queue
 
-@spaces.GPU
 class RichTextStreamer(TextIteratorStreamer):
     def __init__(self, tokenizer, **kwargs):
         super().__init__(tokenizer, **kwargs)
@@ -54,6 +53,7 @@ def chat_with_model(messages):
         return
 
     pad_id = current_tokenizer.pad_token_id
+    eos_id = current_tokenizer.eos_token_id
     if pad_id is None:
         pad_id = current_tokenizer.unk_token_id or 0
 
@@ -66,35 +66,40 @@ def chat_with_model(messages):
 
     streamer = RichTextStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)
 
+    max_new_tokens = 256
+    generated_tokens = 0
+    output_text = ""
+    in_think = False
+
     generation_kwargs = dict(
         **inputs,
-        max_new_tokens=256,
+        max_new_tokens=max_new_tokens,
         do_sample=True,
         streamer=streamer,
-        eos_token_id=current_tokenizer.eos_token_id,
+        eos_token_id=eos_id,
         pad_token_id=pad_id
     )
 
     thread = threading.Thread(target=current_model.generate, kwargs=generation_kwargs)
     thread.start()
 
-    output_text = ""
     messages = messages.copy()
     messages.append({"role": "assistant", "content": ""})
-    in_think = False
 
     for token_info in streamer:
         token_str = token_info["token"]
         token_id = token_info["token_id"]
         is_special = token_info["is_special"]
 
-        if token_id == current_tokenizer.eos_token_id:
-            streamer.end_of_generation.set()  # signal to stop generation thread
+        # Stop immediately at EOS
+        if token_id == eos_id:
             break
 
+        # Optional: skip other special tokens
         if is_special:
             continue
 
+        # Detect reasoning block
         if "<think>" in token_str:
             in_think = True
             token_str = token_str.replace("<think>", "")
@@ -107,8 +112,14 @@ def chat_with_model(messages):
         else:
             output_text += token_str
 
+        # Early stopping if user reappears
         if "\nUser:" in output_text:
             output_text = output_text.split("\nUser:")[0].rstrip()
+            break
+
+        generated_tokens += 1
+        if generated_tokens >= max_new_tokens:
+            break
 
         messages[-1]["content"] = output_text
         yield messages
@@ -118,13 +129,11 @@ def chat_with_model(messages):
         messages[-1]["content"] = output_text
         yield messages
 
-    # Ensure generation thread stops
+    # Wait for thread to finish
     thread.join(timeout=1.0)
-
     current_model.to("cpu")
     torch.cuda.empty_cache()
 
-    return messages
 
 
 # Globals