Update predictor.py

predictor.py

@@ -1,113 +1,141 @@
-# predictor.py
-
 import json
-import numpy as np
-import onnxruntime as ort
-from transformers import BertTokenizer
 import re
+from typing import List, Dict, Set, Tuple
 
-class Predictor:
-    def __init__(self):
+class SentenceExtractor:
+    def __init__(self, main_keywords_path: str, eval_keywords_path: str):
         """
-        在服务启动时，一次性加载所有必要的模型和文件。
+        初始化句子提取器，加载主关键词和评估关键词库
+        :param main_keywords_path: 主关键词JSON文件路径
+        :param eval_keywords_path: 评估关键词库(JSON)文件路径
         """
-        # 1. 加载分词器 (Tokenizer)
-        # Hugging Face Spaces会自动下载git仓库中的所有文件到当前目录
-        self.tokenizer = BertTokenizer.from_pretrained('.')
-
-        # 2. 加载ONNX模型并创建推理会话
-        self.ort_session = ort.InferenceSession('model_quantized.onnx')
-
-        # 3. 加载关键词词集
-        with open('evaluation_keywords2.json', 'r', encoding='utf-8') as f:
-            self.keywords = json.load(f)
+        self.main_keywords = self._load_keywords(main_keywords_path)
+        self.eval_keywords = self._load_eval_keywords(eval_keywords_path)
+        # 合并所有关键词用于快速查找
+        self.all_keywords = self._merge_all_keywords()
         
-        # 4. 定义等级映射
-        self.id2label = {0: 'a', 1: 'b', 2: 'c', 3: 'd', 4: 'e'}
-        self.label2score = {'a': 5, 'b': 4, 'c': 3, 'd': 2, 'e': 1} # 用于计算平均值
-
-    def _extract_relevant_sentences(self, text):
+    def _load_keywords(self, file_path: str) -> Dict[str, List[str]]:
+        """加载主关键词文件"""
+        try:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                return json.load(f)
+        except Exception as e:
+            print(f"加载主关键词文件失败: {e}")
+            return {}
+    
+    def _load_eval_keywords(self, file_path: str) -> Dict[str, Dict[str, List[str]]]:
+        """加载评估关键词库(evaluation_keywords2.json)"""
+        try:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                return json.load(f)
+        except Exception as e:
+            print(f"加载评估关键词库失败: {e}")
+            return {}
+    
+    def _merge_all_keywords(self) -> Set[str]:
+        """合并所有关键词到一个集合中，用于快速查找"""
+        keywords_set = set()
+        
+        # 添加主关键词
+        for category, keywords in self.main_keywords.items():
+            keywords_set.update(keywords)
+            
+        # 添加评估关键词库中的所有关键词
+        for category, types in self.eval_keywords.items():
+            for type_, keywords in types.items():
+                keywords_set.update(keywords)
+                
+        return keywords_set
+    
+    def _split_into_sentences(self, text: str) -> List[str]:
+        """将文本分割成句子"""
+        # 简单的句子分割正则，可根据需要优化
+        sentence_endings = re.compile(r'(?<=[。！？,.!?])\s+')
+        sentences = sentence_endings.split(text)
+        return [s.strip() for s in sentences if s.strip()]
+    
+    def _extract_relevant_sentences(self, text: str) -> Tuple[List[str], Dict[str, List[str]]]:
         """
-        根据关键词提取相关的句子。
+        提取与关键词相关的句子
+        :param text: 输入文本
+        :return: 相关句子列表和按类别分组的句子字典
         """
-        # 使用正则表达式按标点符号分割句子，更准确
-        sentences = re.split(r'[。！？]', text)
+        sentences = self._split_into_sentences(text)
         relevant_sentences = []
-        for sentence in sentences:
-            if not sentence:
-                continue
-            for keyword in self.keywords:
-                if keyword in sentence:
-                    relevant_sentences.append(sentence)
-                    break # 找到一个关键词就添加，避免重复
-        return relevant_sentences
-
-    def _predict_single_sentence(self, sentence):
-        """
-        对单个句子进行模型推理，返回预测的等级标签。
-        """
-        # 使用分词器处理文本
-        inputs = self.tokenizer(sentence, return_tensors="np", padding='max_length', truncation=True, max_length=128)
-        
-        # 准备ONNX模型的输入
-        ort_inputs = {self.ort_session.get_inputs()[0].name: inputs['input_ids']}
+        categorized_sentences = {
+            "main": [],
+            "student_performance": {"positive": [], "negative": [], "nature": [], "suggestion": []},
+            "content_quality": {"positive": [], "negative": [], "nature": [], "suggestion": []},
+            "cross_scene": {"positive": [], "negative": [], "nature": [], "suggestion": []}
+        }
         
