Update utils.py

utils.py

@@ -4,7 +4,6 @@ import re
 from collections import defaultdict
 from transformers import (
     AutoTokenizer,
-    AutoModelForTokenClassification,
     AutoModelForSequenceClassification,
     pipeline,
 )
@@ -15,14 +14,14 @@ from pyvis.network import Network
 # 实体识别模型（NER）
 # -------------------------------
 ner_tokenizer = AutoTokenizer.from_pretrained("ckiplab/bert-base-chinese-ner")
-ner_model = AutoModelForTokenClassification.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")
 
 
 # -------------------------------
-# 人物关系分类模型（BERT 分类器）
+# 人物关系分类模型（使用 RoBERTa）
 # -------------------------------
-rel_model_name = "uer/roberta-base-finetuned-baike-chinese-relation-extraction"
+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)
 rel_model.eval()
@@ -160,4 +159,4 @@ def analyze_chat(file):
 
     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)