Update handler.py

handler.py

@@ -1,45 +1,58 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from typing import Dict, Any
+import logging
+
+# 設置日誌
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class EndpointHandler:
-    def __init__(self):
+    def __init__(self, model_dir=None):
+        logger.info("初始化 EndpointHandler")
         self.tokenizer = None
         self.model = None
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-
+        logger.info(f"使用設備: {self.device}")
+        
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """使 handler 可調用"""
-        inputs = self.preprocess(data)
-        outputs = self.inference(inputs)
-        return self.postprocess(outputs)
+        logger.info("調用 __call__ 方法")
+        return self.inference(self.preprocess(data))
 
     def initialize(self, context):
         """初始化模型和 tokenizer"""
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            "homer7676/FrierenChatbotV1",
-            trust_remote_code=True
-        )
-        self.model = AutoModelForCausalLM.from_pretrained(
-            "homer7676/FrierenChatbotV1",
-            trust_remote_code=True,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-        ).to(self.device)
-        self.model.eval()
-
-    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
+        logger.info("開始初始化模型")
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                "homer7676/FrierenChatbotV1",
+                trust_remote_code=True
+            )
+            logger.info("Tokenizer 載入成功")
+            
+            self.model = AutoModelForCausalLM.from_pretrained(
+                "homer7676/FrierenChatbotV1",
+                trust_remote_code=True,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            ).to(self.device)
+            logger.info("模型載入成功")
+            
+            self.model.eval()
+            logger.info("模型初始化完成")
+            
+        except Exception as e:
+            logger.error(f"模型載入錯誤: {str(e)}")
+            raise
 
     def inference(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
         """執行推理"""
+        logger.info("開始執行推理")
         try:
             message = inputs.get("message", "")
             context = inputs.get("context", "")
-            prompt = f"""你是芙莉蓮，需要遵守以下規則回答：
+            logger.info(f"收到訊息: {message}")
+            
+            input_text = f"""你是芙莉蓮，需要遵守以下規則回答：
 1. 身份設定：
  - 千年精靈魔法師
  - 態度溫柔但帶著些許嘲諷
@@ -57,14 +70,19 @@ class EndpointHandler:
 用戶：{message}
 芙莉蓮："""
 
+            # 記錄 token 數量
+            tokens = self.tokenizer.encode(input_text)
+            logger.info(f"輸入 token 數量: {len(tokens)}")
+            
             inputs = self.tokenizer(
-                prompt,
+                input_text,
                 return_tensors="pt",
                 padding=True,
                 truncation=True,
                 max_length=2048
             ).to(self.device)
-
+            
+            logger.info("開始生成回應")
             with torch.no_grad():
                 outputs = self.model.generate(
                     **inputs,
@@ -72,21 +90,23 @@ class EndpointHandler:
                     temperature=0.7,
                     top_p=0.9,
                     top_k=50,
-                    do_sample=True,
-                    repetition_penalty=1.2,
-                    pad_token_id=self.tokenizer.pad_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id
+                    do_sample=True
                 )
-
+            
             response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             response = response.split("芙莉蓮：")[-1].strip()
+            logger.info(f"生成回應完成，長度: {len(response)}")
             
             return {"generated_text": response}
             
         except Exception as e:
-            print(f"推理過程錯誤: {str(e)}")
+            logger.error(f"推理過程錯誤: {str(e)}")
             return {"error": str(e)}
 
-    def postprocess(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """後處理輸出數據"""
-        return data
+    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