Update handler.py

handler.py

@@ -1,8 +1,6 @@
 from typing import List, Dict
 import torch
 from transformers import AutoConfig, AutoTokenizer, AutoModelForSequenceClassification
-# 移除相对导入
-# from .modeling import BinaryClassifier 
 
 class EndpointHandler:
     def __init__(self, path=""):
@@ -16,71 +14,88 @@ class EndpointHandler:
         # 初始化tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         
-        # 设置最大长度，可以根据你的需求调整
+        # 设置最大长度
         self.max_length = 512
 
     def __call__(self, data: List[Dict[str, str]]) -> List[Dict[str, float]]:
         """
         处理文本推理请求
-        Args:
-            data: 输入数据列表，每个元素是一个字典
-                 例如：[{"inputs": "这是一段测试文本"}]
-        Returns:
-            预测结果列表
         """
-        # 获取所有输入文本
-        texts = [item["inputs"] for item in data]
-        
-        # tokenization
-        encoded_inputs = self.tokenizer(
-            texts,
-            padding=True,
-            truncation=True,
-            max_length=self.max_length,
-            return_tensors="pt"
-        )
-        
-        # 进行预测
-        with torch.no_grad():
-            outputs = self.model(**encoded_inputs)
-            logits = outputs.logits
-            predictions = torch.softmax(logits, dim=-1)
-        
-        # 格式化输出
-        results = []
-        for pred in predictions:
-            label_id = pred.argmax().item()
-            score = pred[label_id].item()
-            results.append({
-                "label": str(label_id),    # 0 或 1
-                "score": float(score)      # 预测概率
-            })
-        
-        return results
+        try:
+            # 获取所有输入文本
+            texts = []
+            for item in data:
+                # 确保我们正确处理输入数据
+                if isinstance(item, dict) and "inputs" in item:
+                    texts.append(item["inputs"])
+                elif isinstance(item, str):
+                    texts.append(item)
+                else:
+                    raise ValueError(f"Unexpected input format: {item}")
+            
+            # tokenization
+            encoded_inputs = self.tokenizer(
+                texts,
+                padding=True,
+                truncation=True,
+                max_length=self.max_length,
+                return_tensors="pt"
+            )
+            
+            # 进行预测
+            with torch.no_grad():
+                outputs = self.model(**encoded_inputs)
+                logits = outputs.logits
+                probabilities = torch.softmax(logits, dim=-1)
+            
+            # 格式化输出
+            results = []
+            for probs in probabilities:
+                label_id = int(torch.argmax(probs).item())
+                confidence = float(probs[label_id].item())
+                results.append({
+                    "label": str(label_id),    # 转换为字符串
+                    "score": confidence        # 预测概率
+                })
+            
+            return results
+            
+        except Exception as e:
+            # 添加错误处理和日志记录
+            print(f"Error in prediction: {str(e)}")
+            return [{"error": str(e)}]
 
     def preprocess(self, text: str) -> Dict[str, torch.Tensor]:
         """
         预处理方法
         """
-        encoded = self.tokenizer(
-            text,
-            padding=True,
-            truncation=True,
-            max_length=self.max_length,
-            return_tensors="pt"
-        )
-        return encoded
+        try:
+            encoded = self.tokenizer(
+                text,
+                padding=True,
+                truncation=True,
+                max_length=self.max_length,
+                return_tensors="pt"
+            )
+            return encoded
+        except Exception as e:
+            print(f"Error in preprocessing: {str(e)}")
+            raise e
 
     def postprocess(self, model_outputs) -> Dict:
         """
         后处理方法
         """
-        logits = model_outputs.logits
-        predictions = torch.softmax(logits, dim=-1)
-        label_id = predictions[0].argmax().item()
-        score = predictions[0][label_id].item()
-        
-        return {
-            "label": str(label_id),
-            "score": float(score)
-        }
+        try:
+            logits = model_outputs.logits
+            probabilities = torch.softmax(logits, dim=-1)
+            label_id = int(torch.argmax(probabilities[0]).item())
+            confidence = float(probabilities[0][label_id].item())
+            
+            return {
+                "label": str(label_id),
+                "score": confidence
+            }
+        except Exception as e:
+            print(f"Error in postprocessing: {str(e)}")
+            raise e