Update handler.py

handler.py

@@ -3,90 +3,99 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from typing import Dict, Any
 import re
 
-SIMPLIFIED_TO_TRADITIONAL = {
-    '发': '發', '书': '書', '记': '記', '亚': '亞', '欧': '歐', '韩': '韓', '边': '邊',
-    '恒': '恆', '说': '說', '话': '話', '东': '東', '车': '車', '马': '馬', '样': '樣',
-    '风': '風', '专': '專', '万': '萬', '劳': '勞', '动': '動', '习': '習', '头': '頭',
-    '们': '們', '为': '為', '产': '產', '场': '場', '实': '實', '观': '觀', '见': '見',
-    '师': '師', '长': '長', '识': '識', '电': '電', '图': '圖', '华': '華', '龙': '龍',
-    '变': '變', '问': '問', '岁': '歲', '义': '義', '还': '還', '报': '報', '乐': '樂',
-    '欢': '歡', '权': '權', '态': '態', '极': '極', '环': '環', '带': '帶', '难': '難'
-}
-
 class EndpointHandler:
-    def __init__(self, model_dir=None):
-        self.tokenizer = None
-        self.model = None
+    def __init__(self, model_dir: str = None):
+        self.model_dir = model_dir
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model_dir = model_dir if model_dir else "homer7676/FrierenChatbotV1"
+        self.model = None
+        self.tokenizer = None
 
-    def initialize(self, context):
-        try:
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                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,
-                trust_remote_code=True,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-            ).to(self.device)
-            
-            self.model.eval()
+    def initialize(self, context: Dict[str, Any] = None):
+        """Initialize the model and tokenizer."""
+        model_id = "homer7676/FrierenChatbotV1"
+        
+        # Initialize tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_id,
+            trust_remote_code=True,
+            padding_side="left"
+        )
+        
+        # Ensure pad token exists
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
             
-        except Exception as e:
-            print(f"模型載入錯誤: {str(e)}")
-            raise
+        # Initialize model
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            trust_remote_code=True,
+            torch_dtype="auto",
+            low_cpu_mem_usage=True
+        ).to(self.device)
+        
+        self.model.eval()
+        return self
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Main prediction pipeline."""
+        inputs = self.preprocess(data)
+        outputs = self.inference(inputs)
+        return self.postprocess(outputs)
 
     def preprocess(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Preprocess the input data."""
+        if isinstance(data, str):
+            return {"message": data}
         inputs = data.pop("inputs", data)
-        if not isinstance(inputs, dict):
-            inputs = {"message": inputs}
-        return inputs
+        return inputs if isinstance(inputs, dict) else {"message": inputs}
 
     def inference(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        """Run the inference."""
         try:
+            # 準備輸入
             message = inputs.get("message", "")
             context = inputs.get("context", "")
             prompt = self._build_prompt(context, message)
-            
-            encoding = self.tokenizer(
+
+            # Tokenize
+            inputs = self.tokenizer(
                 prompt,
                 return_tensors="pt",
-                add_special_tokens=True,
+                padding=True,
                 truncation=True,
-                max_length=2048,
-                padding=True
+                max_length=2048
             ).to(self.device)
-            
+
+            # Generate
             with torch.no_grad():
-                outputs = self.model.generate(
-                    input_ids=encoding["input_ids"],
-                    attention_mask=encoding["attention_mask"],
+                generation_output = self.model.generate(
+                    input_ids=inputs["input_ids"],
+                    attention_mask=inputs["attention_mask"],
                     max_new_tokens=256,
                     temperature=0.7,
                     top_p=0.9,
                     top_k=50,
                     do_sample=True,
-                    repetition_penalty=1.2,
-                    num_beams=4,
-                    early_stopping=True
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    repetition_penalty=1.2
                 )
+
+            response = self.tokenizer.decode(
+                generation_output[0], 
+                skip_special_tokens=True
+            )
             
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # 處理回應
             response = response.split("芙莉蓮：")[-1].strip()
             response = self._process_response(response)
-            return {"response": response}
             
+            return {"response": response}
         except Exception as e:
-            print(f"推理過程錯誤: {str(e)}")
-            return {"response": "抱歉，在處理您的請求時發生了錯誤。請稍後再試。", "error": str(e)}
+            return {"error": f"Inference error: {str(e)}"}
 
     def _build_prompt(self, context: str, query: str) -> str:
+        """Build the prompt for the model."""
         return f"""你是芙莉蓮，需要遵守以下規則回答：
 1. 身份設定：
  - 千年精靈魔法師
@@ -105,22 +114,24 @@ class EndpointHandler:
 用戶：{query}
 芙莉蓮："""
 
-    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:
+        """Process the model's response."""
         if not response or not response.strip():
             return "抱歉，我現在有點恍神，請你再問一次好嗎？"
         
-        response = self._convert_to_traditional(response)
+        # Convert to traditional Chinese
+        for simplified, traditional in SIMPLIFIED_TO_TRADITIONAL.items():
+            response = response.replace(simplified, traditional)
+        
+        # Clean up whitespace
         response = re.sub(r'\s+', '', response)
         
+        # Add ending punctuation if needed
         if not response.endswith(('。', '！', '？', '~', '呢', '啊', '吶')):
             response += '呢。'
             
         return response
 
     def postprocess(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Postprocess the output data."""
         return data