Add Gradio app for NER + RE

app.py

@@ -24,9 +24,11 @@ knowledge_graph = {
 def update_knowledge_graph(entities, relations):
     """更新知识图谱数据"""
     for e in entities:
-        knowledge_graph["entities"].add((e['text'], e['type']))
+        if isinstance(e, dict) and 'text' in e and 'type' in e:
+            knowledge_graph["entities"].add((e['text'], e['type']))
     for r in relations:
-        knowledge_graph["relations"].append((r['head'], r['tail'], r['relation']))
+        if isinstance(r, dict) and all(k in r for k in ("head", "tail", "relation")):
+            knowledge_graph["relations"].append((r['head'], r['tail'], r['relation']))
 
 
 def visualize_kg():
@@ -34,17 +36,21 @@ def visualize_kg():
     net = Network(height="600px", width="100%", notebook=True, directed=True)
     node_map = {}
 
-    # 添加节点（带中文支持）
-    for idx, (name, type_) in enumerate(knowledge_graph["entities"]):
-        node_map[name] = idx
-        net.add_node(idx,
-                     label=name,
-                     title=f"类型：{type_}",
-                     group=type_,
-                     font={'size': 20, 'face': 'SimHei'})  # 使用中文字体
+    # 添加节点
+    idx = 0
+    for ent in knowledge_graph["entities"]:
+        if isinstance(ent, tuple) and len(ent) == 2:
+            name, type_ = ent
+            node_map[name] = idx
+            net.add_node(idx,
+                         label=name,
+                         title=f"类型：{type_}",
+                         group=type_,
+                         font={'size': 20, 'face': 'SimHei'})
+            idx += 1
 
     # 添加边
-    seen_edges = set()  # 防止重复边
+    seen_edges = set()
     for head, tail, relation in knowledge_graph["relations"]:
         if head in node_map and tail in node_map:
             edge_key = f"{head}-{tail}-{relation}"
@@ -55,7 +61,6 @@ def visualize_kg():
                              font={'size': 14})
                 seen_edges.add(edge_key)
 
-    # 优化布局配置
     net.set_options("""
     {
       "nodes": {
@@ -82,24 +87,19 @@ def visualize_kg():
     }
     """)
 
-    # 生成HTML并修复资源引用
     html = net.generate_html()
     html = html.replace('//cdnjs.cloudflare.com', 'https://cdnjs.cloudflare.com')
     html = html.replace('//unpkg.com', 'https://unpkg.com')
     return f'<div class="kg-graph">{html}</div>'
 
 
-# ----------- 增强的NER逻辑 -----------------
 def ner(text):
-    # 优化中文姓名识别（排除常见动词后缀）
     name_pattern = r"([赵钱孙李周吴郑王冯陈褚卫蒋沈韩杨朱秦尤许何吕施张孔曹严华金魏陶姜][\u4e00-\u9fa5]{1,2})(?![的地得啦啊呀])"
-    # 增强ID识别（支持带下划线和数字）
     id_pattern = r"(?<!\S)([a-zA-Z_][a-zA-Z0-9_]{4,})(?![\u4e00-\u9fa5])"
 
     entities = []
     occupied = set()
 
-    # 识别中文姓名
     for match in re.finditer(name_pattern, text):
         start, end = match.start(1), match.end(1)
         if not any(s <= start < e for (s, e) in occupied):
@@ -111,7 +111,6 @@ def ner(text):
             })
             occupied.update(range(start, end))
 
-    # 识别用户ID
     for match in re.finditer(id_pattern, text):
         start, end = match.start(1), match.end(1)
         if not any(s <= start < e for (s, e) in occupied):
@@ -126,7 +125,6 @@ def ner(text):
     return sorted(entities, key=lambda x: x["start"])
 
 
-# ----------- 改进的关系抽取逻辑 -----------------
 def re_extract(entities, text):
     relations = []
     triggers = {
@@ -136,17 +134,14 @@ def re_extract(entities, text):
     }
 
     for i in range(len(entities)):
