Update handler.py

handler.py

@@ -3,16 +3,14 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from typing import Dict, Any
 import re
 
-# 簡繁轉換字典
 SIMPLIFIED_TO_TRADITIONAL = {
     '发': '發', '书': '書', '记': '記', '亚': '亞', '欧': '歐', '韩': '韓', '边': '邊',
     '恒': '恆', '说': '說', '话': '話', '东': '東', '车': '車', '马': '馬', '样': '樣',
     '风': '風', '专': '專', '万': '萬', '劳': '勞', '动': '動', '习': '習', '头': '頭',
     '们': '們', '为': '為', '产': '產', '场': '場', '实': '實', '观': '觀', '见': '見',
     '师': '師', '长': '長', '识': '識', '电': '電', '图': '圖', '华': '華', '龙': '龍',
-    '师': '師', '变': '變', '问': '問', '岁': '歲', '义': '義', '记': '記', '还': '還',
-    '报': '報', '乐': '樂', '欢': '歡', '权': '權', '态': '態', '极': '極', '环': '環',
-    '带': '帶', '难': '難'
+    '变': '變', '问': '問', '岁': '歲', '义': '義', '还': '還', '报': '報', '乐': '樂',
+    '欢': '歡', '权': '權', '态': '態', '极': '極', '环': '環', '带': '帶', '难': '難'
 }
 
 class EndpointHandler:
@@ -23,7 +21,6 @@ class EndpointHandler:
         self.model_dir = model_dir if model_dir else "homer7676/FrierenChatbotV1"
 
     def initialize(self, context):
-        """初始化模型和 tokenizer"""
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(
                 self.model_dir,
@@ -36,9 +33,8 @@ class EndpointHandler:
             self.model = AutoModelForCausalLM.from_pretrained(
                 self.model_dir,
                 trust_remote_code=True,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                device_map="auto"
-            )
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            ).to(self.device)
             
             self.model.eval()
             
@@ -47,14 +43,12 @@ class EndpointHandler:
             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, Any]:
-        """執行推理"""
         try:
             message = inputs.get("message", "")
             context = inputs.get("context", "")
@@ -67,11 +61,7 @@ class EndpointHandler:
                 truncation=True,
                 max_length=2048,
                 padding=True
-            )
-            
-            for key in encoding:
-                if isinstance(encoding[key], torch.Tensor):
-                    encoding[key] = encoding[key].to(self.device)
+            ).to(self.device)
             
             with torch.no_grad():
                 outputs = self.model.generate(
@@ -83,8 +73,6 @@ class EndpointHandler:
                     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,
                     num_beams=4,
                     early_stopping=True
                 )
@@ -92,7 +80,6 @@ class EndpointHandler:
             response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             response = response.split("芙莉蓮：")[-1].strip()
             response = self._process_response(response)
-            
             return {"response": response}
             
         except Exception as e:
@@ -100,7 +87,6 @@ class EndpointHandler:
             return {"response": "抱歉，在處理您的請求時發生了錯誤。請稍後再試。", "error": str(e)}
 
     def _build_prompt(self, context: str, query: str) -> str:
-        """構建提示詞"""
         return f"""你是芙莉蓮，需要遵守以下規則回答：
 1. 身份設定：
  - 千年精靈魔法師
@@ -120,13 +106,11 @@ class EndpointHandler:
 芙莉蓮："""
 
     def _convert_to_traditional(self, text: str) -> str:
-        """將簡體轉換為繁體"""
         for simplified, traditional in SIMPLIFIED_TO_TRADITIONAL.items():
             text = text.replace(simplified, traditional)
         return text
 
     def _process_response(self, response: str) -> str:
-        """處理回應文本"""
         if not response or not response.strip():
             return "抱歉，我現在有點恍神，請你再問一次好嗎？"
         
@@ -139,5 +123,4 @@ class EndpointHandler:
         return response
 
     def postprocess(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """後處理輸出數據"""
         return data