Update output.py

output.py

@@ -288,12 +288,22 @@ class AphasiaInferenceSystem:
 
         self.model = StableAphasiaClassifier(self.config, self.num_labels)
         self.model.bert.resize_token_embeddings(len(self.tokenizer))
-
-        model_path = os.path.join(self.model_dir, "pytorch_model.bin")
-        if not os.path.exists(model_path):
-            raise FileNotFoundError(f"模型權重文件不存在: {model_path}")
-        state = torch.load(model_path, map_location=self.device)
-        self.model.load_state_dict(state)  # (once)
+    
+        sft_path = os.path.join(self.model_dir, "model.safetensors")
+        bin_path = os.path.join(self.model_dir, "pytorch_model.bin")
+        
+        if os.path.exists(sft_path):
+            # safetensors 不使用 pickle，較安全
+            state = load_file(sft_path)  # 讀成 state_dict（CPU）
+            missing, unexpected = self.model.load_state_dict(state, strict=False)
+        elif os.path.exists(bin_path):
+            state = torch.load(bin_path, map_location="cpu")  # 先載到 CPU，再搬到裝置
+            missing, unexpected = self.model.load_state_dict(state, strict=False)
+        else:
+            raise FileNotFoundError(f"找不到模型權重：{sft_path} 或 {bin_path}")
+        
+        if missing or unexpected:
+            print(f"[load_state_dict] missing={missing}, unexpected={unexpected}")
 
         self.model.to(self.device)
         self.model.eval()