Add Gradio app for NER + RE

app.py

@@ -1,31 +1,80 @@
-import torch
-from transformers import BertTokenizer, BertModel
 import gradio as gr
-import re
-import os
-import json
-import pandas as pd
-import chardet
-from pyvis.network import Network
-import networkx as nx
+from transformers import BertTokenizerFast, BertForTokenClassification, BertForSequenceClassification
+import torch
 from pathlib import Path
+from pyvis.network import Network
+
+# 加载模型和分词器
+ner_tokenizer = BertTokenizerFast.from_pretrained("bert-base-chinese")
+ner_model = BertForTokenClassification.from_pretrained("bert-base-chinese", num_labels=10)
+
+re_model = BertForSequenceClassification.from_pretrained("bert-base-chinese", num_labels=5)
+re_tokenizer = BertTokenizerFast.from_pretrained("bert-base-chinese")
 
-# 初始化模型
-model_name = "bert-base-chinese"
-tokenizer = BertTokenizer.from_pretrained(model_name)
-model = BertModel.from_pretrained(model_name)
+# 定义标签和关系类型
+label_list = ["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC", "B-MISC", "I-MISC", "PAD"]
+relation_list = ["no_relation", "per-org", "per-loc", "org-loc", "org-misc"]
 
-# 知识图谱数据存储
+# 用于存储知识图谱
 knowledge_graph = {
-    "entities": set(),
+    "entities": [],
     "relations": []
 }
 
-def update_knowledge_graph(entities, relations):
-    for e in entities:
-        knowledge_graph["entities"].add((e['text'], e['type']))
-    for r in relations:
-        knowledge_graph["relations"].append((r['head'], r['tail'], r['relation']))
+def ner_predict(text):
+    inputs = ner_tokenizer(text, return_tensors="pt", truncation=True)
+    with torch.no_grad():
+        outputs = ner_model(**inputs).logits
+    predictions = torch.argmax(outputs, dim=2)
+    tokens = ner_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
+    predicted_labels = [label_list[label_id] for label_id in predictions[0].numpy()]
+    entities = []
+    current_entity = ""
+    current_label = ""
+    start = None
+    for idx, (token, label) in enumerate(zip(tokens, predicted_labels)):
+        if label.startswith("B-"):
+            if current_entity:
+                entities.append((current_entity, current_label, start, idx))
+            current_entity = token.replace("##", "")
+            current_label = label[2:]
+            start = idx
+        elif label.startswith("I-") and current_label == label[2:]:
+            current_entity += token.replace("##", "")
+        else:
+            if current_entity:
+                entities.append((current_entity, current_label, start, idx))
+                current_entity = ""
+                current_label = ""
+    if current_entity:
+        entities.append((current_entity, current_label, start, len(tokens)))
+    return entities
+
+def re_predict(text, entities):
+    relations = []
+    for i in range(len(entities)):
+        for j in range(len(entities)):
+            if i == j:
+                continue
+            head, tail = entities[i][0], entities[j][0]
+            input_text = f"{head} 和 {tail} 有什么关系？{text}"
+            inputs = re_tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
+            with torch.no_grad():
+                outputs = re_model(**inputs).logits
+            prediction = torch.argmax(outputs, dim=1).item()
+            if relation_list[prediction] != "no_relation":
+                relations.append((head, tail, relation_list[prediction]))
+    return relations
+
+def analyze_text(text):
+    entities = ner_predict(text)
+    relations = re_predict(text, entities)
+    entity_list = [f"{ent[0]} ({ent[1]}) [{ent[2]}, {ent[3]}]" for ent in entities]
+    relation_list_text = [f"{rel[0]} --[{rel[2]}]-> {rel[1]}" for rel in relations]
+    # 更新全局知识图谱
+    knowledge_graph["entities"] = [(ent[0], ent[1]) for ent in entities]
+    knowledge_graph["relations"] = relations
+    return "\n".join(entity_list), "\n".join(relation_list_text)
 
 def visualize_kg():
     net = Network(height="600px", width="100%", notebook=True, directed=True)
