Fix stopping criteria

app.py

@@ -6,6 +6,16 @@ import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import threading
 
+from transformers import TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
+import threading
+
+class StopOnEos(StoppingCriteria):
+    def __init__(self, eos_token_id):
+        self.eos_token_id = eos_token_id
+
+    def __call__(self, input_ids, scores, **kwargs):
+        return input_ids[0, -1].item() == self.eos_token_id
+
 @spaces.GPU
 def chat_with_model(messages):
     global current_model, current_tokenizer
@@ -18,15 +28,17 @@ def chat_with_model(messages):
     prompt = format_prompt(messages)
     inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
 
-    streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)
+    streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=True)
+    stopping_criteria = StoppingCriteriaList([StopOnEos(current_tokenizer.eos_token_id)])
+
     generation_kwargs = dict(
         **inputs,
         max_new_tokens=256,
         do_sample=True,
-        streamer=streamer
+        streamer=streamer,
+        stopping_criteria=stopping_criteria
     )
 
-    # Launch generation in a background thread
     thread = threading.Thread(target=current_model.generate, kwargs=generation_kwargs)
     thread.start()
 
@@ -37,15 +49,13 @@ def chat_with_model(messages):
     for new_text in streamer:
         output_text += new_text
         messages[-1]["content"] = output_text
-        if current_tokenizer.eos_token and current_tokenizer.eos_token in output_text:
-            break
         yield messages
-        
 
     current_model.to("cpu")
     torch.cuda.empty_cache()
 
 
+
 # Globals
 current_model = None
 current_tokenizer = None