Update app.py

app.py

@@ -3,7 +3,6 @@ from threading import Thread
 from typing import Iterator
 
 import gradio as gr
-import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
@@ -16,7 +15,6 @@ if torch.cuda.is_available():
     model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
     tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-@spaces.GPU
 def generate(
     message: str,
     chat_history: list[tuple[str, str]],
@@ -41,7 +39,7 @@ def generate(
 
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
-        {"input_ids": input_ids},
+        input_ids=input_ids,
         streamer=streamer,
         max_new_tokens=max_new_tokens,
         do_sample=True,
@@ -58,51 +56,53 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
-chat_interface = gr.ChatInterface(
+chat_interface = gr.Interface(
     fn=generate,
-    additional_inputs=[
-        gr.Textbox(label="System prompt", lines=6),
+    inputs=[
+        gr.Textbox(label="User Input", lines=5, placeholder="Enter your message..."),
+        gr.Textbox(label="System Prompt", lines=5, placeholder="Enter system prompt (optional)..."),
         gr.Slider(
-            label="Max new tokens",
+            label="Max New Tokens",
             minimum=1,
             maximum=MAX_MAX_NEW_TOKENS,
             step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
+            default=DEFAULT_MAX_NEW_TOKENS,
         ),
         gr.Slider(
             label="Temperature",
             minimum=0.01,
             maximum=1.0,
             step=0.01,
-            value=0.7,
+            default=0.7,
         ),
         gr.Slider(
-            label="Top-p (nucleus sampling)",
+            label="Top-p (Nucleus Sampling)",
             minimum=0.05,
             maximum=1.0,
             step=0.01,
-            value=1.0,
+            default=1.0,
         ),
         gr.Slider(
-            label="Repetition penalty",
+            label="Repetition Penalty",
             minimum=1.0,
             maximum=2.0,
             step=0.05,
-            value=1.1,
+            default=1.1,
         ),
     ],
-    stop_btn=None,
+    outputs=gr.Textbox(label="Chat Output", lines=10),
+    title="🦣MAmmoTH2",
+    description="A simple web interactive chat demo based on gradio.",
     examples=[
         ["Hello there! How are you doing?"],
         ["Can you explain briefly to me what is the Python programming language?"],
         ["Explain the plot of Cinderella in a sentence."],
-        ["How many hours does it take a man to eat a Helicopter?"],
+        ["How many hours does it take a man to eat a helicopter?"],
         ["Write a 100-word article on 'Benefits of Open-Source in AI research'"],
     ],
+    theme="compact",
+    live=True,
+    capture_session=True,
 )
 
-with gr.Blocks(css="style.css") as demo:
-    chat_interface.render()
-
-if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+chat_interface.launch()