@@ -52,159 +101,31 @@ def visualize_kg():
     }
     """)
 
-    temp_file = "kg.html"
-    net.save_graph(temp_file)
-
-    # ✅ 返回 HTML 内容而不是文件路径
-    html_content = Path(temp_file).read_text(encoding="utf-8")
-    return html_content
-
-# ----------- NER 和 RE 抽取逻辑 -----------------
-def ner(text):
-    pattern_name = r"[赵钱孙李周吴郑王冯陈褚卫蒋沈韩杨朱秦尤许何吕施张孔曹严华金魏陶姜][\u4e00-\u9fa5]{1,2}"
-    pattern_id = r"\b[a-zA-Z_][a-zA-Z0-9_]{4,}\b"
-    entities = []
-
-    for match in re.finditer(pattern_name, text):
-        entities.append({
-            "text": match.group(),
-            "start": match.start(),
-            "end": match.end(),
-            "type": "PersonName"
-        })
-
-    for match in re.finditer(pattern_id, text):
-        if not any(e["start"] == match.start() for e in entities):
-            entities.append({
-                "text": match.group(),
-                "start": match.start(),
-                "end": match.end(),
-                "type": "UserID"
-            })
-
-    return sorted(entities, key=lambda x: x["start"])
-
-def re_extract(entities, text):
-    relations = []
-    if len(entities) >= 2:
-        for i in range(len(entities) - 1):
-            head = entities[i]["text"]
-            tail = entities[i + 1]["text"]
-            context = text[entities[i]["end"]:entities[i + 1]["start"]]
-            if "推荐" in context or "找" in context:
-                relation = "recommend"
-            elif "发送" in context or "发给" in context:
-                relation = "send_to"
-            elif "提到" in context or "说" in context:
-                relation = "mention"
-            else:
-                relation = "knows"
-            relations.append({
-                "head": head,
-                "tail": tail,
-                "relation": relation
-            })
-    return relations
-
-# ----------- 文本处理逻辑 -----------------
-def process_text(text):
-    entities = ner(text)
-    relations = re_extract(entities, text)
-    update_knowledge_graph(entities, relations)
-    kg_html = visualize_kg()
-    entity_output = "\n".join([f"{e['text']} ({e['type']}) [{e['start']}, {e['end']}]" for e in entities])
-    relation_output = "\n".join([f"{r['head']} --[{r['relation']}]-> {r['tail']}" for r in relations])
-    return entity_output, relation_output, kg_html
-
-# ----------- 文件上传处理逻辑 -----------------
-def detect_encoding(file_path):
-    with open(file_path, 'rb') as f:
-        raw_data = f.read(4096)
-    result = chardet.detect(raw_data)
-    return result['encoding'] if result['encoding'] else 'utf-8'
-
-def process_file(file):
-    ext = os.path.splitext(file.name)[-1].lower()
-    full_text = ""
-    warning = ""
-
-    try:
-        encoding = detect_encoding(file.name)
-
-        if ext == ".txt":
-            with open(file.name, "r", encoding=encoding) as f:
-                full_text = f.read()
-
-        elif ext == ".jsonl":
-            with open(file.name, "r", encoding=encoding) as f:
-                lines = f.readlines()
-                texts = []
-                skipped_lines = []
-                for i, line in enumerate(lines, start=1):
-                    try:
-                        obj = json.loads(line)
-                        texts.append(obj.get("text", ""))
-                    except Exception:
-                        skipped_lines.append(i)
-                full_text = "\n".join(texts)
-                if skipped_lines:
-                    warning = f"⚠️ 跳过 {len(skipped_lines)} 行无效 JSON（如第 {skipped_lines[0]} 行）\n\n"
-
-        elif ext == ".json":
-            with open(file.name, "r", encoding=encoding) as f:
-                data = json.load(f)
-                if isinstance(data, list):
-                    full_text = "\n".join([str(item.get("text", "")) for item in data])
-                elif isinstance(data, dict):
-                    full_text = data.get("text", "")
-                else:
-                    return "❌ JSON 文件格式无法解析", "", ""
-
-        elif ext == ".csv":
-            df = pd.read_csv(file.name, encoding=encoding)
-            if "text" in df.columns:
-                full_text = "\n".join(df["text"].astype(str))
-            else:
-                return "❌ CSV 中未找到 'text' 列", "", ""
-
-        else:
-            return f"❌ 不支持的文件格式：{ext}", "", ""
-
-    except Exception as e:
-        return f"❌ 文件读取错误：{str(e)}", "", ""
-
-    entity_out, relation_out, kg_html = process_text(full_text)
-    return warning + entity_out, relation_out, kg_html
-
-# ----------- Gradio UI -----------------
-with gr.Blocks() as demo:
-    gr.Markdown("""# 📱 微信聊天记录分析系统  
-    功能包括：实体识别（NER）、关系抽取（RE）和知识图谱可视化""")
-
-    with gr.Tab("✍️ 直接输入文本"):
-        with gr.Row():
-            input_text = gr.Textbox(label="输入聊天内容", lines=8, placeholder="请输入中文微信聊天记录")
-        analyze_btn = gr.Button("分析文本")
-        with gr.Row():
-            entity_output1 = gr.Textbox(label="识别出的实体")
-            relation_output1 = gr.Textbox(label="抽取的关系")
-        kg_html1 = gr.HTML(label="知识图谱可视化")
-        analyze_btn.click(fn=process_text, inputs=[input_text], outputs=[entity_output1, relation_output1, kg_html1])
-
-    with gr.Tab("📁 上传文件"):
-        file_input = gr.File(label="上传聊天记录文件", file_types=[".txt", ".jsonl", ".json", ".csv"])
-        analyze_file_btn = gr.Button("分析文件")
-        with gr.Row():
-            entity_output2 = gr.Textbox(label="识别出的实体")
-            relation_output2 = gr.Textbox(label="抽取的关系")
-        kg_html2 = gr.HTML(label="知识图谱可视化")
-        analyze_file_btn.click(fn=process_file, inputs=[file_input], outputs=[entity_output2, relation_output2, kg_html2])
-
-    with gr.Tab("🗺️ 完整知识图谱"):
-        gr.Markdown("## 当前累计构建的知识图谱")
-        refresh_btn = gr.Button("刷新图谱")
-        full_kg = gr.HTML()
-        refresh_btn.click(fn=lambda: visualize_kg(), outputs=full_kg)
-
+    # 保存 HTML 到 Hugging Face Spaces 可访问路径
+    file_path = "/home/user/kg.html"
+    net.save_graph(file_path)
+
+    # 返回 iframe HTML
+    return f'<iframe src="/file=kg.html" width="100%" height="600px" frameborder="0"></iframe>'
+
+# 搭建 Gradio 界面
+with gr.Blocks(title="Wechat Ner Re") as demo:
+    gr.Markdown("## 微信聊天记录结构化系统（NER + RE + 知识图谱）")
+    with gr.Row():
+        input_text = gr.Textbox(lines=5, label="请输入文本")
+        analyze_button = gr.Button("分析文本")
+    with gr.Row():
+        ner_output = gr.Textbox(label="识别出的实体")
+        re_output = gr.Textbox(label="抽取的关系")
+    analyze_button.click(analyze_text, inputs=input_text, outputs=[ner_output, re_output])
+
+    # 显示知识图谱
+    gr.Markdown("## 知识图谱可视化")
+    with gr.Row():
+        kg_button = gr.Button("生成知识图谱")
+    kg_html1 = gr.HTML(label="知识图谱可视化", show_label=True)
+    kg_button.click(fn=visualize_kg, outputs=kg_html1)
+
+# 启动应用
 if __name__ == "__main__":
     demo.launch()