Update app.py

app.py

@@ -5,13 +5,7 @@ from typing import Iterator
 import gradio as gr
 import spaces
 import torch
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    StoppingCriteria,
-    StoppingCriteriaList,
-    TextIteratorStreamer,
-)
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
@@ -46,19 +40,14 @@ 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([StoppingCriteria(stops=stop_words_ids)])
-
     generate_kwargs = dict(
-        input_ids=model_inputs,
+        {"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,
+        num_beams=1,
         repetition_penalty=repetition_penalty,
     )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
@@ -69,11 +58,10 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
-stop_button = gr.Button(text="Stop")
 chat_interface = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
-        gr.Textbox(label="System prompt", lines=6),
+        gr.Textbox(label="System prompt", lines=6, width=800),  # Adjust width here
         gr.Slider(
             label="Max new tokens",
             minimum=1,
@@ -103,7 +91,7 @@ chat_interface = gr.ChatInterface(
             value=1.1,
         ),
     ],
-    stop_btn=stop_button,  # Use the created stop button instance
+    stop_btn=None,
     examples=[
         ["Hello there! How are you doing?"],
         ["Can you explain briefly to me what is the Python programming language?"],
@@ -113,7 +101,6 @@ chat_interface = gr.ChatInterface(
     ],
 )
 
-
 with gr.Blocks(css="style.css") as demo:
     chat_interface.render()