add app.py and requirements.txt

app.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import AutoTokenizer, AutoModel, BertTokenizer, BertModel
+from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline, AutoModel
 import gradio as gr
 import re
 import os
@@ -10,8 +10,9 @@ import time
 
 # ======================== 模型加载 ========================
 bert_model_name = "bert-base-chinese"
-bert_tokenizer = BertTokenizer.from_pretrained(bert_model_name)
-bert_model = BertModel.from_pretrained(bert_model_name)
+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")
 
 chatglm_model, chatglm_tokenizer = None, None
 use_chatglm = False
@@ -50,32 +51,25 @@ def ner(text, model_type="bert"):
     if model_type == "chatglm" and use_chatglm:
         try:
             prompt = f"请从以下文本中识别所有实体，用JSON格式返回:[{text}]"
-            response, _ = chatglm_model.chat(chatglm_tokenizer, prompt, temperature=0.1)
+            response = chatglm_model.chat(chatglm_tokenizer, prompt, temperature=0.1)
+            if isinstance(response, tuple):
+                response = response[0]
             entities = json.loads(response)
             return entities, time.time() - start_time
         except Exception as e:
             print(f"❌ ChatGLM 实体识别失败：{e}")
             return [], time.time() - start_time
 
-    name_pattern = r"([\u4e00-\u9fa5]{2,4})(?![的等地得啦啊哦])"
-    id_pattern = r"(?<!\S)([a-zA-Z_][a-zA-Z0-9_]{4,})(?![\u4e00-\u9fa5])"
-    entities, occupied = [], set()
-
-    def is_occupied(start, end):
-        return any(s <= start < e or s < end <= e for s, e in occupied)
-
-    for match in re.finditer(name_pattern, text):
-        start, end = match.start(1), match.end(1)
-        if not is_occupied(start, end):
-            entities.append({"text": match.group(1), "start": start, "end": end, "type": "人名"})
-            occupied.add((start, end))
-
-    for match in re.finditer(id_pattern, text):
-        start, end = match.start(1), match.end(1)
-        if not is_occupied(start, end):
-            entities.append({"text": match.group(1), "start": start, "end": end, "type": "用户ID"})
-            occupied.add((start, end))
-
+    # 使用微调的 BERT 中文 NER 模型
+    raw_results = bert_ner_pipeline(text)
+    entities = []
+    for r in raw_results:
+        entities.append({
+            "text": r["word"],
+            "start": r["start"],
+            "end": r["end"],
+            "type": r["entity_group"]
+        })
     return entities, time.time() - start_time
 
 # ======================== 关系抽取（RE） ========================
@@ -86,7 +80,9 @@ def re_extract(entities, text):
         entity_list = [e['text'] for e in entities]
         prompt = f"分析以下实体之间的关系：{entity_list}\n文本上下文：{text}"
         if use_chatglm:
-            response, _ = chatglm_model.chat(chatglm_tokenizer, prompt, temperature=0.1)
+            response = chatglm_model.chat(chatglm_tokenizer, prompt, temperature=0.1)
+            if isinstance(response, tuple):
+                response = response[0]
             return json.loads(response)
     except Exception as e:
         print(f"❌ ChatGLM 关系抽取失败：{e}")
@@ -103,32 +99,38 @@ def process_text(text, model_type="bert"):
     kg_text = visualize_kg_text()
     return ent_text, rel_text, kg_text, f"{duration:.2f} 秒"
 
-def process_file(file, model_type="bert"):
-    content = file.read()
-    if len(content) > 5 * 1024 * 1024:
-        return "❌ 文件太大", "", "", ""
-    encoding = chardet.detect(content)['encoding'] or 'utf-8'
-    text = content.decode(encoding)
-    return process_text(text, model_type)
-
 # ======================== 模型评估与自动标注 ========================
 def convert_telegram_json_to_eval_format(path):
     with open(path, encoding="utf-8") as f:
         data = json.load(f)
-    result = []
-    for m in data.get("messages", []):
-        if isinstance(m.get("text"), str):
-            result.append({"text": m["text"], "entities": []})
-        elif isinstance(m.get("text"), list):
-            txt = ''.join([x["text"] if isinstance(x, dict) else x for x in m["text"]])
-            result.append({"text": txt, "entities": []})
-    return result
+    if isinstance(data, dict) and "text" in data:
+        return [{"text": data["text"], "entities": [
+            {"text": data["text"][e["start"]:e["end"]]} for e in data.get("entities", [])
+        ]}]
+    elif isinstance(data, list):
+        return data
+    elif isinstance(data, dict) and "messages" in data:
+        result = []
+        for m in data.get("messages", []):
+            if isinstance(m.get("text"), str):
+                result.append({"text": m["text"], "entities": []})
+            elif isinstance(m.get("text"), list):
+                txt = ''.join([x["text"] if isinstance(x, dict) else x for x in m["text"]])
+                result.append({"text": txt, "entities": []})
+        return result
+    return []
 
 def evaluate_ner_model(data, model_type):
     y_true, y_pred = [], []
     for item in data:
-        gold = set(e['text'] for e in item['entities'])
-        pred, _ = ner(item['text'], model_type)
+        text = item["text"]
+        gold = set()
+        for e in item.get("entities", []):
+            if "text" in e:
+                gold.add(e["text"])
+            elif "start" in e and "end" in e:
+                gold.add(text[e["start"]:e["end"]])
+        pred, _ = ner(text, model_type)
         pred = set(e['text'] for e in pred)
         for ent in gold.union(pred):
             y_true.append(1 if ent in gold else 0)
@@ -143,7 +145,7 @@ def auto_annotate(file, model_type):
     return json.dumps(data, ensure_ascii=False, indent=2)
 
 def save_json(json_text):
-    fname = "auto_labeled.json"
+    fname = f"auto_labeled_{int(time.time())}.json"
     with open(fname, "w", encoding="utf-8") as f:
         f.write(json_text)
     return fname