Update app.py

app.py

@@ -5,7 +5,7 @@ from typing import Iterator
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteriaList, StoppingCriteriaSub
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
@@ -25,7 +25,6 @@ def generate(
     temperature: float = 0.7,
     top_p: float = 1.0,
     repetition_penalty: float = 1.1,
-    stop_ids: list[list[int]] = [[2, 6, 7, 8]]
 ) -> Iterator[str]:
     conversation = []
     if system_prompt:
@@ -40,18 +39,21 @@ def generate(
         gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
     input_ids = input_ids.to(model.device)
 
+    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)])
+
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        {"input_ids": input_ids},
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        temperature=temperature,
-        num_beams=1,
-        repetition_penalty=repetition_penalty,
-        stop_ids=stop_ids
-    )
+    generate_kwargs = {
+        "input_ids": input_ids,
+        "streamer": streamer,
+        "max_new_tokens": max_new_tokens,
+        "do_sample": True,
+        "top_p": top_p,
+        "temperature": temperature,
+        "stopping_criteria": stopping_criteria,
+        "repetition_penalty": repetition_penalty,
+    }
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
@@ -63,7 +65,7 @@ def generate(
 chat_interface = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
-        gr.Textbox(label="System prompt", lines=8),  # Increase lines for larger box
+        gr.Textbox(label="System prompt", lines=6),
         gr.Slider(
             label="Max new tokens",
             minimum=1,
@@ -93,7 +95,7 @@ chat_interface = gr.ChatInterface(
             value=1.1,
         ),
     ],
-    stop_token=None,  # Remove stop token
+    stop_words=stop_words,  # Set the stop words
     examples=[
         ["Hello there! How are you doing?"],
         ["Can you explain briefly to me what is the Python programming language?"],
@@ -107,4 +109,4 @@ with gr.Blocks(css="style.css") as demo:
     chat_interface.render()
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()