Update predictor.py

predictor.py

@@ -1,142 +1,113 @@
+# predictor.py
+
 import json
+import numpy as np
+import onnxruntime as ort
+from transformers import BertTokenizer
 import re
-from typing import List, Dict, Set, Tuple
 
-class SentenceExtractor:
-    def __init__(self, main_keywords_path: str, eval_keywords_path: str):
+class Predictor:
+    def __init__(self):
         """
-        初始化句子提取器，加载主关键词和评估关键词库
-        :param main_keywords_path: 主关键词JSON文件路径
-        :param eval_keywords_path: 评估关键词库(JSON)文件路径
+        在服务启动时，一次性加载所有必要的模型和文件。
         """
-        self.main_keywords = self._load_keywords(main_keywords_path)
-        self.eval_keywords = self._load_eval_keywords(eval_keywords_path)
-        # 合并所有关键词用于快速查找
-        self.all_keywords = self._merge_all_keywords()
-        
-    def _load_keywords(self, file_path: str) -> Dict[str, List[str]]:
-        """加载主关键词文件"""
-        try:
-            with open(file_path, 'r', encoding='utf-8') as f:
-                return json.load(f)
-        except Exception as e:
-            print(f"加载主关键词文件失败: {e}")
-            return {}
-    
-    def _load_eval_keywords(self, file_path: str) -> Dict[str, Dict[str, List[str]]]:
-        """加载评估关键词库(evaluation_keywords2.json)"""
-        try:
-            with open(file_path, 'r', encoding='utf-8') as f:
-                return json.load(f)
-        except Exception as e:
-            print(f"加载评估关键词库失败: {e}")
-            return {}
-    
-    def _merge_all_keywords(self) -> Set[str]:
-        """合并所有关键词到一个集合中，用于快速查找"""
-        keywords_set = set()
+        # 1. 加载分词器 (Tokenizer)
+        # Hugging Face Spaces会自动下载git仓库中的所有文件到当前目录
+        self.tokenizer = BertTokenizer.from_pretrained('.')
+
+        # 2. 加载ONNX模型并创建推理会话
+        self.ort_session = ort.InferenceSession('model_quantized.onnx')
+
+        # 3. 加载关键词词集
+        with open('evaluation_keywords2.json', 'r', encoding='utf-8') as f:
+            self.keywords = json.load(f)
         
-        # 添加主关键词
-        for category, keywords in self.main_keywords.items():
-            keywords_set.update(keywords)
-            
-        # 添加评估关键词库中的所有关键词
-        for category, types in self.eval_keywords.items():
-            for type_, keywords in types.items():
-                keywords_set.update(keywords)
-                
-        return keywords_set
-    
-    def _split_into_sentences(self, text: str) -> List[str]:
-        """将文本分割成句子"""
-        # 简单的句子分割正则，可根据需要优化
-        sentence_endings = re.compile(r'(?<=[。！？,.!?])\s+')
-        sentences = sentence_endings.split(text)
-        return [s.strip() for s in sentences if s.strip()]
-    
-    def _extract_relevant_sentences(self, text: str) -> Tuple[List[str], Dict[str, List[str]]]:
+        # 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 _extract_relevant_sentences(self, text):
         """
-        提取与关键词相关的句子
-        :param text: 输入文本
-        :return: 相关句子列表和按类别分组的句子字典
+        根据关键词提取相关的句子。
         """
-        sentences = self._split_into_sentences(text)
+        # 使用正则表达式按标点符号分割句子，更准确
+        sentences = re.split(r'[。！？]', text)
         relevant_sentences = []
-        categorized_sentences = {
-            "main": [],
-            "student_performance": {"positive": [], "negative": [], "nature": [], "suggestion": []},
-            "content_quality": {"positive": [], "negative": [], "nature": [], "suggestion": []},
-            "cross_scene": {"positive": [], "negative": [], "nature": [], "suggestion": []}
-        }
-        
         for sentence in sentences:
