Upload 2 files

app.py

@@ -10,13 +10,13 @@ import time
 # ======================== 数据库模块 ========================
 import pymysql
 from configparser import ConfigParser
+from contextlib import contextmanager
 
-
+@contextmanager
 def get_db_connection():
     config = ConfigParser()
     config.read('db_config.ini')
-
-    return pymysql.connect(
+    conn = pymysql.connect(
         host=config.get('mysql', 'host'),
         user=config.get('mysql', 'user'),
         password=config.get('mysql', 'password'),
@@ -25,7 +25,10 @@ def get_db_connection():
         charset=config.get('mysql', 'charset', fallback='utf8mb4'),
         cursorclass=pymysql.cursors.DictCursor
     )
-
+    try:
+        yield conn
+    finally:
+        conn.close()
 
 def save_to_db(table, data):
     conn = None
@@ -122,30 +125,27 @@ def visualize_kg_text():
 
 # ======================== 实体识别（NER） ========================
 def merge_adjacent_entities(entities):
-    merged = []
-    for entity in entities:
-        if not merged:
-            merged.append(entity)
-            continue
+    if not entities:
+        return entities
 
+    merged = [entities[0]]
+    for entity in entities[1:]:
         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"]
-            }
+                entity["start"] == last["end"]):
+            last["text"] += entity["text"]
+            last["end"] = entity["end"]
         else:
             merged.append(entity)
+
     return merged
 
 
 def ner(text, model_type="bert"):
     start_time = time.time()
+
+    # 如果使用的是 ChatGLM 模型，执行 ChatGLM 的NER
     if model_type == "chatglm" and use_chatglm:
         try:
             prompt = f"""请从以下文本中识别所有实体，严格按照JSON列表格式返回，每个实体包含text、type、start、end字段：
@@ -155,47 +155,35 @@ def ner(text, model_type="bert"):
             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)
+                valid_entities = [ent for ent in entities if all(k in ent for k in ("text", "type", "start", "end"))]
                 return valid_entities, time.time() - start_time
             except Exception as e:
-                print(f"JSON 解析失败: {e}")
+                print(f"JSON解析失败: {e}")
                 return [], time.time() - start_time
         except Exception as e:
-            print(f"ChatGLM 调用失败：{e}")
+            print(f"ChatGLM调用失败: {e}")
             return [], time.time() - start_time
 
-    # 使用微调的 BERT 中文 NER 模型
+    # 使用BERT NER
+    text_chunks = [text[i:i + 510] for i in range(0, len(text), 510)]  # 安全分段
     raw_results = []
-    max_len = 510  # 安全一点，留一点空余
-    text_chunks = [text[i:i + max_len] for i in range(0, len(text), max_len)]
-
     for idx, chunk in enumerate(text_chunks):
         chunk_results = bert_ner_pipeline(chunk)
-        # 修正每个 chunk 里识别的实体在整体文本中的位置
         for r in chunk_results:
-            r["start"] += idx * max_len
-            r["end"] += idx * max_len
+            r["start"] += idx * 510
+            r["end"] += idx * 510
         raw_results.extend(chunk_results)
 
-    entities = []
-    for r in raw_results:
-        mapped_type = LABEL_MAPPING.get(r['entity_group'], r['entity_group'])
-        entities.append({
-            "text": r['word'].replace(' ', ''),
-            "start": r['start'],
-            "end": r['end'],
-            "type": mapped_type
-        })
-
-    # 执行合并处理
+    entities = [{
+        "text": r['word'].replace(' ', ''),
+        "start": r['start'],
+        "end": r['end'],
+        "type": LABEL_MAPPING.get(r['entity_group'], r['entity_group'])
+    } for r in raw_results]
+
     entities = merge_adjacent_entities(entities)
     return entities, time.time() - start_time
 
@@ -349,7 +337,7 @@ def process_file(file, model_type="bert"):
             text = content.decode(encoding)
         except UnicodeDecodeError:
             # 尝试常见中文编码
-            for enc in ['gb18030', 'utf-16', 'big5']:
+            for enc in ['gb18030', 'utf-16', 'big5'] :
                 try:
                     text = content.decode(enc)
                     break
@@ -363,6 +351,7 @@ def process_file(file, model_type="bert"):
         return f"❌ 文件处理错误: {str(e)}", "", "", ""
 
 
+
 # ======================== 模型评估与自动标注 ========================
 def convert_telegram_json_to_eval_format(path):
     with open(path, encoding="utf-8") as f:

db_config.ini

@@ -0,0 +1,6 @@
+[mysql]
+host = localhost
+user = root
+password = 123456
+database = entity_kg
+charset = utf8mb4