Update predictor.py

predictor.py

@@ -32,13 +32,28 @@ class SentenceExtractor:
 
     def _preprocess_text(self, text: str) -> np.ndarray:
         """
-        预处理文本为模型输入格式
+        预处理文本为模型输入格式 - 使用BERT tokenizer
         """
-        max_seq_length = 128
-        features = np.zeros((1, max_seq_length), dtype=np.float32)
-        for i, ch in enumerate(text[:max_seq_length]):
-            features[0, i] = (ord(ch) % 256) / 255.0
-        return features
+        try:
+            from transformers import AutoTokenizer
+            # 使用与学生模型相同的tokenizer
+            tokenizer = AutoTokenizer.from_pretrained("uer/chinese_roberta_L-4_H-256")
+            inputs = tokenizer(
+                text,
+                truncation=True,
+                padding=True,
+                max_length=512,
+                return_tensors='np'
+            )
+            return inputs
+        except Exception as e:
+            print(f"Tokenizer预处理失败: {e}")
+            # 降级到简单字符编码
+            max_seq_length = 128
+            features = np.zeros((1, max_seq_length), dtype=np.float32)
+            for i, ch in enumerate(text[:max_seq_length]):
+                features[0, i] = (ord(ch) % 256) / 255.0
+            return features
     
     def _predict_grade_with_model(self, text: str) -> str:
         """
@@ -47,8 +62,20 @@ class SentenceExtractor:
         try:
             if not self.ort_session:
                 return "C"
-            input_data = self._preprocess_text(text)
-            outputs = self.ort_session.run([self.output_name], {self.input_name: input_data})
+            inputs = self._preprocess_text(text)
+            
+            # 检查是否是tokenizer输出格式
+            if isinstance(inputs, dict) and 'input_ids' in inputs:
+                # BERT tokenizer格式
+                input_data = {
+                    'input_ids': inputs['input_ids'],
+                    'attention_mask': inputs['attention_mask']
+                }
+            else:
+                # 简单字符编码格式
+                input_data = {self.input_name: inputs}
+            
+            outputs = self.ort_session.run([self.output_name], input_data)
             predictions = outputs[0]
             grade_index = int(np.argmax(predictions))
             grades = ['A', 'B', 'C', 'D', 'E']
@@ -261,4 +288,5 @@ if __name__ == "__main__":
     
     for i, item in enumerate(result['scored_sentences'], 1):
         print(f"句子{i}加评分等级：{item['sentence']} - {item['grade']}")
+