Update app.py

app.py

@@ -5,25 +5,8 @@ from typing import Iterator
 import gradio as gr
 import spaces
 import torch
-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
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers.generation_stopping_criteria import StoppingCriteria, StoppingCriteriaList
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
@@ -57,21 +40,22 @@ 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)
 
+    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)])
+    stopping_criteria = StoppingCriteriaList([StoppingCriteria(stops=stop_words_ids)])
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-    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,
-    }
+    generate_kwargs = dict(
+        input_ids=model_inputs,
+        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()
 
@@ -113,7 +97,7 @@ chat_interface = gr.ChatInterface(
             value=1.1,
         ),
     ],
-    stop_words=stop_words,  # Set the stop words
+    stop_button=True,  # Changed stop button to True
     examples=[
         ["Hello there! How are you doing?"],
         ["Can you explain briefly to me what is the Python programming language?"],