Upload app.py

app.py

@@ -422,55 +422,135 @@ def re_extract(entities, text, use_bert_model=True):
 # ======================== 文本分析主流程 ========================
 def create_knowledge_graph(entities, relations):
     """
-    创建知识图谱的文本表示
+    创建交互式网络图形式的知识图谱
     """
-    result = []
-    
-    # 去重实体
-    unique_entities = []
-    seen = set()
-    for e in entities:
-        if (e['text'], e['type']) not in seen:
-            seen.add((e['text'], e['type']))
-            unique_entities.append(e)
+    # 创建一个新的网络图
+    net = Network(height="600px", width="100%", bgcolor="#ffffff", font_color="black", directed=True)
     
-    # 去重关系
-    unique_relations = []
-    seen = set()
-    for r in relations:
-        if (r['head'], r['relation'], r['tail']) not in seen:
-            seen.add((r['head'], r['relation'], r['tail']))
-            unique_relations.append(r)
+    # 设置实体类型的颜色映射
+    entity_colors = {
+        'PER': '#FF6B6B',  # 人物-红色
+        'ORG': '#4ECDC4',  # 组织-青色
+        'LOC': '#45B7D1',  # 地点-蓝色
+        'TIME': '#96CEB4', # 时间-绿色
+        'TITLE': '#D4A5A5' # 职位-粉色
+    }
     
-    # 添加标题
-    result.append("# 知识图谱展示\n")
+    # 添加实体节点
+    added_nodes = set()
+    for entity in entities:
+        if entity['text'] not in added_nodes:
+            node_color = entity_colors.get(entity['type'], '#D3D3D3')
+            net.add_node(
+                entity['text'],
+                label=entity['text'],
+                title=f"类型: {entity['type']}",
+                color=node_color,
+                size=20,
+                font={'size': 16}
+            )
+            added_nodes.add(entity['text'])
     
-    # 添加图例
-    result.append("## 图例")
-    result.append("- 👤 人物 (PER)")
-    result.append("- 🏢 组织 (ORG)")
-    result.append("- 📍 地点 (LOC)")
-    result.append("- 👔 职位 (TITLE)\n")
+    # 添加关系边
+    for relation in relations:
+        if relation['head'] in added_nodes and relation['tail'] in added_nodes:
+            net.add_edge(
+                relation['head'],
+                relation['tail'],
+                label=relation['relation'],
+                title=relation['relation'],
+                arrows={'to': {'enabled': True, 'type': 'arrow'}},
+                color={'color': '#666666'},
+                font={'size': 12}
+            )
     
-    # 添加实体
-    result.append("## 实体列表")
-    icons = {
-        'PER': '👤',
-        'ORG': '🏢',
-        'LOC': '📍',
-        'TIME': '⏰',
-        'TITLE': '👔'
+    # 设置物理布局参数
+    net.set_options('''
+    {
+        "nodes": {
+            "shape": "dot",
+            "shadow": true
+        },
+        "edges": {
+            "smooth": {
+                "type": "continuous",
+                "forceDirection": "none"
+            },
+            "shadow": true
+        },
+        "physics": {
+            "barnesHut": {
+                "gravitationalConstant": -2000,
+                "centralGravity": 0.3,
+                "springLength": 200,
+                "springConstant": 0.04,
+                "damping": 0.09
+            },
+            "minVelocity": 0.75
+        },
+        "interaction": {
+            "hover": true,
+            "navigationButtons": true,
+            "keyboard": true
+        }
     }
-    for e in unique_entities:
-        icon = icons.get(e['type'], '•')
-        result.append(f"- {icon} **{e['text']}** `{e['type']}`")
-    
-    # 添加关系
-    result.append("\n## 关系列表")
-    for r in unique_relations:
-        result.append(f"- **{r['head']}** ➡️ *{r['relation']}* ➡️ **{r['tail']}**")
+    ''')
     
-    return "\n".join(result)
+    # 生成HTML文件
+    try:
+        # 创建临时目录（如果不存在）
+        temp_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "temp")
+        os.makedirs(temp_dir, exist_ok=True)
+        
+        # 生成唯一的文件名
+        output_path = os.path.join(temp_dir, f"kg_{int(time.time())}.html")
+        
+        # 保存图谱
+        net.save_graph(output_path)
+        
+        # 读取生成的HTML文件内容
+        with open(output_path, 'r', encoding='utf-8') as f:
+            html_content = f.read()
+        
+        # 删除临时文件
+        os.remove(output_path)
+        
+        # 修改HTML内容以适应Gradio界面
+        html_content = html_content.replace('height: 600px', 'height: 700px')
+        
+        # 添加图例
+        legend_html = f"""
+        <div style="margin-bottom: 10px; padding: 10px; background-color: #f8f9fa; border-radius: 5px;">
+            <div style="font-weight: bold; margin-bottom: 5px;">图例说明：</div>
+            <div style="display: flex; gap: 15px; flex-wrap: wrap;">
+                <div style="display: flex; align-items: center; gap: 5px;">
+                    <div style="width: 15px; height: 15px; background: {entity_colors['PER']}; border-radius: 50%;"></div>
+                    <span>人物 (PER)</span>
+                </div>
+                <div style="display: flex; align-items: center; gap: 5px;">
+                    <div style="width: 15px; height: 15px; background: {entity_colors['ORG']}; border-radius: 50%;"></div>
+                    <span>组织 (ORG)</span>
+                </div>
+                <div style="display: flex; align-items: center; gap: 5px;">
+                    <div style="width: 15px; height: 15px; background: {entity_colors['LOC']}; border-radius: 50%;"></div>
+                    <span>地点 (LOC)</span>
+                </div>
+                <div style="display: flex; align-items: center; gap: 5px;">
+                    <div style="width: 15px; height: 15px; background: {entity_colors['TITLE']}; border-radius: 50%;"></div>
+                    <span>职位 (TITLE)</span>
+                </div>
+            </div>
+        </div>
+        """
+        
+        # 将图例添加到HTML内容中
+        html_content = legend_html + html_content
+        
+        return html_content
+        
+    except Exception as e:
+        logging.error(f"生成知识图谱失败: {str(e)}")
+        return f"<div class='error'>生成知识图谱失败: {str(e)}</div>"
 
 def process_text(text, model_type="bert"):
     """