Update utils.py

utils.py

@@ -1,6 +1,7 @@
 import json
 import tempfile
 import re
+import os
 from collections import defaultdict
 from transformers import (
     AutoTokenizer,
@@ -11,25 +12,38 @@ import torch
 from pyvis.network import Network
 
 # -------------------------------
-# 实体识别模型（NER）
+# 模型配置
 # -------------------------------
-ner_tokenizer = AutoTokenizer.from_pretrained("ckiplab/bert-base-chinese-ner")
-ner_model = AutoModelForSequenceClassification.from_pretrained("ckiplab/bert-base-chinese-ner")
-ner_pipeline = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple")
+# 使用环境变量配置模型名称，便于在Hugging Face上部署时修改
+NER_MODEL_NAME = os.environ.get("NER_MODEL_NAME", "ckiplab/bert-base-chinese-ner")
+REL_MODEL_NAME = os.environ.get("REL_MODEL_NAME", "hfl/chinese-roberta-wwm-ext")
 
+# -------------------------------
+# 实体识别模型（NER）
+# -------------------------------
+try:
+    ner_tokenizer = AutoTokenizer.from_pretrained(NER_MODEL_NAME)
+    ner_model = AutoModelForSequenceClassification.from_pretrained(NER_MODEL_NAME)
+    ner_pipeline = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple")
+except Exception as e:
+    print(f"NER模型加载失败: {e}")
+    # 可以添加备选方案或错误处理逻辑
 
 # -------------------------------
 # 人物关系分类模型（使用 RoBERTa）
 # -------------------------------
-rel_model_name = "hfl/chinese-roberta-wwm-ext"  # 推荐的中文 RoBERTa 模型
-rel_tokenizer = AutoTokenizer.from_pretrained(rel_model_name)
-rel_model = AutoModelForSequenceClassification.from_pretrained(
-    rel_model_name, 
-    num_labels=6,  # 确保标签数量匹配
-    id2label={0: "夫妻", 1: "父子", 2: "朋友", 3: "师生", 4: "同事", 5: "其他"},
-    label2id={"夫妻": 0, "父子": 1, "朋友": 2, "师生": 3, "同事": 4, "其他": 5}
-)
-rel_model.eval()
+try:
+    rel_tokenizer = AutoTokenizer.from_pretrained(REL_MODEL_NAME)
+    rel_model = AutoModelForSequenceClassification.from_pretrained(
+        REL_MODEL_NAME, 
+        num_labels=6,  # 确保标签数量匹配
+        id2label={0: "夫妻", 1: "父子", 2: "朋友", 3: "师生", 4: "同事", 5: "其他"},
+        label2id={"夫妻": 0, "父子": 1, "朋友": 2, "师生": 3, "同事": 4, "其他": 5}
+    )
+    rel_model.eval()
+except Exception as e:
+    print(f"关系分类模型加载失败: {e}")
+    # 可以添加备选方案或错误处理逻辑
 
 # 关系分类的标签映射
 relation_id2label = {
@@ -41,55 +55,52 @@ legal_id2label = {
     0: "无违法", 1: "赌博", 2: "毒品", 3: "色情", 4: "诈骗", 5: "暴力"
 }
 
-
-def classify_relation_bert(e1, e2, context):
-    prompt = f"{e1}和{e2}的关系是？{context}"
-    inputs = rel_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
-    with torch.no_grad():
-        logits = rel_model(**inputs).logits
-    pred = torch.argmax(logits, dim=1).item()
-    probs = torch.nn.functional.softmax(logits, dim=1)
-    confidence = probs[0, pred].item()
-    return f"{relation_id2label[pred]}（置信度 {confidence:.2f}）"
-
-
 # -------------------------------
 # 聊天输入解析
 # -------------------------------
 def parse_input_file(file):
-    filename = file.name
-    if filename.endswith(".json"):
-        return json.load(file)
-    elif filename.endswith(".txt"):
-        content = file.read().decode("utf-8")
-        lines = content.strip().splitlines()
-        chat_data = []
-        for line in lines:
-            match = re.match(r"(\d{4}-\d{2}-\d{2}.*?) (.*?): (.*)", line)
-            if match:
-                _, sender, message = match.groups()
-                chat_data.append({"sender": sender, "message": message})
-        return chat_data
-    else:
-        raise ValueError("不支持的文件格式，请上传 JSON 或 TXT 文件")
-
+    """解析聊天文件，支持JSON和TXT格式"""
+    try:
+        filename = file.name
+        if filename.endswith(".json"):
+            return json.load(file)
+        elif filename.endswith(".txt"):
+            content = file.read().decode("utf-8")
+            lines = content.strip().splitlines()
+            chat_data = []
+            for line in lines:
+                match = re.match(r"(\d{4}-\d{2}-\d{2}.*?) (.*?): (.*)", line)
+                if match:
+                    _, sender, message = match.groups()
+                    chat_data.append({"sender": sender, "message": message})
+            return chat_data
+        else:
+            raise ValueError("不支持的文件格式，请上传JSON或TXT文件")
+    except Exception as e:
+        print(f"文件解析错误: {e}")
+        raise
 
 # -------------------------------
 # 实体提取函数
 # -------------------------------
 def extract_entities(text):
-    results = ner_pipeline(text)
-    people = set()
-    for r in results:
-        if r["entity_group"] == "PER":
-            people.add(r["word"])
-    return list(people)
-
+    """从文本中提取人物实���"""
+    try:
+        results = ner_pipeline(text)
+        people = set()
+        for r in results:
+            if r["entity_group"] == "PER":
+                people.add(r["word"])
+        return list(people)
+    except Exception as e:
+        print(f"实体提取错误: {e}")
+        return []
 
 # -------------------------------
 # 关系抽取函数（共现 + BERT 分类）
 # -------------------------------
 def extract_relations(chat_data, entities):
+    """分析人物之间的关系"""
     relations = defaultdict(lambda: defaultdict(lambda: {"count": 0, "contexts": []}))
 
     for entry in chat_data:
@@ -116,62 +127,84 @@ def extract_relations(chat_data, entities):
                 edges.append((e1, e2, relations[e1][e2]["count"], label))
     return edges
 
-
 # -------------------------------
 # 法律风险分析（黄赌毒等）函数
 # -------------------------------
-def classify_illegal_behavior(chat_context):
-    prompt = f"请分析以下聊天记录，判断是否涉及以下违法行为：赌博、毒品、色情、诈骗、暴力行为。\n聊天内容：{chat_context}\n请回答："
-    inputs = rel_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
-    with torch.no_grad():
-        logits = rel_model(**inputs).logits
-    pred = torch.argmax(logits, dim=1).item()
-    probs = torch.nn.functional.softmax(logits, dim=1)
-    confidence = probs[0, pred].item()
-    
-    return f"违法行为判断结果：{legal_id2label.get(pred, '未知')}（置信度 {confidence:.2f}）"
+def classify_relation_bert(e1, e2, context):
+    """使用BERT模型分析人物关系"""
+    try:
+        prompt = f"{e1}和{e2}的关系是？{context}"
+        inputs = rel_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+        with torch.no_grad():
+            logits = rel_model(**inputs).logits
+        pred = torch.argmax(logits, dim=1).item()
+        probs = torch.nn.functional.softmax(logits, dim=1)
+        confidence = probs[0, pred].item()
+        return f"{relation_id2label[pred]}（置信度 {confidence:.2f}）"
+    except Exception as e:
+        print(f"关系分类错误: {e}")
+        return "其他（置信度 0.00）"
 
+def classify_illegal_behavior(chat_context):
+    """分析聊天内容中的法律风险"""
+    try:
+        prompt = f"请分析以下聊天记录，判断是否涉及以下违法行为：赌博、毒品、色情、诈骗、暴力行为。\n聊天内容：{chat_context}\n请回答："
+        inputs = rel_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+        with torch.no_grad():
+            logits = rel_model(**inputs).logits
+        pred = torch.argmax(logits, dim=1).item()
+        probs = torch.nn.functional.softmax(logits, dim=1)
+        confidence = probs[0, pred].item()
+        return f"违法行为判断结果：{legal_id2label.get(pred, '未知')}（置信度 {confidence:.2f}）"
+    except Exception as e:
+        print(f"法律风险分析错误: {e}")
+        return "违法行为判断结果：未知（置信度 0.00）"
 
 # -------------------------------
 # 图谱绘制
 # -------------------------------
 def draw_graph(entities, relations):
-    g = Network(height="600px", width="100%", notebook=False)
-    g.barnes_hut()
-    for ent in entities:
-        g.add_node(ent, label=ent)
-    for e1, e2, weight, label in relations:
-        g.add_edge(e1, e2, value=weight, title=f"{label}（互动{weight}次）", label=label)
-    tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
-    g.show(tmp_file.name)
-    with open(tmp_file.name, 'r', encoding='utf-8') as f:
-        return f.read()
-
+    """生成人物关系图谱"""
+    try:
+        g = Network(height="600px", width="100%", notebook=False)
+        g.barnes_hut()
+        for ent in entities:
+            g.add_node(ent, label=ent)
+        for e1, e2, weight, label in relations:
+            g.add_edge(e1, e2, value=weight, title=f"{label}（互动{weight}次）", label=label)
+        tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
+        g.show(tmp_file.name)
+        with open(tmp_file.name, 'r', encoding='utf-8') as f:
+            return f.read()
+    except Exception as e:
+        print(f"图谱绘制错误: {e}")
+        return "<h3>图谱生成失败</h3><p>请检查输入数据是否有效</p>"
 
 # -------------------------------
 # 主流程函数
 # -------------------------------
 def analyze_chat(file):
+    """分析聊天记录的主函数"""
     if file is None:
-        return "请上传聊天文件", "", ""
+        return "请上传聊天文件", "", "", ""
     
     try:
         content = parse_input_file(file)
     except Exception as e:
-        return f"读取文件失败: {e}", "", ""
+        return f"读取文件失败: {e}", "", "", ""
 
     text = "\n".join([entry["sender"] + ": " + entry["message"] for entry in content])
     entities = extract_entities(text)
     if not entities:
-        return "未识别到任何人物实体", "", ""
+        return "未识别到任何人物实体", "", "", ""
 
     relations = extract_relations(content, entities)
     if not relations:
-        return "未发现人物之间的关系", "", ""
+        return "未发现人物之间的关系", "", "", ""
 
     # 法律风险分析
     illegal_behavior_results = [classify_illegal_behavior(msg["message"]) for msg in content]
 
     graph_html = draw_graph(entities, relations)
 
-    return str(entities), str(relations), graph_html, "\n".join(illegal_behavior_results)        
+    return str(entities), str(relations), graph_html, "\n".join(illegal_behavior_results)