Update models.py

models.py

@@ -130,6 +130,36 @@ class StudentForAudioClassification(PreTrainedModel):
         self.post_init()
 
     def forward(self, input_values, **kwargs):
-        # input_values: already processed W2V embeddings or raw audio features
         logits, feat = self.student(input_values)
         return SequenceClassifierOutput(logits=logits)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *args,
+        trust_remote_code=False,
+        **kwargs
+    ):
+        # 1) 기본 Config 로드 (trust_remote_code=True 권장)
+        config = AutoConfig.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs
+        )
+        # 2) 빈 래퍼 인스턴스 생성
+        model = cls(config)
+
+        # 3) 원본 체크포인트 state_dict 로드
+        #    (safetensors라면 safetensors 라이브러리로 읽으세요)
+        sd = torch.load(
+            f"{pretrained_model_name_or_path}/pytorch_model.bin",
+            map_location="cpu",
+            weights_only=True
+        )
+        # 4) 모든 키에 `student.` prefix 추가
+        prefixed_sd = {f"student.{k}": v for k, v in sd.items()}
+
+        # 5) 래퍼 모델에 가중치 로드
+        model.load_state_dict(prefixed_sd, strict=True)
+        return model