add app.py and requirements.txt

app.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import BertTokenizer, BertModel, LlamaTokenizer, LlamaForCausalLM
+from transformers import AutoTokenizer, AutoModel, BertTokenizer, BertModel
 import gradio as gr
 import re
 import os
@@ -9,25 +9,23 @@ import chardet
 from pyvis.network import Network
 import time
 
-# 初始化模型（适配 Hugging Face Secrets）
+# 初始化模型
 bert_model_name = "bert-base-chinese"
 bert_tokenizer = BertTokenizer.from_pretrained(bert_model_name)
 bert_model = BertModel.from_pretrained(bert_model_name)
 
-llama_model_name = os.getenv("LLAMA_MODEL_NAME", "meta-llama/Llama-2-7b-chat-hf")
-access_token = os.getenv("HUGGINGFACE_HUB_TOKEN")
-
-llama_tokenizer = LlamaTokenizer.from_pretrained(
-    llama_model_name,
-    use_auth_token=access_token  # 使用旧参数名确保兼容性
-)
-
-llama_model = LlamaForCausalLM.from_pretrained(
-    llama_model_name,
-    use_auth_token=access_token,
-    torch_dtype=torch.float16,  # 添加量化配置
-    device_map="auto"
+# 加载中文模型 ChatGLM3-6B
+chatglm_model_name = "THUDM/chatglm3-6b"
+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="auto",
+    torch_dtype=torch.float16
+).eval()
 
 # 知识图谱数据存储
 knowledge_graph = {
@@ -106,14 +104,25 @@ def visualize_kg():
 def ner(text, model_type="bert"):
     start_time = time.time()
     if model_type == "bert":
-        tokenizer = bert_tokenizer
-        model = bert_model
-    elif model_type == "llama":
-        tokenizer = llama_tokenizer
-        model = llama_model
+        # BERT 中文实体识别（原逻辑保留）
+        name_pattern = r"([\u4e00-\u9fa5]{2,4})(?![的等地得啦啊哦])"
+        id_pattern = r"(?<!\S)([a-zA-Z_][a-zA-Z0-9_]{4,})(?![\\u4e00-\\u9fa5])"
+    else:
+        # ChatGLM 增强实体识别
+        response, _ = chatglm_model.chat(
+            chatglm_tokenizer,
+            f"请从以下文本中识别所有实体，用JSON格式返回:[{text}]",
+            temperature=0.1
+        )
+        try:
+            entities = json.loads(response)
+            return entities, time.time() - start_time
+        except:
+            pass
 
-    name_pattern = r"([\u4e00-\u9fa5]{2,4})(?![\u7684\u5730\u5f97\u5566\u554a\u5440])"
-    id_pattern = r"(?<!\S)([a-zA-Z_][a-zA-Z0-9_]{4,})(?![\u4e00-\u9fa5])"
+        # 如果模型响应失败，使用备用正则
+        name_pattern = r"([\\u4e00-\\u9fa5]{2,4})(?![的等地得啦啊哦])"
+        id_pattern = r"(?<!\S)([a-zA-Z_][a-zA-Z0-9_]{4,})"
 
     entities = []
     occupied = set()
@@ -128,7 +137,7 @@ def ner(text, model_type="bert"):
                 "text": match.group(1),
                 "start": start,
                 "end": end,
-                "type": "PersonName"
+                "type": "人名"
             })
             occupied.add((start, end))
 
@@ -139,7 +148,7 @@ def ner(text, model_type="bert"):
                 "text": match.group(1),
                 "start": start,
                 "end": end,
-                "type": "UserID"
+                "type": "用户ID"
             })
             occupied.add((start, end))
 
@@ -149,15 +158,26 @@ def ner(text, model_type="bert"):
 
 def re_extract(entities, text):
     relations = []
-    for i, entity1 in enumerate(entities):
-        for j, entity2 in enumerate(entities):
-            if i != j:
-                relation = {
-                    "head": entity1['text'],
-                    "tail": entity2['text'],
-                    "relation": "联系"
-                }
-                relations.append(relation)
+    if len(entities) < 2:
+        return relations
+
+    # 使用ChatGLM分析关系
+    entity_list = [e['text'] for e in entities]
+    prompt = f"分析以下实体之间的关系：{entity_list}\n文本上下文：{text}"
+    response, _ = chatglm_model.chat(chatglm_tokenizer, prompt, temperature=0.1)
+
+    try:
+        relations = json.loads(response)
+    except:
+        # 备用简单关系生成
+        for i in range(len(entities)):
+            for j in range(i + 1, len(entities)):
+                relations.append({
+                    "head": entities[i]['text'],
+                    "tail": entities[j]['text'],
+                    "relation": "相关"
+                })
+
     return relations
 
 
