Update handler.py

handler.py

@@ -7,18 +7,18 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 class EndpointHandler:
-    def __init__(self):
+    def __init__(self, model_dir: str = None):  # 修改這裡，添加類型提示和默認值
+        logger.info(f"初始化 EndpointHandler，model_dir: {model_dir}")
+        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}")
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]:  # 修改返回類型
-        """處理輸入並返回符合要求的格式"""
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]:
         try:
             inputs = self.preprocess(data)
             outputs = self.inference(inputs)
-            # 確保返回值是列表格式
             return [outputs]
         except Exception as e:
             logger.error(f"處理過程錯誤: {str(e)}")
@@ -29,12 +29,12 @@ class EndpointHandler:
         logger.info("開始初始化模型")
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(
-                "homer7676/FrierenChatbotV1",
+                self.model_dir,  # 使用 model_dir
                 trust_remote_code=True
             )
             
             self.model = AutoModelForCausalLM.from_pretrained(
-                "homer7676/FrierenChatbotV1",
+                self.model_dir,  # 使用 model_dir
                 trust_remote_code=True,
                 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
             ).to(self.device)