Update predictor.py

predictor.py

@@ -4,6 +4,9 @@ import onnxruntime as ort
 import numpy as np
 from typing import List, Dict, Set, Tuple
 
+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"):
         """
@@ -148,36 +151,53 @@ class SentenceExtractor:
         return relevant_sentences, categorized_sentences
     
     def extract(self, text: str) -> Dict[str, any]:
-        """
-        提取文本中与关键词相关的句子并使用模型评分
-        :param text: 输入文本
-        :return: 包含相关句子、分类信息和评分的字典
-        """
-        if not text:
-            return {
-                "relevant_sentences": [], 
-                "scored_sentences": [],
-                "categorized_sentences": {}
-            }
-            
-        relevant_sentences, categorized_sentences = self._extract_relevant_sentences(text)
-        
-        # 为每个相关句子评分
-        scored_sentences = []
-        for sentence in relevant_sentences:
-            grade = self._predict_grade(sentence)
-            scored_sentences.append({
-                "sentence": sentence,
-                "grade": grade
-            })
-        
+    """
+    提取文本中与关键词相关的句子并使用模型评分
+    :param text: 输入文本
+    :return: 包含相关句子、分类信息、评分及综合等级的字典
+    """
+    if not text:
         return {
-            "relevant_sentences": relevant_sentences,
-            "scored_sentences": scored_sentences,
-            "categorized_sentences": categorized_sentences,
-            "count": len(relevant_sentences)
+            "relevant_sentences": [], 
+            "scored_sentences": [],
+            "categorized_sentences": {},
+            "count": 0,
+            "overall_grade": "C"  # 空文本默认C
         }
 
+    relevant_sentences, categorized_sentences = self._extract_relevant_sentences(text)
+    
+    # 为每个相关句子评分
+    scored_sentences = []
+    score_map = {'A': 5, 'B': 4, 'C': 3, 'D': 2, 'E': 1}
+    total_score = 0
+
+    for sentence in relevant_sentences:
+        grade = self._predict_grade(sentence)
+        scored_sentences.append({
+            "sentence": sentence,
+            "grade": grade
+        })
+        total_score += score_map.get(grade, 3)  # 没匹配默认C=3
+
+    # 计算综合得分（四舍五入）
+    overall_grade = "C"
+    if relevant_sentences:
+        avg_score = total_score / len(relevant_sentences)
+        rounded_score = int(round(avg_score))  # 四舍五入到整数1-5
+        # 反向映射
+        reverse_map = {5:'A', 4:'B', 3:'C', 2:'D', 1:'E'}
+        overall_grade = reverse_map.get(rounded_score, "C")
+
+    return {
+        "relevant_sentences": relevant_sentences,
+        "scored_sentences": scored_sentences,
+        "categorized_sentences": categorized_sentences,
+        "count": len(relevant_sentences),
+        "overall_grade": overall_grade
+    }
+
+
 # 使用示例
 if __name__ == "__main__":
     # 假设主关键词文件名为main_keywords.json