add app.py and requirements.txt

app.py

@@ -9,46 +9,59 @@ from sklearn.metrics import precision_score, recall_score, f1_score
 import time
 
 # ======================== 模型加载 ========================
-bert_model_name = "bert-base-chinese"
-bert_tokenizer = AutoTokenizer.from_pretrained(bert_model_name)
-bert_ner_model = AutoModelForTokenClassification.from_pretrained("ckiplab/bert-base-chinese-ner")
-bert_ner_pipeline = pipeline("ner", model=bert_ner_model, tokenizer=bert_tokenizer, aggregation_strategy="simple")
+NER_MODEL_NAME = "uer/roberta-base-finetuned-cluener2020-chinese"
+bert_tokenizer = AutoTokenizer.from_pretrained(NER_MODEL_NAME)
+bert_ner_model = AutoModelForTokenClassification.from_pretrained(NER_MODEL_NAME)
+bert_ner_pipeline = pipeline(
+    "ner",
+    model=bert_ner_model,
+    tokenizer=bert_tokenizer,
+    aggregation_strategy="first"
+)
+
+LABEL_MAPPING = {
+    "address": "LOC",
+    "company": "ORG",
+    "name": "PER",
+    "organization": "ORG",
+    "position": "TITLE"
+}
 
 chatglm_model, chatglm_tokenizer = None, None
 use_chatglm = False
 try:
-    chatglm_model_name = "THUDM/chatglm-6b-int4"  # 4-bit量化版本
-    chatglm_tokenizer = AutoTokenizer.from_pretrained(
-        chatglm_model_name,
-        trust_remote_code=True
-    )
+    chatglm_model_name = "THUDM/chatglm-6b-int4"
+    chatglm_tokenizer = AutoTokenizer.from_pretrained(chatglm_model_name, trust_remote_code=True)
     chatglm_model = AutoModel.from_pretrained(
         chatglm_model_name,
         trust_remote_code=True,
         device_map="cpu",
-        torch_dtype=torch.float32  # 必须使用float32
+        torch_dtype=torch.float32
     ).eval()
     use_chatglm = True
-    print("✅ 4-bit量化版ChatGLM加载成功（需6GB内存）")
+    print("✅ 4-bit量化版ChatGLM加载成功")
 except Exception as e:
-    print(f"❌ 量化模型加载失败: {e}")
+    print(f"❌ ChatGLM加载失败: {e}")
 
 # ======================== 知识图谱结构 ========================
 knowledge_graph = {"entities": set(), "relations": set()}
 
-
 def update_knowledge_graph(entities, relations):
     for e in entities:
         if isinstance(e, dict) and 'text' in e and 'type' in e:
             knowledge_graph["entities"].add((e['text'], e['type']))
-
+    # 修改4：添加关系去重逻辑
+    existing_relations = {frozenset({r[0], r[1], r[2]}) for r in knowledge_graph["relations"]}
     for r in relations:
         if isinstance(r, dict) and all(k in r for k in ("head", "tail", "relation")):
-            # 标准化关系方向
-            relation_tuple = (r['head'], r['tail'], r['relation'])
-            reverse_tuple = (r['tail'], r['head'], r['relation'])
-            if reverse_tuple not in knowledge_graph["relations"]:
-                knowledge_graph["relations"].add(relation_tuple)
+            new_rel = frozenset({r['head'], r['tail'], r['relation']})
+            if new_rel not in existing_relations:
+                knowledge_graph["relations"].add((r['head'], r['tail'], r['relation']))
+
+def visualize_kg_text():
+    nodes = [f"{ent[0]} ({ent[1]})" for ent in knowledge_graph["entities"]]
+    edges = [f"{h} --[{r}]-> {t}" for h, t, r in knowledge_graph["relations"]]
+    return "\n".join(["📌 实体:"] + nodes + ["", "📎 关系:"] + edges)
 
 
 def visualize_kg_text():
@@ -58,50 +71,57 @@ def visualize_kg_text():
 
 
 # ======================== 实体识别（NER） ========================
+def merge_adjacent_entities(entities):
+    merged = []
+    for entity in entities:
+        if not merged:
+            merged.append(entity)
+            continue
+
+        last = merged[-1]
+        # 合并相邻的同类型实体
+        if (entity["type"] == last["type"] and
+                entity["start"] == last["end"] and
+                entity["text"] not in last["text"]):
+            merged[-1] = {
+                "text": last["text"] + entity["text"],
+                "type": last["type"],
+                "start": last["start"],
+                "end": entity["end"]
+            }
+        else:
+            merged.append(entity)
+    return merged
+
+
 def ner(text, model_type="bert"):
     start_time = time.time()
     if model_type == "chatglm" and use_chatglm:
