Update predictor.py

predictor.py

@@ -1,87 +1,115 @@
+# predictor.py
+
 import json
-import re
-from typing import List, Dict, Tuple
 import numpy as np
-from onnxruntime import InferenceSession
+import onnxruntime as ort
+from transformers import BertTokenizer
+import re
+
+class Predictor:
+    def __init__(self):
+        """
+        在服务启动时，一次性加载所有必要的模型和文件。
+        """
+        # 1. 加载分词器 (Tokenizer)
+        # Hugging Face Spaces会自动下载git仓库中的所有文件到当前目录
+        self.tokenizer = BertTokenizer.from_pretrained('.')
 
-# 加载ONNX模型
-MODEL_SESSION = InferenceSession('model_quantized.onnx')
+        # 2. 加载ONNX模型并创建推理会话
+        self.ort_session = ort.InferenceSession('model_quantized.onnx')
 
-# 评分等级映射
-GRADE_MAPPING = {
-    0: 'A',
-    1: 'B',
-    2: 'C', 
-    3: 'D',
-    4: 'E'
-}
+        # 3. 加载关键词词集
+        with open('evaluation_keywords2.json', 'r', encoding='utf-8') as f:
+            self.keywords = json.load(f)
+        
+        # 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 analyze_text_with_keywords(text: str, keywords_data: Dict) -> List[str]:
-    """
-    使用关键词库分析文本，提取相关句子
-    """
-    # 按句子分割（简单的中文分句）
-    sentences = re.split(r'[。！？；]', text)
-    sentences = [s.strip() for s in sentences if len(s.strip()) > 5]  # 过滤短句
-    
-    relevant_sentences = []
-    
-    # 检查每个句子是否包含关键词
-    for sentence in sentences:
-        for category in keywords_data.values():
-            for keyword in category:
+    def _extract_relevant_sentences(self, text):
+        """
+        根据关键词提取相关的句子。
+        """
+        # 使用正则表达式按标点符号分割句子，更准确
+        sentences = re.split(r'[。！？]', 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 preprocess_text(sentence: str) -> np.ndarray:
-    """
-    文本预处理（根据你的模型需要修改）
-    """
-    # 这里应该添加你的tokenizer逻辑
-    # 示例: 返回随机向量 (实际使用时替换为真实预处理)
-    return np.random.rand(1, 768).astype(np.float32)
+                    break # 找到一个关键词就添加，避免重复
+        return relevant_sentences
 
-def predict_grade(sentences: List[str]) -> Tuple[str, List[str]]:
-    """
-    预测句子等级并计算平均分
-    """
-    if not sentences:
-        return "N/A", []
-    
-    sentence_grades = []
-    grade_scores = []
-    
-    for sentence in sentences:
-        # 1. 预处理
-        inputs = preprocess_text(sentence)
+    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']}
+        
+        # 执行推理
+        ort_outs = self.ort_session.run(None, ort_inputs)
         
-        # 2. 运行ONNX模型
-        outputs = MODEL_SESSION.run(None, {'input': inputs})
-        pred_class = np.argmax(outputs[0])
+        # 处理输出结果
+        prediction = np.argmax(ort_outs[0], axis=1)[0]
+        return self.id2label[prediction]
+
+    def predict(self, text):
+        """
+        执行完整的预测流程：提取句子 -> 逐句评分 -> 计算平均等级。
+        这是暴露给app.py调用的主方法。
+        """
+        # 步骤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)
         
-        # 3. 记录结果
-        grade = GRADE_MAPPING.get(pred_class, 'E')
-        sentence_grades.append(grade)
-        grade_scores.append(pred_class)
-    
-    # 4. 计算平均分
-    avg_score = np.mean(grade_scores)
-    final_grade = GRADE_MAPPING.get(int(round(avg_score)), 'E')
-    
-    return final_grade, sentence_grades
+        # 将平均分四舍五入后映射回最终等级
+        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)
+        }
+
 
-# 测试用
-if __name__ == '__main__':
-    test_text = "这个学生表现很好。创新性不足。逻辑清晰。完成度一般。"
-    with open('evaluation_keywords2.json', 'r', encoding='utf-8') as f:
-        keywords = json.load(f)
-    
-    relevant = analyze_text_with_keywords(test_text, keywords)
-    print(f"匹配的句子: {relevant}")
-    
-    grade, details = predict_grade(relevant)
-    print(f"最终等级: {grade}")
-    print(f"详细评分: {details}")