-        # 执行推理
-        ort_outs = self.ort_session.run(None, ort_inputs)
+        for sentence in sentences:
+            # 检查是否包含主关键词
+            main_keyword_matched = False
+            for category, keywords in self.main_keywords.items():
+                for keyword in keywords:
+                    if keyword in sentence:
+                        relevant_sentences.append(sentence)
+                        categorized_sentences["main"].append(sentence)
+                        main_keyword_matched = True
+                        break
+                if main_keyword_matched:
+                    break
+            
+            # 检查评估关键词库中的关键词
+            for category in ["student_performance", "content_quality", "cross_scene"]:
+                if category not in self.eval_keywords:
+                    continue
+                    
+                for sentiment in ["positive", "negative", "nature", "suggestion"]:
+                    if sentiment not in self.eval_keywords[category]:
+                        continue
+                        
+                    for keyword in self.eval_keywords[category][sentiment]:
+                        if keyword in sentence and sentence not in categorized_sentences[category][sentiment]:
+                            # 如果还没添加到相关句子列表，则添加
+                            if sentence not in relevant_sentences:
+                                relevant_sentences.append(sentence)
+                            categorized_sentences[category][sentiment].append(sentence)
         
-        # 处理输出结果
-        prediction = np.argmax(ort_outs[0], axis=1)[0]
-        return self.id2label[prediction]
-
-    def predict(self, text):
+        return relevant_sentences, categorized_sentences
+    
+    def extract(self, text: str) -> Dict[str, any]:
         """
-        执行完整的预测流程：提取句子 -> 逐句评分 -> 计算平均等级。
-        这是暴露给app.py调用的主方法。
+        提取文本中与关键词相关的句子
+        :param text: 输入文本
+        :return: 包含相关句子和分类信息的字典
         """
-        # 步骤1: 提取包含关键词的句子
-        relevant_sentences = self._extract_relevant_sentences(text)
-
-        if not relevant_sentences:
-            return {
-                "grade": "c",  # 如果没有找到相关句子，返回一个默认的中间等级
-                "summary": "文本中未检测到可用于评价的关键词句，无法进行有效分析。",
-                "analyzed_sentences_count": 0
-            }
-
-        # 步骤2: 对每个相关句子进行评分
-        scores = []
-        for sentence in relevant_sentences:
-            label = self._predict_single_sentence(sentence)
-            scores.append(self.label2score[label])
-
-        # 步骤3: 计算平均分并转换为最终等级
-        if not scores:
-             return {
-                "grade": "c",
-                "summary": "虽然找到相关句子，但模型未能给出评分。",
-                "analyzed_sentences_count": len(relevant_sentences)
-            }
-
-        average_score = sum(scores) / len(scores)
+        if not text:
+            return {"relevant_sentences": [], "categorized_sentences": {}}
+            
+        relevant_sentences, categorized_sentences = self._extract_relevant_sentences(text)
         
-        # 将平均分四舍五入后映射回最终等级
-        final_grade = ""
-        if average_score >= 4.5:
-            final_grade = "a"
-        elif average_score >= 3.5:
-            final_grade = "b"
-        elif average_score >= 2.5:
-            final_grade = "c"
-        elif average_score >= 1.5:
-            final_grade = "d"
-        else:
-            final_grade = "e"
-
-        # 步骤4: 生成总结性文本
-        summary = f"系统分析了 {len(relevant_sentences)} 个关键句子，综合评定等级为“{final_grade.upper()}”。"
-
         return {
-            "grade": final_grade,
-            "summary": summary,
-            "analyzed_sentences_count": len(relevant_sentences)
-        }
+            "relevant_sentences": relevant_sentences,
+            "categorized_sentences": categorized_sentences,
+            "count": len(relevant_sentences)
+        }
+
+# 使用示例
+if __name__ == "__main__":
+    # 假设主关键词文件名为main_keywords.json
+    extractor = SentenceExtractor(
+        main_keywords_path="main_keywords.json",
+        eval_keywords_path="evaluation_keywords2.json"
+    )
+    
+    sample_text = """
+    该学生表现优异，团队合作能力强，在项目中展现了很强的创新能力。
+    但代码质量不高，存在安全漏洞，需要加强测试。
+    项目文档完整，符合行业标准，具有很好的应用价值。
+    建议加强代码审查，提高系统安全性，优化算法效率。
+    """
+    
+    result = extractor.extract(sample_text)
+    print(f"提取到 {result['count']} 个相关句子：")
+    for i, sent in enumerate(result['relevant_sentences'], 1):
+        print(f"{i}. {sent}")
+    
+    print("\n按类别分组：")
+    print(json.dumps(result['categorized_sentences'], ensure_ascii=False, indent=2))