-        try:
-            prompt = f"""请从以下文本中识别所有实体，严格按照JSON列表格式返回，每个实体包含text、type、start、end字段：
-示例：[{{"text": "北京", "type": "LOC", "start": 0, "end": 2}}]
-文本：{text}"""
-            response = chatglm_model.chat(chatglm_tokenizer, prompt, temperature=0.1)
-            if isinstance(response, tuple):
-                response = response[0]
-
-            # 增强 JSON 解析
-            try:
-                json_str = re.search(r'\[.*\]', response, re.DOTALL).group()
-                entities = json.loads(json_str)
-                # 验证字段
-                valid_entities = []
-                for ent in entities:
-                    if all(k in ent for k in ("text", "type", "start", "end")):
-                        valid_entities.append(ent)
-                return valid_entities, time.time() - start_time
-            except Exception as e:
-                print(f"JSON 解析失败: {e}")
-                return [], time.time() - start_time
-        except Exception as e:
-            print(f"ChatGLM 调用失败：{e}")
-            return [], time.time() - start_time
+    # ... [原有ChatGLM代码保持不变] ...
 
-    # 使用微调的 BERT 中文 NER 模型
+    # 修改6：优化BERT模型处理流程
     raw_results = bert_ner_pipeline(text)
     entities = []
     for r in raw_results:
+        mapped_type = LABEL_MAPPING.get(r['entity_group'], r['entity_group'])
         entities.append({
-            "text": r["word"],
-            "start": r["start"],
-            "end": r["end"],
-            "type": r["entity_group"]
+            "text": r['word'].replace(' ', ''),
+            "start": r['start'],
+            "end": r['end'],
+            "type": mapped_type
         })
-    return entities, time.time() - start_time
 
+    # 执行合并处理
+    entities = merge_adjacent_entities(entities)
+    return entities, time.time() - start_time
 
 # ======================== 关系抽取（RE） ========================
 def re_extract(entities, text):
-    if len(entities) < 2:
+    # 修改7：添加实体类型过滤
+    valid_entity_types = {"PER", "LOC", "ORG"}
+    filtered_entities = [e for e in entities if e["type"] in valid_entity_types]
+
+    if len(filtered_entities) < 2:
         return []
 
     relations = []
@@ -204,30 +224,55 @@ def convert_telegram_json_to_eval_format(path):
 
 def evaluate_ner_model(data, model_type):
     y_true, y_pred = [], []
+    POS_TOLERANCE = 1  # 允许的位置误差
+
     for item in data:
         text = item["text"]
         gold_entities = []
         for e in item.get("entities", []):
             if "text" in e and "type" in e:
-                # 使用哈希避免重复
-                gold_entities.append(f"{e['text']}|{e['type']}|{e.get('start', -1)}|{e.get('end', -1)}")
-
-        pred_entities = []
-        pred, _ = ner(text, model_type)
-        for e in pred:
-            pred_entities.append(f"{e['text']}|{e['type']}|{e['start']}|{e['end']}")
+                # 标准化标签
+                norm_type = LABEL_MAPPING.get(e["type"], e["type"])
+                gold_entities.append({
+                    "text": e["text"],
+                    "type": norm_type,
+                    "start": e.get("start", -1),
+                    "end": e.get("end", -1)
+                })
+
+        pred_entities, _ = ner(text, model_type)
+
+        # 构建对比集合
+        all_entities = set()
+        # 处理标注数据
+        for g in gold_entities:
+            key = f"{g['text']}|{g['type']}|{g['start']}|{g['end']}"
+            all_entities.add(key)
+
+        # 处理预测结果
+        pred_set = set()
+        for p in pred_entities:
+            # 允许位置误差
+            matched = False
+            for g in gold_entities:
+                if (p["text"] == g["text"] and
+                        p["type"] == g["type"] and
+                        abs(p["start"] - g["start"]) <= POS_TOLERANCE and
+                        abs(p["end"] - g["end"]) <= POS_TOLERANCE):
+                    matched = True
+                    break
+            pred_set.add(matched)
 
-        # 创建所有可能的实体集合
-        all_entities = set(gold_entities + pred_entities)
-        for ent in all_entities:
-            y_true.append(1 if ent in gold_entities else 0)
-            y_pred.append(1 if ent in pred_entities else 0)
+        # 构建指标
+        y_true.extend([1] * len(gold_entities))
+        y_pred.extend([1 if m else 0 for m in pred_set])
 
     if not y_true:
         return "⚠️ 无有效标注数据"
 
-    return f"Precision: {precision_score(y_true, y_pred):.2f}\nRecall: {recall_score(y_true, y_pred):.2f}\nF1: {f1_score(y_true, y_pred):.2f}"
-
+    return (f"Precision: {precision_score(y_true, y_pred, zero_division=0):.2f}\n"
+            f"Recall: {recall_score(y_true, y_pred, zero_division=0):.2f}\n"
+            f"F1: {f1_score(y_true, y_pred, zero_division=0):.2f}")
 
 def auto_annotate(file, model_type):
     data = convert_telegram_json_to_eval_format(file.name)
@@ -245,7 +290,10 @@ def save_json(json_text):
 
 
 # ======================== Gradio 界面 ========================
-with gr.Blocks(css=".kg-graph {height: 500px;}") as demo:
+with gr.Blocks(css="""
+    .kg-graph {height: 500px; overflow-y: auto;}
+    .warning {color: #ff6b6b;}
+""") as demo:
     gr.Markdown("# 🤖 聊天记录实体关系识别系统")
 
     with gr.Tab("📄 文本分析"):