Update app.py

app.py

@@ -250,7 +250,27 @@ def stream_single(prompt, sys_prompt, temperature, max_new_tokens, device,
     except Exception as e:
         yield "", f"推理失败：{e}", state
 
+def run_single(prompt, sys_prompt, temperature, max_new_tokens, device,
+               model_path, device_map, torch_dtype, state):
+    if not prompt or not str(prompt).strip():
+        return "", "请先输入提示词。", state
 
+    t0 = time.time()
+    state = ensure_enhancer(state, model_path, device_map, torch_dtype)
+    enhancer = state["enhancer"]
+    try:
+        out = enhancer.predict(
+            prompt_cot=prompt,
+            sys_prompt=sys_prompt,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            device=device
+        )
+        dt = time.time() - t0
+        return out, f"耗时：{dt:.2f}s", state
+    except Exception as e:
+        return "", f"推理失败：{e}", state
+        
 # 示例数据
 test_list_zh = [
     "第三人称视角，赛车在城市赛道上飞驰，左上角是小地图，地图下面是当前名次，右下角仪表盘显示当前速度。",
@@ -313,12 +333,18 @@ with gr.Blocks(title="Prompt Enhancer_V2") as demo:
                 out_text = gr.Textbox(label="重写结果", lines=10)
                 out_info = gr.Markdown("准备就绪。")
 
+        # run_btn.click(
+        #     stream_single,
+        #     inputs=[prompt, sys_prompt, temperature, max_new_tokens, device,
+        #     model_path, device_map, torch_dtype, state],
+        #     outputs=[out_text, out_info, state]
+        #     )
         run_btn.click(
-            stream_single,
+            run_single,
             inputs=[prompt, sys_prompt, temperature, max_new_tokens, device,
-            model_path, device_map, torch_dtype, state],
+                    model_path, device_map, torch_dtype, state],
             outputs=[out_text, out_info, state]
-            )
+        )
 
     gr.Markdown(
         "提示：如有任何问题可email联系：linqing1995@buaa.edu.cn"