Update predictor.py

predictor.py

@@ -1,7 +1,5 @@
 import json
 import re
-import os
-import hashlib
 import onnxruntime as ort
 import numpy as np
 from typing import List, Dict, Set, Optional
@@ -24,17 +22,6 @@ class SentenceExtractor:
         word_score_plus_threshold: int = 1,
         word_score_minus_threshold: int = -1,
     ):
-        # 统一以文件所在目录为根，避免工作目录不同导致找不到资源
-        self.base_dir = os.path.dirname(os.path.abspath(__file__))
-        self.tokenizer_dir = self.base_dir
-
-        # 允许传相对路径：自动转绝对
-        if not os.path.isabs(model_path):
-            model_path = os.path.join(self.base_dir, model_path)
-
-        if not os.path.isabs(eval_keywords_path):
-            eval_keywords_path = os.path.join(self.base_dir, eval_keywords_path)
-
         self.eval_keywords = self._load_eval_keywords(eval_keywords_path)
         self.all_keywords = self._extract_all_keywords()
 
@@ -50,43 +37,21 @@ class SentenceExtractor:
         self.word_score_plus_threshold = int(word_score_plus_threshold)
         self.word_score_minus_threshold = int(word_score_minus_threshold)
         try:
-            # 强制使用 CPU provider，避免某些环境下选择到不可用的 GPU provider 导致加载失败
-            self.ort_session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
+            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
             print("ONNX 模型加载成功")
-            self.model_loaded: bool = True
         except Exception as e:
             print(f"ONNX 模型加载失败: {e}")
             self.ort_session = None
-            self.model_loaded = False
-
-        # 记录模型文件信息，便于排查“用错模型”问题
-        try:
-            self.model_path_abs: Optional[str] = os.path.abspath(model_path)
-            self.model_sha256: Optional[str] = None
-            if os.path.exists(model_path):
-                sha = hashlib.sha256()
-                with open(model_path, 'rb') as f:
-                    for chunk in iter(lambda: f.read(8192), b''):
-                        sha.update(chunk)
-                self.model_sha256 = sha.hexdigest()
-        except Exception:
-            self.model_path_abs = None
-            self.model_sha256 = None
 
     def _preprocess_text(self, text: str) -> np.ndarray:
         try:
             from transformers import AutoTokenizer
-            # 1) 优先从与脚本同目录加载本地 tokenizer（部署一起带上 tokenizer.json 等文件）
             try:
-                tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_dir, local_files_only=True)
+                tokenizer = AutoTokenizer.from_pretrained(".", local_files_only=True)
             except Exception:
-                try:
-                    tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_dir)
-                except Exception:
-                    # 2) 兜底：在线模型（需要外网）
-                    tokenizer = AutoTokenizer.from_pretrained("uer/chinese_roberta_L-4_H-256")
+                tokenizer = AutoTokenizer.from_pretrained("uer/chinese_roberta_L-4_H-256")
             inputs = tokenizer(
                 text,
                 truncation=True,
@@ -103,16 +68,15 @@ class SentenceExtractor:
                 features[0, i] = (ord(ch) % 256) / 255.0
             return features
 
-    def _predict_grade_with_model(self, text: str) -> Dict[str, any]:
+    def _predict_grade_with_model(self, text: str) -> str:
         try:
             if not self.ort_session:
                 word_score = self._calculate_word_scores(text)["total_score"]
-                grade = "C"
                 if word_score > 1:
-                    grade = "B"
+                    return "B"
                 if word_score < -1:
-                    grade = "D"
-                return {"grade": grade, "source": "rule", "word_score_total": word_score}
+                    return "D"
+                return "C"
 
             inputs = self._preprocess_text(text)
 
@@ -142,29 +106,15 @@ class SentenceExtractor:
             predictions = outputs[0]
             grade_index = int(np.argmax(predictions))
             grades = ['A', 'B', 'C', 'D', 'E']
-            probs = self._softmax(predictions)[0].tolist()
-            return {
-                "grade": grades[grade_index],
-                "source": "model",
-                "prob": float(probs[grade_index]),
-                "probs": probs,
-                "logits": predictions[0].tolist(),
-            }
+            return grades[grade_index]
         except Exception as e:
             print(f"模型预测出错: {e}")
             word_score = self._calculate_word_scores(text)["total_score"]
-            grade = "C"
             if word_score > 1:
-                grade = "B"
+                return "B"
             if word_score < -1:
-                grade = "D"
-            return {"grade": grade, "source": "rule", "word_score_total": word_score}
-
-    @staticmethod
-    def _softmax(x: np.ndarray) -> np.ndarray:
-        x = x - np.max(x, axis=-1, keepdims=True)
-        exp_x = np.exp(x)
-        return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
+                return "D"
+            return "C"
 
     def _load_eval_keywords(self, file_path: str) -> Dict[str, Dict[str, List[str]]]:
         try:
@@ -365,17 +315,9 @@ class SentenceExtractor:
         scored_sentences = []
         total_sentence_score = 0
         for sentence in relevant_sentences:
-            info = self._predict_grade_with_model(sentence)
-            grade = info.get("grade", "C")
+            grade = self._predict_grade_with_model(sentence)
             score = score_map.get(grade, 3)
-            # 附带调试信息
-            scored_sentences.append({
-                "sentence": sentence,
-                "grade": grade,
-                "source": info.get("source", "unknown"),
-                "prob": info.get("prob"),
-                "word_score_total": info.get("word_score_total"),
-            })
+            scored_sentences.append({"sentence": sentence, "grade": grade})
             total_sentence_score += score
 
         comprehensive_grade = "C"
@@ -404,19 +346,4 @@ class SentenceExtractor:
             "neutral_word_count": word_scores["neutral_count"],
             "scored_sentences": scored_sentences,
             "count": len(relevant_sentences),
-            # 调试字段
-            "debug": {
-                "model_loaded": getattr(self, "model_loaded", False),
-                "model_path_abs": getattr(self, "model_path_abs", None),
-                "model_sha256": getattr(self, "model_sha256", None),
-                "aggregation_mode": self.aggregation_mode,
-                "min_sentence_char_len": self.min_sentence_char_len,
-                "merge_leading_punct": self.merge_leading_punct,
-                "word_score_plus_threshold": self.word_score_plus_threshold,
-                "word_score_minus_threshold": self.word_score_minus_threshold,
-                "relevant_sentences": relevant_sentences,
-                "word_score_total": word_scores["total_score"],
-            }
-        }
-
-
+        }