Update app.py

app.py

@@ -1,82 +1,147 @@
 import gradio as gr
 import requests
 import json
+import pandas as pd
+import io
 
 # ----------------- 配置你的服务器API地址 -----------------
 # 将IP地址和端口替换为你自己的
 SERVER_URL = "http://103.235.229.133:7000/predict"
 # -------------------------------------------------------
 
-# 默认的输入数据，方便用户测试
-DEFAULT_CSV_DATA = "id,subject,relation,object,prompt_source\n001,France,capital city of,Paris,she traveled to France for the first time\n002,Germany,capital city of,Berlin,the capital of Germany is Berlin"
+# 默认的模型名称和输入数据，方便用户测试
+DEFAULT_MODEL_NAME = "/data/tsq/MODELS/Qwen2.5-1.5B-Instruct"
+DEFAULT_DATAFRAME_VALUE = [
+    ["France", "capital city of", "Paris", "she traveled to France for the first time"],
+    ["Germany", "capital city of", "Berlin", "the capital of Germany is Berlin"]
+]
 
-def call_api(csv_data):
+def call_api(model_name, only_final_result, data_df):
     """
-    这个函数会被Gradio调用，它负责向后端API发送请求。
+    这个函数会被Gradio调用，它负责收集UI上的输入，
+    将其转换成API需要的格式，然后发送请求。
     """
-    if not csv_data.strip():
-        return "错误：输入数据不能为空！", None
+    # 1. 数据校验
+    if not model_name.strip():
+        # 使用Gradio的错误提示功能
+        raise gr.Error("模型名称不能为空！")
+    if data_df is None or data_df.empty:
+        raise gr.Error("输入数据不能为空！请在表格中至少添加一行数据。")
 
-    # 构建发送到API的JSON数据体
+    # 2. 将输入的DataFrame转换成API需要的CSV字符串
+    # API需要 'id' 列，我们在这里自动生成
+    data_df_with_id = data_df.copy()
+    # 为id列生成唯一的、格式化的ID，例如 000, 001, 002...
+    data_df_with_id.insert(0, 'id', [f'{i:03d}' for i in range(len(data_df_with_id))])
+    
+    # 将DataFrame转换为CSV字符串
+    # 使用io.StringIO来在内存中处理字符串，避免写入文件
+    output = io.StringIO()
+    data_df_with_id.to_csv(output, index=False)
+    csv_data = output.getvalue()
+    
+    # 3. 构建发送到API的JSON数据体
     payload = {
-        "model_name": "/data/tsq/MODELS/Qwen2.5-1.5B-Instruct",
+        "model_name": model_name,
         "data_csv": csv_data,
-        "only_final_result": False # 你可以根据需要修改这个参数
+        "only_final_result": only_final_result
     }
 
     try:
-        # 发送POST请求
-        response = requests.post(SERVER_URL, json=payload, timeout=60) # 设置60秒超时
+        # 4. 发送POST请求
+        # 添加流式传输和更长的超时时间以应对长时间运行的任务
+        response = requests.post(SERVER_URL, json=payload, timeout=300, stream=True)
+        response.raise_for_status() # 如果状态码不是2xx，则抛出HTTPError
 
-        # 检查响应状态码
-        if response.status_code == 200:
-            # 解析返回的JSON数据
-            result_json = response.json()
-            # 为了美观地显示，我们将JSON格式化
-            pretty_json = json.dumps(result_json, indent=2, ensure_ascii=False)
-            
-            # 假设结果在 'result' 键下，并且是一个列表
-            # 我们可以直接将其传递给Gradio的DataFrame组件
-            dataframe_result = result_json.get("result", [])
-            
-            return pretty_json, dataframe_result
+        # 解析返回的JSON数据
+        result_json = response.json()
+        
+        # 5. 处理和返回结果
+        # 从结果中提取'result'键，它应该是一个字典列表
+        dataframe_result_data = result_json.get("result", [])
+        
+        # 将字典列表转换为Pandas DataFrame，以便Gradio的DataFrame组件更好地显示
+        if isinstance(dataframe_result_data, list) and len(dataframe_result_data) > 0:
+            output_df = pd.DataFrame(dataframe_result_data)
         else:
-            # 如果服务器返回错误
-            error_message = f"API请求失败，状态码: {response.status_code}\n错误详情: {response.text}"
-            return error_message, None
+            output_df = pd.DataFrame() # 如果没有结果，返回一个空的DataFrame
+            
+        return result_json, output_df
 
