Update models.py

models.py

@@ -1,11 +1,10 @@
+import os
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, PretrainedConfig, AutoConfig
 from transformers.modeling_outputs import SequenceClassifierOutput
-from huggingface_hub import hf_hub_download
-import os
 
-# Hard-coded configuration constants
+# ===================== 하드코딩된 설정 ========================
 DIM_ECG = 46
 DIM_RR = 62
 DIM_EDA = 24
@@ -22,24 +21,16 @@ class TeacherNet(nn.Module):
         hidden = TEACHER_FEAT_DIM
         dropout = 0.4
         self.ecg_encoder = nn.Sequential(
-            nn.Linear(DIM_ECG, hidden),
-            nn.ReLU(inplace=True),
-            nn.Dropout(p=dropout)
+            nn.Linear(DIM_ECG, hidden), nn.ReLU(inplace=True), nn.Dropout(p=dropout)
         )
         self.rr_encoder = nn.Sequential(
-            nn.Linear(DIM_RR, hidden),
-            nn.ReLU(inplace=True),
-            nn.Dropout(p=dropout)
+            nn.Linear(DIM_RR, hidden), nn.ReLU(inplace=True), nn.Dropout(p=dropout)
         )
         self.eda_encoder = nn.Sequential(
-            nn.Linear(DIM_EDA, hidden),
-            nn.ReLU(inplace=True),
-            nn.Dropout(p=dropout)
+            nn.Linear(DIM_EDA, hidden), nn.ReLU(inplace=True), nn.Dropout(p=dropout)
         )
         self.video_encoder = nn.Sequential(
-            nn.Linear(DIM_VIDEO, hidden),
-            nn.ReLU(inplace=True),
-            nn.Dropout(p=dropout)
+            nn.Linear(DIM_VIDEO, hidden), nn.ReLU(inplace=True), nn.Dropout(p=dropout)
         )
         self.classifier = nn.Sequential(
             nn.Linear(4 * hidden, hidden),
@@ -87,6 +78,8 @@ class StudentNet(nn.Module):
         logits = self.classifier(feat)
         return logits, feat
 
+# ==== Transformers Compatibility ==== #
+
 class StressConfig(PretrainedConfig):
     model_type = "audio-classification"
     def __init__(self, **kwargs):
@@ -114,7 +107,6 @@ class StudentForAudioClassification(PreTrainedModel):
         trust_remote_code=False,
         **kwargs
     ):
-        # 1) Config 로드
         config = AutoConfig.from_pretrained(
             pretrained_model_name_or_path,
             trust_remote_code=trust_remote_code,
@@ -122,33 +114,16 @@ class StudentForAudioClassification(PreTrainedModel):
         )
         model = cls(config)
 
-        # 2) 파일 위치 분기 처리 (로컬 디렉토리 vs. 허브)
+        # 🟢 [핵심] 경로가 폴더(로컬)면 직접 파일 찾기
         if os.path.isdir(pretrained_model_name_or_path):
-            # safetensors 우선, 없으면 pytorch_model.bin
-            safetensor_path = os.path.join(pretrained_model_name_or_path, "model.safetensors")
             bin_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
-            if os.path.isfile(safetensor_path):
-                from safetensors.torch import load_file as safetensors_load
-                sd = safetensors_load(safetensor_path)
-            else:
-                sd = torch.load(bin_path, map_location="cpu", weights_only=True)
         else:
-            # 허브에서 다운로드
-            try:
-                from safetensors.torch import load_file as safetensors_load
-                safetensor_path = hf_hub_download(
-                    repo_id=pretrained_model_name_or_path,
-                    filename="model.safetensors"
-                )
-                sd = safetensors_load(safetensor_path)
-            except Exception:
-                bin_path = hf_hub_download(
-                    repo_id=pretrained_model_name_or_path,
-                    filename="pytorch_model.bin"
-                )
-                sd = torch.load(bin_path, map_location="cpu", weights_only=True)
-
-        # 3) state_dict prefix
+            from huggingface_hub import hf_hub_download
+            bin_path = hf_hub_download(
+                repo_id=pretrained_model_name_or_path,
+                filename="pytorch_model.bin",
+            )
+        sd = torch.load(bin_path, map_location="cpu", weights_only=True)
         prefixed_sd = {f"student.{k}": v for k, v in sd.items()}
         model.load_state_dict(prefixed_sd, strict=True)
         return model