Update predictor.py

predictor.py

@@ -2,20 +2,40 @@ import json
 import re
 import onnxruntime as ort
 import numpy as np
-from typing import List, Dict, Set
+from typing import List, Dict, Set, Optional
 
 
 score_map = {'A': 5, 'B': 4, 'C': 3, 'D': 2, 'E': 1}
 
 
 class SentenceExtractor:
-    def __init__(self, eval_keywords_path: str, model_path: str = "distilled_model.onnx"):
+    def __init__(
+        self,
+        eval_keywords_path: str,
+        model_path: str = "distilled_model.onnx",
+        *,
+        # 分句与聚合相关的可配置开关
+        merge_leading_punct: bool = True,
+        min_sentence_char_len: int = 6,
+        aggregation_mode: str = "max",  # 可选："max" | "mean"
+        # 加减号阈值（>0 / <0 为原逻辑；建议适度提高到 2/-2）
+        word_score_plus_threshold: int = 1,
+        word_score_minus_threshold: int = -1,
+    ):
         self.eval_keywords = self._load_eval_keywords(eval_keywords_path)
         self.all_keywords = self._extract_all_keywords()
 
         self.ort_session = None
         self.input_name = None
         self.output_name = None
+        # 配置项
+        self.merge_leading_punct = merge_leading_punct
+        self.min_sentence_char_len = max(0, int(min_sentence_char_len))
+        self.aggregation_mode = aggregation_mode.lower().strip()
+        if self.aggregation_mode not in {"max", "mean"}:
+            self.aggregation_mode = "max"
+        self.word_score_plus_threshold = int(word_score_plus_threshold)
+        self.word_score_minus_threshold = int(word_score_minus_threshold)
         try:
             self.ort_session = ort.InferenceSession(model_path)
             self.input_name = self.ort_session.get_inputs()[0].name
@@ -114,17 +134,48 @@ class SentenceExtractor:
     def _split_into_sentences(self, text: str) -> List[str]:
         if not text:
             return []
+
+        # 先按强终止符切分
         normalized = re.sub(r'([。！？.!?])', r'\1\n', text)
         normalized = re.sub(r'[；;]\s*', '；\n', normalized)
         candidates = [s.strip() for s in re.split(r'[\r\n]+', normalized) if s.strip()]
-        sentences: List[str] = []
+
+        # 长句再按逗号细分
+        rough_sentences: List[str] = []
         for s in candidates:
             if len(s) > 80 and not re.search(r'[。！？.!?；;]', s):
                 parts = re.split(r'[，,]', s)
-                sentences.extend([p.strip() for p in parts if p.strip()])
+                rough_sentences.extend([p.strip() for p in parts if p.strip()])
             else:
-                sentences.append(s)
-        return sentences
+                rough_sentences.append(s)
+
+        # 合并以标点开头的碎片，并过滤超短句
+        sentences: List[str] = []
+        leading_punct_pattern = r'^[，,。；;：:、\s]+'
+        for s in rough_sentences:
+            if self.merge_leading_punct and re.match(leading_punct_pattern, s):
+                # 去掉前缀标点后并入上一句
+                cleaned = re.sub(leading_punct_pattern, '', s)
+                if sentences:
+                    sentences[-1] = f"{sentences[-1]}{cleaned}"
+                else:
+                    if cleaned:
+                        sentences.append(cleaned)
+                continue
+
+            # 过滤极短句（去标点长度）
+            plain = re.sub(r'[，,。；;：:、!！?？\s]', '', s)
+            if self.min_sentence_char_len > 0 and len(plain) < self.min_sentence_char_len:
+                # 不直接丢弃：若有上一句，合并
+                if sentences:
+                    sentences[-1] = f"{sentences[-1]}{s}"
+                else:
+                    sentences.append(s)
+                continue
+
+            sentences.append(s)
+
+        return [s.strip() for s in sentences if s and s.strip()]
 
     def _fuzzy_match_keyword(self, sentence: str, keyword: str) -> bool:
         """更严格的中文关键词匹配。
@@ -271,16 +322,21 @@ class SentenceExtractor:
 
         comprehensive_grade = "C"
         if relevant_sentences:
-            avg_score = total_sentence_score / len(relevant_sentences)
-            rounded_score = int(round(avg_score))
             reverse_map = {5: 'A', 4: 'B', 3: 'C', 2: 'D', 1: 'E'}
-            comprehensive_grade = reverse_map.get(rounded_score, "C")
+            if self.aggregation_mode == "max":
+                # 取最高等级（更鲁棒，避免短碎句拉低均值）
+                max_score = max(score_map.get(item["grade"], 3) for item in scored_sentences)
+                comprehensive_grade = reverse_map.get(max_score, "C")
+            else:
+                avg_score = total_sentence_score / len(relevant_sentences)
+                rounded_score = int(round(avg_score))
+                comprehensive_grade = reverse_map.get(rounded_score, "C")
 
         word_scores = self._calculate_word_scores(text)
         final_grade = comprehensive_grade
-        if word_scores["total_score"] > 0:
+        if word_scores["total_score"] > self.word_score_plus_threshold:
             final_grade = comprehensive_grade + "+"
-        elif word_scores["total_score"] < 0:
+        elif word_scores["total_score"] < self.word_score_minus_threshold:
             final_grade = comprehensive_grade + "-"
 
         return {