Update app.py

app.py

@@ -1,7 +1,5 @@
-import json
-import spaces
-import subprocess
 import gradio as gr
+import spaces
 from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
 
@@ -17,25 +15,46 @@ hf_hub_download(
     local_dir="./models",
 )
 
-llm = Llama(
-    model_path="models/Qwen2.5-Math-7B-Instruct-Q6_K_L.gguf",
-    flash_attn=True,
-    n_ctx=8192,
-    n_batch=1024,
-    chat_format="chatml",
-)
+# 模型加载函数明确包装GPU环境
+@spaces.GPU
+def load_llm():
+    llm = Llama(
+        model_path="models/Qwen2.5-Math-7B-Instruct-Q6_K_L.gguf",
+        flash_attn=True,
+        n_ctx=4096,
+        n_batch=512,
+        chat_format="chatml",
+        n_gpu_layers=-1,
+    )
+    return llm
+
+# 全局加载模型（仅加载一次，提升速度）
+llm = load_llm()
 
-# Gradio 组件
+# 模型推理函数也使用GPU环境
+@spaces.GPU
+def generate_response(messages, max_tokens, temperature, top_p):
+    response = llm.create_chat_completion(
+        messages=messages,
+        stream=True,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+    )
+
+    message_repl = ""
+    for chunk in response:
+        delta = chunk['choices'][0]["delta"]
+        if "content" in delta:
+            message_repl += delta["content"]
+            yield message_repl
+
+# Gradio组件设置
 output_md = gr.Markdown(
     label="Answer",
     value="Answer will be presented here",
     latex_delimiters=[
         {"left": "\\(", "right": "\\)", "display": True},
-        {"left": "\\begin\{equation\}", "right": "\\end\{equation\}", "display": True},
-        {"left": "\\begin\{align\}", "right": "\\end\{align\}", "display": True},
-        {"left": "\\begin\{alignat\}", "right": "\\end\{alignat\}", "display": True},
-        {"left": "\\begin\{gather\}", "right": "\\end\{gather\}", "display": True},
-        {"left": "\\begin\{CD\}", "right": "\\end\{CD\}", "display": True},
         {"left": "\\[", "right": "\\]", "display": True},
     ],
     elem_id="qwen-md",
@@ -43,6 +62,7 @@ output_md = gr.Markdown(
     container=True,
     render=False,
 )
+
 target_lang = gr.Dropdown(
     choices=["Chinese", "English"],
     value="Chinese",
@@ -50,61 +70,45 @@ target_lang = gr.Dropdown(
     interactive=True,
     render=False,
 )
+
 new_tokens = gr.Slider(
-    minimum=1, maximum=8192, value=2048, step=1, label="Max new tokens", render=False
+    minimum=1, maximum=1024, value=128, step=1, label="Max new tokens", render=False
 )
+
 temperature = gr.Slider(
-    minimum=0, maximum=2.0, value=0.5, step=0.1, label="Temperature", render=False
+    minimum=0, maximum=1.0, value=0.2, step=0.1, label="Temperature", render=False
 )
-top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.05, label="Top P", render=False)
+
+top_p = gr.Slider(
+    minimum=0.0, maximum=1.0, value=0.8, step=0.05, label="Top P", render=False
+)
+
 input_text = gr.Textbox(label="Ask math questions here", render=False)
 submit_btn = gr.Button(value="Ask", render=False)
+
 banner = gr.Markdown(value="""
 # 📖 Qwen2.5-Math GGUF
-This WebUI is based on Qwen2.5-Math-7B-Instruct-GGUF for mathematical reasoning. You can input texts of mathematical or arithmetic problems.
-"""
-)
-
+Fast mathematical reasoning with ZeroGPU optimized.
+""")
 
-# Gradio 函数
-def respond(
-    input_text,
-    lang="Chinese",
-    max_tokens=2048,
-    temperature=0.5,
-    top_p=0.95,
-):
-    if lang == "Chinese":
-        sys_msg = "你是一个乐于助人的数学助手. 你使用中文回答问题"
-    else:
-        sys_msg = "You are a helpful math assistant. You should always provide your answer in English."
+# Gradio响应函数，调用GPU推理
+def respond(input_text, lang, max_tokens, temperature, top_p):
+    sys_msg = (
+        "你是一个乐于助人的数学助手. 你使用中文回答问题"
+        if lang == "Chinese" else
+        "You are a helpful math assistant. You should always provide your answer in English."
+    )
     messages = [
-        {
-            "role": "system",
-            "content": sys_msg,
-        },
+        {"role": "system", "content": sys_msg},
         {"role": "user", "content": input_text},
     ]
 
-    response = ""
-    response = llm.create_chat_completion(
-        messages=messages,
-        stream=True,
-        max_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-    )
-    message_repl = ""
-    for chunk in response:
-        if len(chunk['choices'][0]["delta"]) != 0 and "content" in chunk['choices'][0]["delta"]:
-            message_repl = message_repl + \
-                chunk['choices'][0]["delta"]["content"]
-        yield message_repl
-
+    # GPU生成器直接返回
+    yield from generate_response(messages, max_tokens, temperature, top_p)
 
 with gr.Blocks(css=CSS, theme="NoCrypt/miku") as demo:
     submit_btn.click(
-        fn=respond,
+        respond,
         inputs=[input_text, target_lang, new_tokens, temperature, top_p],
         outputs=output_md,
     )
@@ -122,4 +126,4 @@ with gr.Blocks(css=CSS, theme="NoCrypt/miku") as demo:
                 output_md.render()
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()