Update handler.py

handler.py

@@ -7,61 +7,56 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 class EndpointHandler:
-    def __init__(self, model_dir: str = None):  # 修改這裡，添加類型提示和默認值
-        logger.info(f"初始化 EndpointHandler，model_dir: {model_dir}")
+    def __init__(self, model_dir: str = None):
         self.model_dir = model_dir if model_dir else "homer7676/FrierenChatbotV1"
         self.tokenizer = None
         self.model = None
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info(f"使用設備: {self.device}")
+        logger.info(f"初始化 EndpointHandler，使用設備: {self.device}")
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]:
         try:
             inputs = self.preprocess(data)
             outputs = self.inference(inputs)
+            # 確保輸出不為空
+            if not outputs or "generated_text" not in outputs:
+                raise ValueError("No text was generated")
             return [outputs]
         except Exception as e:
             logger.error(f"處理過程錯誤: {str(e)}")
             return [{"error": str(e)}]
 
     def initialize(self, context):
-        """初始化模型和 tokenizer"""
         logger.info("開始初始化模型")
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(
-                self.model_dir,  # 使用 model_dir
+                self.model_dir,
                 trust_remote_code=True
             )
             
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            
             self.model = AutoModelForCausalLM.from_pretrained(
-                self.model_dir,  # 使用 model_dir
+                self.model_dir,
                 trust_remote_code=True,
                 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
             ).to(self.device)
             
-            if self.tokenizer.pad_token is None:
-                self.tokenizer.pad_token = self.tokenizer.eos_token
-                
             self.model.eval()
             logger.info("模型初始化完成")
         except Exception as e:
             logger.error(f"模型載入錯誤: {str(e)}")
             raise
 
-    def preprocess(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """預處理輸入數據"""
-        inputs = data.pop("inputs", data)
-        if not isinstance(inputs, dict):
-            inputs = {"message": inputs}
-        return inputs
-
     def inference(self, inputs: Dict[str, Any]) -> Dict[str, str]:
-        """執行推理"""
         logger.info("開始執行推理")
         try:
             message = inputs.get("message", "")
             context = inputs.get("context", "")
+            logger.info(f"處理訊息: {message}")
             
+            # 構建提示詞
             prompt = f"""你是芙莉蓮，需要遵守以下規則回答：
 1. 身份設定：
  - 千年精靈魔法師
@@ -80,29 +75,57 @@ class EndpointHandler:
 用戶：{message}
 芙莉蓮："""
 
-            inputs = self.tokenizer(
+            # 記錄提示詞長度
+            logger.info(f"提示詞長度: {len(prompt)}")
+            
+            # Tokenize
+            encoding = self.tokenizer.encode_plus(
                 prompt,
+                add_special_tokens=True,
                 return_tensors="pt",
                 padding=True,
                 truncation=True,
                 max_length=2048
-            ).to(self.device)
+            )
+            
+            # 移動到正確的設備
+            input_ids = encoding["input_ids"].to(self.device)
+            attention_mask = encoding["attention_mask"].to(self.device)
+            
+            logger.info(f"輸入 token 數量: {input_ids.shape[-1]}")
 
+            # 生成回應
             with torch.no_grad():
                 outputs = self.model.generate(
-                    **inputs,
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
                     max_new_tokens=256,
                     temperature=0.7,
                     top_p=0.9,
                     top_k=50,
                     do_sample=True,
                     pad_token_id=self.tokenizer.pad_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    num_return_sequences=1,
+                    no_repeat_ngram_size=3
                 )
+                
+                logger.info(f"生成的 token 數量: {outputs.shape[-1]}")
 
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            response = response.split("芙莉蓮：")[-1].strip()
-            logger.info("生成回應完成")
+            # 解碼回應
+            full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # 分離出模型的回應部分
+            if "芙莉蓮：" in full_response:
+                response = full_response.split("芙莉蓮：")[-1].strip()
+            else:
+                response = full_response.split("用戶：")[-1].strip()
+            
+            logger.info(f"生成回應長度: {len(response)}")
+            
+            # 確保回應不為空
+            if not response:
+                response = "抱歉，我似乎有點恍神了。能請你再說一次嗎？"
             
             return {
                 "generated_text": response
@@ -110,4 +133,11 @@ class EndpointHandler:
 
         except Exception as e:
             logger.error(f"推理過程錯誤: {str(e)}")
-            return {"error": str(e)}
+            raise
+
+    def preprocess(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        logger.info(f"預處理輸入數據: {data}")
+        inputs = data.pop("inputs", data)
+        if not isinstance(inputs, dict):
+            inputs = {"message": inputs}
+        return inputs