Upload app.py

app.py

@@ -225,25 +225,32 @@ def ner(text, model_type="bert"):
 
 
 # ======================== 关系抽取（RE） ========================
-def re_extract(entities, text):
+def re_extract(entities, text, use_bert_model=True, bert_model=None):
     # 参数校验
     if not entities or not text:
+        print("[DEBUG] 参数校验失败，实体或文本为空")
         return []
 
     # 实体类型过滤（根据业务需求调整）
     valid_entity_types = {"PER", "LOC", "ORG", "TITLE"}
     filtered_entities = [e for e in entities if e.get("type") in valid_entity_types]
 
+    if not filtered_entities:
+        print("[DEBUG] 未找到有效的实体")
+        return []
+
     # --------------------- 处理单实体场景 ---------------------
     if len(filtered_entities) == 1:
         single_relations = []
         ent = filtered_entities[0]
+        print(f"[DEBUG] 处理单实体：{ent['text']}，类型：{ent['type']}")
 
         # 规则1：人物职位检测
         if ent["type"] == "PER":
             position_keywords = ["CEO", "经理", "总监", "工程师", "教授"]
             for keyword in position_keywords:
                 if keyword in text:
+                    print(f"[DEBUG] 发现职位关键词：{keyword}")
                     single_relations.append({
                         "head": ent["text"],
                         "tail": keyword,
@@ -258,6 +265,7 @@ def re_extract(entities, text):
                 if verb in text:
                     match = re.search(fr"{ent['text']}{verb}(.*?)[，。]", text)
                     if match:
+                        print(f"[DEBUG] 发现位置关系：{ent['text']} {verb} {match.group(1)}")
                         single_relations.append({
                             "head": ent["text"],
                             "tail": match.group(1).strip(),
@@ -269,9 +277,10 @@ def re_extract(entities, text):
     # --------------------- 多实体关系抽取 ---------------------
     relations = []
 
-    # 方案1：使用ChatGLM抽取关系
-    if use_chatglm and len(filtered_entities) >= 2:
+    # 方案1：使用BERT模型进行关系抽取
+    if use_bert_model and len(filtered_entities) >= 2:
         try:
+            # 假设 BERT 模型是基于你自己训练的模型进行关系抽取
             entity_list = [e["text"] for e in filtered_entities]
             prompt = f"""请分析以下文本中的实体关系，严格按照JSON列表格式返回：
 文本内容：{text}
@@ -282,9 +291,8 @@ def re_extract(entities, text):
 3. 示例格式：[{{"head":"实体1", "tail":"实体2", "relation":"关系类型"}}]
 请直接返回JSON，不要多余内容："""
 
-            response = chatglm_model.chat(chatglm_tokenizer, prompt, temperature=0.01)
-            if isinstance(response, tuple):
-                response = response[0]
+            # 使用BERT模型进行关系抽取（这里假设模型函数是 `bert_model.predict`，具体调用方式按你模型接口调整）
+            response = bert_model.predict(prompt)
 
             # 增强JSON解析
             try:
@@ -305,11 +313,13 @@ def re_extract(entities, text):
                                 rel.get("relation") in valid_rel_types):
                             valid_relations.append(rel)
                     relations = valid_relations
+                else:
+                    print("[DEBUG] 未能解析出关系JSON")
             except Exception as e:
                 print(f"[DEBUG] 关系解析失败: {str(e)}")
 
         except Exception as e:
-            print(f"ChatGLM关系抽取异常: {str(e)}")
+            print(f"[DEBUG] BERT模型关系抽取异常: {str(e)}")
 
     # 方案2：规则兜底（当模型不可用或未抽取出关系时）
     if len(relations) == 0:
@@ -317,18 +327,21 @@ def re_extract(entities, text):
         location_matches = re.finditer(r'([^\s，。]+)[位于|坐落于|地处]([^\s，。]+)', text)
         for match in location_matches:
             head, tail = match.groups()
+            print(f"[DEBUG] 发现位于关系：{head} 位于 {tail}")
             relations.append({"head": head, "tail": tail, "relation": "位于"})
 
         # 规则2：A属于B
         belong_matches = re.finditer(r'([^\s，。]+)(属于|隶属于)([^\s，。]+)', text)
         for match in belong_matches:
             head, _, tail = match.groups()
+            print(f"[DEBUG] 发现属于关系：{head} 属于 {tail}")
             relations.append({"head": head, "tail": tail, "relation": "属于"})
 
         # 规则3：人物-机构关系
         person_org_pattern = r'([\u4e00-\u9fa5]{2,4})(现任|担任|就职于)([\u4e00-\u9fa5]+?公司|[\u4e00-\u9fa5]+?大学)'
         for match in re.finditer(person_org_pattern, text):
             head, _, tail = match.groups()
+            print(f"[DEBUG] 发现人物职位关系：{head} {tail}")
             relations.append({"head": head, "tail": tail, "relation": "任职于"})
 
     # 后��理：去重和验证
@@ -341,8 +354,11 @@ def re_extract(entities, text):
             head_exists = any(e["text"] == rel["head"] for e in filtered_entities)
             tail_exists = any(e["text"] == rel["tail"] for e in filtered_entities)
             if head_exists and tail_exists:
+                print(f"[DEBUG] 添加有效关系：{rel}")
                 final_relations.append(rel)
                 seen.add(key)
+            else:
+                print(f"[DEBUG] 无效关系：{rel}")
 
     return final_relations