Update app.py

app.py

@@ -16,8 +16,7 @@ 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(
+def generate_and_display(
     message: str,
     chat_history: list[tuple[str, str]],
     system_prompt: str,
@@ -25,7 +24,7 @@ def generate(
     temperature: float = 0.7,
     top_p: float = 1.0,
     repetition_penalty: float = 1.1,
-) -> Iterator[str]:
+) -> str:
     conversation = []
     if system_prompt:
         conversation.append({"role": "system", "content": system_prompt})
@@ -36,7 +35,7 @@ def generate(
     input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
+        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)
@@ -52,48 +51,60 @@ def generate(
     )
 
     outputs = []
-    with torch.no_grad():
-        model_outputs = model.generate(**generate_kwargs)
-        for text in streamer.generate_from_iterator(model_outputs):
-            outputs.append(text)
-            yield "".join(outputs)
+    model_outputs = model.generate(**generate_kwargs)
+    for text in streamer.generate_from_iterator(model_outputs):
+        outputs.append(text)
+    return "".join(outputs)
 
-chat_interface = gr.Interface(
-    fn=generate,
-    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",
-            minimum=1,
-            maximum=MAX_MAX_NEW_TOKENS,
-            step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
-        ),
-        gr.Slider(
-            label="Temperature",
-            minimum=0.01,
-            maximum=1.0,
-            step=0.01,
-            value=0.7,
-        ),
-        gr.Slider(
-            label="Top-p (Nucleus Sampling)",
-            minimum=0.05,
-            maximum=1.0,
-            step=0.01,
-            value=1.0,
-        ),
-        gr.Slider(
-            label="Repetition Penalty",
-            minimum=1.0,
-            maximum=2.0,
-            step=0.05,
-            value=1.1,
-        ),
-        "generate"  # This is a placeholder for the button
-    ],
-    outputs=gr.Textbox(label="Chat Output", lines=10),
+def generate_response():
+    outputs = generate_and_display(
+        input_textbox.value,
+        chat_history=[],
+        system_prompt=system_prompt_textbox.value,
+        max_new_tokens=max_new_tokens_slider.value,
+        temperature=temperature_slider.value,
+        top_p=top_p_slider.value,
+        repetition_penalty=repetition_penalty_slider.value,
+    )
+    chat_output_textbox.value = outputs
+
+input_textbox = gr.Textbox(label="User Input", lines=5, placeholder="Enter your message...")
+system_prompt_textbox = gr.Textbox(label="System Prompt", lines=5, placeholder="Enter system prompt (optional)...")
+max_new_tokens_slider = gr.Slider(
+    label="Max New Tokens",
+    minimum=1,
+    maximum=MAX_MAX_NEW_TOKENS,
+    step=1,
+    value=DEFAULT_MAX_NEW_TOKENS,
+)
+temperature_slider = gr.Slider(
+    label="Temperature",
+    minimum=0.01,
+    maximum=1.0,
+    step=0.01,
+    value=0.7,
+)
+top_p_slider = gr.Slider(
+    label="Top-p (Nucleus Sampling)",
+    minimum=0.05,
+    maximum=1.0,
+    step=0.01,
+    value=1.0,
+)
+repetition_penalty_slider = gr.Slider(
+    label="Repetition Penalty",
+    minimum=1.0,
+    maximum=2.0,
+    step=0.05,
+    value=1.1,
+)
+generate_button = gr.Button(label="Generate Response", command=generate_response)
+chat_output_textbox = gr.Textbox(label="Chat Output", lines=10)
+
+gr.Interface(
+    generate_and_display,
+    inputs=[input_textbox, system_prompt_textbox, max_new_tokens_slider, temperature_slider, top_p_slider, repetition_penalty_slider],
+    outputs=chat_output_textbox,
     title="🦣MAmmoTH2",
     description="A simple web interactive chat demo based on gradio.",
     examples=[
@@ -105,6 +116,4 @@ chat_interface = gr.Interface(
     ],
     theme="default",
     live=True,
-)
-
-chat_interface.launch()
+).launch()