-            # 检查是否包含主关键词
-            main_keyword_matched = False
-            for category, keywords in self.main_keywords.items():
-                for keyword in keywords:
-                    if keyword in sentence:
-                        relevant_sentences.append(sentence)
-                        categorized_sentences["main"].append(sentence)
-                        main_keyword_matched = True
-                        break
-                if main_keyword_matched:
-                    break
-            
-            # 检查评估关键词库中的关键词
-            for category in ["student_performance", "content_quality", "cross_scene"]:
-                if category not in self.eval_keywords:
-                    continue
-                    
-                for sentiment in ["positive", "negative", "nature", "suggestion"]:
-                    if sentiment not in self.eval_keywords[category]:
-                        continue
-                        
-                    for keyword in self.eval_keywords[category][sentiment]:
-                        if keyword in sentence and sentence not in categorized_sentences[category][sentiment]:
-                            # 如果还没添加到相关句子列表，则添加
-                            if sentence not in relevant_sentences:
-                                relevant_sentences.append(sentence)
-                            categorized_sentences[category][sentiment].append(sentence)
+            if not sentence:
+                continue
+            for keyword in self.keywords:
+                if keyword in sentence:
+                    relevant_sentences.append(sentence)
+                    break # 找到一个关键词就添加，避免重复
+        return relevant_sentences
+
+    def _predict_single_sentence(self, sentence):
+        """
+        对单个句子进行模型推理，返回预测的等级标签。
+        """
+        # 使用分词器处理文本
+        inputs = self.tokenizer(sentence, return_tensors="np", padding='max_length', truncation=True, max_length=128)
         
-        return relevant_sentences, categorized_sentences
-    
-    def extract(self, text: str) -> Dict[str, any]:
+        # 准备ONNX模型的输入
+        ort_inputs = {self.ort_session.get_inputs()[0].name: inputs['input_ids']}
+        
+        # 执行推理
+        ort_outs = self.ort_session.run(None, ort_inputs)
+        
+        # 处理输出结果
+        prediction = np.argmax(ort_outs[0], axis=1)[0]
+        return self.id2label[prediction]
+
+    def predict(self, text):
         """
-        提取文本中与关键词相关的句子
-        :param text: 输入文本
-        :return: 包含相关句子和分类信息的字典
+        执行完整的预测流程：提取句子 -> 逐句评分 -> 计算平均等级。
+        这是暴露给app.py调用的主方法。
         """
-        if not text:
-            return {"relevant_sentences": [], "categorized_sentences": {}}
-            
-        relevant_sentences, categorized_sentences = self._extract_relevant_sentences(text)
+        # 步骤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)
         
-        return {
-            "relevant_sentences": relevant_sentences,
-            "categorized_sentences": categorized_sentences,
-            "count": len(relevant_sentences)
-        }
+        # 将平均分四舍五入后映射回最终等级
+        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"
 
-# 使用示例
-if __name__ == "__main__":
-    # 假设主关键词文件名为main_keywords.json
-    extractor = SentenceExtractor(
-        main_keywords_path="main_keywords.json",
-        eval_keywords_path="evaluation_keywords2.json"
-    )
-    
-    sample_text = """
-    该学生表现优异，团队合作能力强，在项目中展现了很强的创新能力。
-    但代码质量不高，存在安全漏洞，需要加强测试。
-    项目文档完整，符合行业标准，具有很好的应用价值。
-    建议加强代码审查，提高系统安全性，优化算法效率。
-    """
-    
-    result = extractor.extract(sample_text)
-    print(f"提取到 {result['count']} 个相关句子：")
-    for i, sent in enumerate(result['relevant_sentences'], 1):
-        print(f"{i}. {sent}")
-    
-    print("\n按类别分组：")
-    print(json.dumps(result['categorized_sentences'], ensure_ascii=False, indent=2))
+        # 步骤4: 生成总结性文本
+        summary = f"系统分析了 {len(relevant_sentences)} 个关键句子，综合评定等级为“{final_grade.upper()}”。"
 
+        return {
+            "grade": final_grade,
+            "summary": summary,
+            "analyzed_sentences_count": len(relevant_sentences)
+        }