utils.py

import json
import tempfile
import re
import os
from collections import defaultdict
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    pipeline,
)
import torch
from pyvis.network import Network

# -------------------------------
# 模型配置
# -------------------------------
# 使用环境变量配置模型名称，便于在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）
# -------------------------------
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 = {
    0: "夫妻", 1: "父子", 2: "朋友", 3: "师生", 4: "同事", 5: "其他"
}

# 法律风险分析的标签映射
legal_id2label = {
    0: "无违法", 1: "赌博", 2: "毒品", 3: "色情", 4: "诈骗", 5: "暴力"
}

# -------------------------------
# 聊天输入解析
# -------------------------------
def parse_input_file(file):
    """解析聊天文件，支持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):
    """从文本中提取人物实体"""
    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:
        msg = entry["message"]
        found = [e for e in entities if e in msg]
        for i in range(len(found)):
            for j in range(i + 1, len(found)):
                e1, e2 = found[i], found[j]
                relations[e1][e2]["count"] += 1
                relations[e1][e2]["contexts"].append(msg)
                relations[e2][e1]["count"] += 1
                relations[e2][e1]["contexts"].append(msg)

    edges = []
    for e1 in relations:
        for e2 in relations[e1]:
            if e1 < e2:
                context_text = " ".join(relations[e1][e2]["contexts"])
                # 截断过长的文本
                max_context_length = 500  # 根据需要调整
                if len(context_text) > max_context_length:
                    context_text = context_text[:max_context_length] + "..."
                label = classify_relation_bert(e1, e2, context_text)
                edges.append((e1, e2, relations[e1][e2]["count"], label))
    return edges

# -------------------------------
# 法律风险分析（黄赌毒等）函数
# -------------------------------
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):
    """生成人物关系图谱"""
    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 "请上传聊天文件", "", "", ""
    
    try:
        content = parse_input_file(file)
    except Exception as e:
        return f"读取文件失败: {e}", "", "", ""

    text = "\n".join([entry["sender"] + ": " + entry["message"] for entry in content])
    entities = extract_entities(text)
    if not entities:
        return "未识别到任何人物实体", "", "", ""

    relations = extract_relations(content, entities)
    if not relations:
        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)