@@ -177,7 +197,7 @@ def process_text(text, model_type="bert"):
         )
         kg_html = visualize_kg()
 
-        return entity_output, relation_output, gr.HTML(kg_html), f"\u5904\u7406\u65f6\u95f4：{processing_time:.2f}\u79d2"
+        return entity_output, relation_output, gr.HTML(kg_html), f"处理时间：{processing_time:.2f}秒"
 
     except Exception as e:
         return f"处理出错: {str(e)}", "", gr.HTML(), ""
@@ -185,10 +205,7 @@ def process_text(text, model_type="bert"):
 
 def process_file(file, model_type="bert"):
     try:
-        # 读取文件内容（适配 Hugging Face 文件系统）
         content_bytes = file.read()
-
-        # 文件大小限制（5MB）
         if len(content_bytes) > 5 * 1024 * 1024:
             return "❌ 文件大小超过5MB限制", "", gr.HTML(), ""
 
@@ -226,46 +243,31 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# 🚀 智能聊天记录分析系统\n**功能**: 实体识别 → 关系抽取 → 动态知识图谱")
+    gr.Markdown("# 🚀 智能聊天记录分析系统（ChatGLM3-6B版）")
 
     with gr.Tab("✍️ 文本分析"):
-        gr.Markdown("### 直接输入聊天内容")
         input_text = gr.Textbox(label="输入内容", lines=8,
                                 placeholder="示例：张三@李四 请把需求文档_v2发送给王五")
-        model_type = gr.Radio(["bert", "llama"], label="选择模型", value="bert")
+        model_type = gr.Radio(["bert", "chatglm"], label="选择模型", value="bert")
         analyze_btn = gr.Button("开始分析", variant="primary")
 
         with gr.Row():
-            entity_output = gr.Textbox(label="识别的实体", lines=6, interactive=False)
-            relation_output = gr.Textbox(label="提取的关系", lines=6, interactive=False)
+            entity_output = gr.Textbox(label="识别的实体", lines=6)
+            relation_output = gr.Textbox(label="提取的关系", lines=6)
         kg_output = gr.HTML(label="知识图谱")
-        time_output = gr.Textbox(label="处理时间", interactive=False)
-
-        analyze_btn.click(
-            process_text,
-            inputs=[input_text, model_type],
-            outputs=[entity_output, relation_output, kg_output, time_output],
-            show_progress="full"
-        )
+        time_output = gr.Textbox(label="处理时间")
 
     with gr.Tab("📄 文件分析"):
-        gr.Markdown("### 上传文件进行分析（支持 .txt, .jsonl, .json, .csv 格式）")
-        file_input = gr.File(label="选择文件", type="file")
+        file_input = gr.File(label="选择文件", file_types=[".txt", ".csv", ".json"])
         analyze_file_btn = gr.Button("开始分析文件", variant="primary")
-        file_entity_output = gr.Textbox(label="识别的实体", lines=6, interactive=False)
-        file_relation_output = gr.Textbox(label="提取的关系", lines=6, interactive=False)
+        file_entity_output = gr.Textbox(label="识别的实体", lines=6)
+        file_relation_output = gr.Textbox(label="提取的关系", lines=6)
         file_kg_output = gr.HTML(label="知识图谱")
-        file_time_output = gr.Textbox(label="处理时间", interactive=False)
+        file_time_output = gr.Textbox(label="处理时间")
 
-        analyze_file_btn.click(
-            process_file,
-            inputs=[file_input, model_type],
-            outputs=[file_entity_output, file_relation_output, file_kg_output, file_time_output],
-            show_progress="full"
-        )
+    analyze_btn.click(process_text, [input_text, model_type],
+                      [entity_output, relation_output, kg_output, time_output])
+    analyze_file_btn.click(process_file, [file_input, model_type],
+                           [file_entity_output, file_relation_output, file_kg_output, file_time_output])
 
-demo.launch(
-    server_name="0.0.0.0",
-    server_port=7860,
-    debug=False
-)
+demo.launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -1,10 +1,11 @@
 gradio==3.50.2
 transformers==4.39.3
 torch>=2.1.0
-accelerate>=0.27.0
-sentencepiece>=0.2.0 
 pandas>=2.0.0
 chardet>=5.0.0
 networkx>=3.0
 pyvis>=0.3.2
-python-dotenv>=1.0.0
+python-dotenv>=1.0.0
+sentencepiece>=0.2.0
+cpm-kernels>=1.0.11 
+accelerate>=0.27.0