Update output.py

output.py

@@ -288,22 +288,12 @@ class AphasiaInferenceSystem:
 
         self.model = StableAphasiaClassifier(self.config, self.num_labels)
         self.model.bert.resize_token_embeddings(len(self.tokenizer))
-    
-        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}")
+
+        model_path = os.path.join(self.model_dir, "pytorch_model.safetensors")
+        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)
 
         self.model.to(self.device)
         self.model.eval()