Upload app.py

app.py

@@ -508,13 +508,42 @@ def process_text(text, model_type="bert"):
     return ent_text, rel_text, kg_text, f"{total_duration:.2f} 秒"
 
 
+from matplotlib.figure import Figure
+from io import BytesIO
+
+def generate_kg_image(entities, relations):
+    """
+    生成知识图谱的图片
+    """
+    fig = Figure(figsize=(10, 8))
+    ax = fig.add_subplot(1, 1, 1)
+    ax.axis("off")
+    
+    # 创建网络图
+    G = nx.DiGraph()
+    for entity in entities:
+        G.add_node(entity["text"], label=entity["type"])
+    for relation in relations:
+        G.add_edge(relation["head"], relation["tail"], label=relation["relation"])
+    
+    # 绘制网络图
+    pos = nx.spring_layout(G)
+    nx.draw(G, pos, with_labels=True, node_color="lightblue", ax=ax)
+    nx.draw_networkx_edge_labels(G, pos, edge_labels={(u, v): d["label"] for u, v, d in G.edges(data=True)}, ax=ax)
+    
+    # 保存为图片
+    buf = BytesIO()
+    fig.savefig(buf, format="png")
+    buf.seek(0)
+    return buf
+
 def process_file(file, model_type="bert"):
     try:
         with open(file.name, 'rb') as f:
             content = f.read()
 
         if len(content) > 5 * 1024 * 1024:
-            return "❌ 文件太大", "", "", ""
+            return "❌ 文件太大", "", "", "", None
 
         # 检测编码
         try:
@@ -529,15 +558,21 @@ def process_file(file, model_type="bert"):
                 except:
                     continue
             else:
-                return "❌ 编码解析失败", "", "", ""
+                return "❌ 编码解析失败", "", "", "", None
 
-        # 直接调用process_text处理文本
-        return process_text(text, model_type)
+        # 调用现有流程处理文本
+        ent_text, rel_text, kg_text, duration = process_text(text, model_type)
+        
+        # 生成知识图谱图片
+        entities, _ = ner(text, model_type)
+        relations, _ = re_extract(entities, text)
+        kg_image = generate_kg_image(entities, relations)
+        
+        return ent_text, rel_text, kg_text, duration, kg_image
         
     except Exception as e:
         logging.error(f"文件处理错误: {str(e)}")
-        return f"❌ 文件处理错误: {str(e)}", "", "", ""
-
+        return f"❌ 文件处理错误: {str(e)}", "", "", "", None
 
 
 # ======================== 模型评估与自动标注 ========================
@@ -665,8 +700,8 @@ with gr.Blocks(css="""
     with gr.Tab("🗂 文件分析"):
         file_input = gr.File(file_types=[".txt", ".json"])
         file_btn = gr.Button("上传并分析")
-        fout1, fout2, fout3, fout4 = gr.Textbox(), gr.Textbox(), gr.Textbox(), gr.Textbox()
-        file_btn.click(fn=process_file, inputs=[file_input, model_type], outputs=[fout1, fout2, fout3, fout4])
+        fout1, fout2, fout3, fout4, fout5 = gr.Textbox(), gr.Textbox(), gr.Textbox(), gr.Textbox(), gr.File(label="下载知识图谱图片")
+        file_btn.click(fn=process_file, inputs=[file_input, model_type], outputs=[fout1, fout2, fout3, fout4, fout5])
 
     with gr.Tab("📊 模型评估"):
         eval_file = gr.File(label="上传标注 JSON")