+    except requests.exceptions.HTTPError as e:
+        # 更详细地捕获HTTP错误
+        error_details = f"API返回错误 (状态码: {e.response.status_code}):\n{e.response.text}"
+        raise gr.Error(f"请求失败: {error_details}")
     except requests.exceptions.RequestException as e:
-        # 如果网络连接等问题出错
-        error_message = f"连接服务器失败: {e}"
-        return error_message, None
+        # 捕获网络连接等问题
+        raise gr.Error(f"连接服务器失败: {e}")
+    except Exception as e:
+        # 捕获其他未知错误
+        raise gr.Error(f"发生未知错误: {e}")
+
 
-# 使用Gradio Blocks创建更灵活的界面
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 Ml_patch 函数推理接口")
-    gr.Markdown("在下方文本框中输入CSV格式的数据，然后点击“提交”按钮，调用云服务器进行推理。")
+# --- 使用Gradio Blocks创建更灵活、更美观的界面 ---
+with gr.Blocks(theme=gr.themes.Glass(), css=".gradio-container {max-width: 1280px !important; margin: auto;}") as demo:
+    gr.Markdown(
+        """
+        # 🚀 Ml_patch 函数推理接口
+        一个为 `Ml_patch` 函数设计的交互式Web界面。请在下方配置模型参数并输入您要测试的数据。
+        """
+    )
     
     with gr.Row():
-        with gr.Column(scale=1):
-            input_csv = gr.Textbox(
-                lines=10, 
-                label="输入数据 (CSV格式)",
-                info="第一行必须是表头，例如: id,subject,relation,object,prompt_source",
-                value=DEFAULT_CSV_DATA
+        # 左侧列：输入配置
+        with gr.Column(scale=2):
+            gr.Markdown("### 1. 配置模型和参数")
+            with gr.Group():
+                model_name_input = gr.Textbox(
+                    label="模型名称 (model_name)",
+                    value=DEFAULT_MODEL_NAME,
+                    info="请输入部署在服务器上的模型路径。"
+                )
+                only_final_result_input = gr.Checkbox(
+                    label="仅返回最终结果 (only_final_result)",
+                    value=False,
+                    info="勾选此项，API将只返回最终聚合结果，而不是所有层的详细信息。"
+                )
+            
+            gr.Markdown("### 2. 输入数据")
+            # 使用可编辑的DataFrame作为输入，用户可以动态增删行
+            data_input = gr.DataFrame(
+                label="数据表格",
+                headers=["subject", "relation", "object", "prompt_source"],
+                value=DEFAULT_DATAFRAME_VALUE,
+                row_count=(1, "dynamic"), # 至少1行，动态增加
+                col_count=(4, "fixed"),
+                interactive=True,
+                info="点击表格下方的 'Edit' (铅笔图标)，然后点击 '+' 来添加新行，或点击垃圾桶图标删除行。"
             )
-            submit_btn = gr.Button("提交推理", variant="primary")
+            
+            submit_btn = gr.Button("提交推理", variant="primary", scale=1)
 
-        with gr.Column(scale=2):
-            gr.Markdown("### 原始JSON输出")
-            output_json = gr.Textbox(label="API 返回的完整JSON结果", interactive=False, lines=10)
-            gr.Markdown("### 结果表格展示")
-            output_dataframe = gr.DataFrame(label="推理结果", interactive=False)
+        # 右侧列：输出结果
+        with gr.Column(scale=3):
+            gr.Markdown("### 3. 查看推理结果")
+            with gr.Tabs():
+                with gr.TabItem("结果表格"):
+                    output_dataframe = gr.DataFrame(
+                        label="推理结果详情", 
+                        interactive=False,
+                        # 根据您的示例输出定义列，Gradio会自动匹配
+                        headers=["subject", "prompt_source", "prompt_target", "layer_source", "is_correct_patched", "generations"]
+                    )
+                with gr.TabItem("原始JSON响应"):
+                    # 使用gr.JSON组件可以更好地格式化显示JSON
+                    output_json = gr.JSON(label="API 返回的完整JSON")
             
+    # 定义按钮点击事件
     submit_btn.click(
         fn=call_api,
-        inputs=input_csv,
+        inputs=[model_name_input, only_final_result_input, data_input],
         outputs=[output_json, output_dataframe]
     )
 
 # 启动Gradio应用
 if __name__ == "__main__":
+    # 使用share=True可以生成一个公开链接，方便他人访问
     demo.launch()