Update app.py

app.py

@@ -5,7 +5,25 @@ from typing import Iterator
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    StoppingCriteria,
+    StoppingCriteriaList,
+    TextIteratorStreamer,
+)
+
+class StoppingCriteriaSub(StoppingCriteria):
+    def __init__(self, stops = [], encounters=1):
+        super().__init__()
+        # self.stops = [stop.to("cuda") for stop in stops]
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
+        last_token = input_ids[0][-1]
+        for stop in self.stops:
+            if tokenizer.decode(stop) == tokenizer.decode(last_token):
+                return True
+        return False
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
@@ -17,6 +35,7 @@ if torch.cuda.is_available():
     tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 @spaces.GPU
+User
 def generate(
     message: str,
     chat_history: list[tuple[str, str]],
@@ -40,6 +59,10 @@ def generate(
     input_ids = input_ids.to(model.device)
 
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+    stop_words = ["</s>"]
+    stop_words_ids = [tokenizer(stop_word, return_tensors='pt', add_special_tokens=False)['input_ids'].squeeze() for stop_word in stop_words]
+    stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])
+    
     generate_kwargs = dict(
         {"input_ids": input_ids},
         streamer=streamer,
@@ -48,6 +71,7 @@ def generate(
         top_p=top_p,
         temperature=temperature,
         num_beams=1,
+        stopping_criteria=stopping_criteria,
         repetition_penalty=repetition_penalty,
     )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
@@ -56,11 +80,7 @@ def generate(
     outputs = []
     for text in streamer:
         outputs.append(text)
-        generated_text = "".join(outputs)
-        if "<s>" in generated_text:
-            yield generated_text[:generated_text.index("<s>")+3]
-            break
-        yield generated_text
+        yield "".join(outputs)
 
 
 chat_interface = gr.ChatInterface(