-        # 检查前后两个窗口范围
         for j in range(max(0, i - 2), min(len(entities), i + 3)):
             if i == j:
                 continue
 
-            # 获取上下文内容
             ctx_start = entities[i]["end"]
             ctx_end = entities[j]["start"]
             context = text[ctx_start:ctx_end].strip()
 
-            # 处理@提及的情况
             if text.startswith('@', entities[i]["start"] - 1):
                 relations.append({
                     "head": entities[i]["text"],
@@ -155,7 +150,6 @@ def re_extract(entities, text):
                 })
                 continue
 
-            # 关系判断
             relation_type = "knows"
             for rel_type, keywords in triggers.items():
                 if any(kw in context for kw in keywords):
@@ -168,7 +162,6 @@ def re_extract(entities, text):
                 "relation": relation_type
             })
 
-    # 去重
     unique_relations = []
     seen = set()
     for rel in relations:
@@ -180,7 +173,6 @@ def re_extract(entities, text):
     return unique_relations
 
 
-# ----------- 文本处理流程 -----------------
 def process_text(text):
     try:
         entities = ner(text)
@@ -197,19 +189,12 @@ def process_text(text):
         )
         kg_html = visualize_kg()
 
-        # 调试日志
-        print(f"Entities: {entities}")
-        print(f"Relations: {relations}")
-        with open("debug_kg.html", "w", encoding="utf-8") as f:
-            f.write(kg_html)
-
         return entity_output, relation_output, gr.HTML(kg_html)
 
     except Exception as e:
         return f"处理出错: {str(e)}", "", gr.HTML()
 
 
-# ----------- 文件处理模块 -----------------
 def detect_encoding(file_path):
     with open(file_path, 'rb') as f:
         return chardet.detect(f.read(4096))['encoding'] or 'utf-8'
@@ -218,7 +203,6 @@ def detect_encoding(file_path):
 def process_file(file):
     ext = os.path.splitext(file.name)[-1].lower()
     full_text = ""
-    warning = ""
 
     try:
         encoding = detect_encoding(file.name)
@@ -262,7 +246,7 @@ def process_file(file):
         return f"❌ 文件处理错误: {str(e)}", "", gr.HTML()
 
 
-# ----------- Gradio界面 -----------------
+# Gradio UI
 css = """
 .kg-container {
     border: 1px solid #e0e0e0;
@@ -285,13 +269,12 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("""# 🚀 智能聊天记录分析系统  
-    **功能**: 实体识别 → 关系抽取 → 动态知识图谱""")
+    gr.Markdown("# 🚀 智能聊天记录分析系统\n**功能**: 实体识别 → 关系抽取 → 动态知识图谱")
 
     with gr.Tab("✍️ 文本分析"):
         gr.Markdown("### 直接输入聊天内容")
         input_text = gr.Textbox(label="输入内容", lines=8,
-                                placeholder="示例：\n张三@李四 请把需求文档_v2发送给王五\n李四回复：已发送至[email protected]")
+                                placeholder="示例：张三@李四 请把需求文档_v2发送给王五")
         analyze_btn = gr.Button("开始分析", variant="primary")
 
         with gr.Row():
@@ -315,8 +298,7 @@ with gr.Blocks(css=css) as demo:
 
     with gr.Tab("📁 文件分析"):
         gr.Markdown("### 上传聊天记录文件")
-        file_input = gr.File(label="选择文件",
-                             file_types=[".txt", ".json", ".jsonl", ".csv"])
+        file_input = gr.File(label="选择文件", file_types=[".txt", ".json", ".jsonl", ".csv"])
         file_btn = gr.Button("分析文件", variant="primary")
 
         with gr.Row():
@@ -343,4 +325,4 @@ with gr.Blocks(css=css) as demo:
         )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)