Update predictor.py

predictor.py

@@ -8,7 +8,7 @@ score_map = {'A': 5, 'B': 4, 'C': 3, 'D': 2, 'E': 1}
 
 
 class SentenceExtractor:
-    def __init__(self, main_keywords_path: str, eval_keywords_path: str, model_path: str = "model_quantized.onnx"):
+    def __init__(self, main_keywords_path: str, eval_keywords_path: str, model_path: str = "model_quantized.onnx", use_model: bool = False):
         """
         初始化句子提取器，加载主关键词、评估关键词库和评分模型
         :param main_keywords_path: 主关键词JSON文件路径
@@ -19,50 +19,80 @@ class SentenceExtractor:
         self.eval_keywords = self._load_eval_keywords(eval_keywords_path)
         self.all_keywords = self._merge_all_keywords()
         
-        # 加载ONNX评分模型
-        self.ort_session = ort.InferenceSession(model_path)
-        self.input_name = self.ort_session.get_inputs()[0].name
-        self.output_name = self.ort_session.get_outputs()[0].name
+        # 加载ONNX评分模型（可选）
+        self.use_model = use_model
+        self.ort_session = None
+        self.input_name = None
+        self.output_name = None
+        if self.use_model:
+            try:
+                self.ort_session = ort.InferenceSession(model_path)
+                self.input_name = self.ort_session.get_inputs()[0].name
+                self.output_name = self.ort_session.get_outputs()[0].name
+            except Exception as e:
+                print(f"ONNX 模型加载失败，回退到启发式打分: {e}")
+                self.use_model = False
 
     def _preprocess_text(self, text: str) -> np.ndarray:
         """
-        预处理文本，为模型输入做准备
-        注意：这里需要根据实际模型的输入要求进行调整
+        预处理文本（占位实现）。若启用 ONNX，请根据训练时的 tokenizer/embedding 改造。
         """
-        # 示例预处理 - 实际实现需与训练模型时的预处理一致
-        # 这里假设模型接受固定长度的词向量或嵌入
-        # 以下为示例代码，需根据实际模型修改
         max_seq_length = 128
-        # 简单的哈希特征示例（实际应使用与模型匹配的预处理）
         features = np.zeros((1, max_seq_length), dtype=np.float32)
-        for i, char in enumerate(text[:max_seq_length]):
-            features[0, i] = hash(char) % 1000 / 1000.0
+        for i, ch in enumerate(text[:max_seq_length]):
+            features[0, i] = (ord(ch) % 256) / 255.0
         return features
     
-    def _predict_grade(self, text: str) -> str:
-        """
-        使用ONNX模型预测文本评分等级
-        :param text: 输入句子
-        :return: 评分等级(A/B/C/D/E)
-        """
+    def _predict_grade_with_model(self, text: str) -> str:
         try:
-            # 预处理文本
+            if not self.ort_session:
+                return "C"
             input_data = self._preprocess_text(text)
-            
-            # 模型推理
             outputs = self.ort_session.run([self.output_name], {self.input_name: input_data})
-            
-            # 解析模型输出获取等级
-            # 假设模型输出是概率分布，取最大概率对应的等级
             predictions = outputs[0]
-            grade_index = np.argmax(predictions)
-            
-            # 将索引映射到等级A-E
+            grade_index = int(np.argmax(predictions))
             grades = ['A', 'B', 'C', 'D', 'E']
             return grades[grade_index]
         except Exception as e:
             print(f"模型预测出错: {e}")
-            return "C"  # 出错时返回默认等级
+            return "C"
+
+    def _predict_grade_heuristic(self, text: str) -> str:
+        score = 0
+        hit_any = False
+        for category in ["student_performance", "content_quality", "cross_scene"]:
+            cat_dict = self.eval_keywords.get(category, {})
+            for sentiment, weight in [["positive", 2], ["suggestion", 1], ["negative", -2], ["nature", 0]]:
+                for kw in cat_dict.get(sentiment, []):
+                    if kw and kw in text:
+                        score += weight
+                        hit_any = True
+        if not hit_any:
+            for _, kws in self.main_keywords.items():
+                if any(kw in text for kw in kws):
+                    return "C"
+            return "C"
+
+        if score >= 3:
+            return "A"
+        if score >= 1:
+            return "B"
+        if score == 0:
+            return "C"
+        if score <= -3:
+            return "E"
+        return "D"
+
+    def _predict_grade(self, text: str) -> str:
+        grade = self._predict_grade_heuristic(text)
+        if self.use_model:
+            model_grade = self._predict_grade_with_model(text)
+            # 简单融合策略：若模型比启发式高两档以上，则提升一档
+            order = {"A":5, "B":4, "C":3, "D":2, "E":1}
+            if order.get(model_grade,3) - order.get(grade,3) >= 2:
+                return ["A","B","C","D","E"][max(0, 5 - (order.get(grade,3)+1))]
+            return grade
+        return grade
     
     def _load_keywords(self, file_path: str) -> Dict[str, List[str]]:
         """加载主关键词文件"""