wrong class fix

app.py

@@ -19,27 +19,55 @@ class WhisperEncoderOnlyConfig(WhisperConfig):
         self.n_super = n_super
         self.n_code = n_code
 
-class WhisperEncoderOnlyForClassification(nn.Module):
+class WhisperEncoderOnlyForClassification(WhisperPreTrainedModel):
     config_class = WhisperEncoderOnlyConfig
-    
+
     def __init__(self, config):
-        super().__init__()
-        self.config = config
+        super().__init__(config)
+
         self.encoder = WhisperEncoder(config)
-        
+
         hidden = config.d_model
         self.fam_head = nn.Linear(hidden, config.n_fam)
         self.super_head = nn.Linear(hidden, config.n_super)
         self.code_head = nn.Linear(hidden, config.n_code)
-    
-    def forward(self, input_features):
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        """Whisper doesn't have token embeddings"""
+        return None
+
+    def set_input_embeddings(self, value):
+        """Ignore"""
+        pass
+
+    def enable_input_require_grads(self):
+        return
+
+    def forward(self, input_features, labels=None):
         enc_out = self.encoder(input_features=input_features)
         pooled = enc_out.last_hidden_state.mean(dim=1)
-        
+
+        fam_logits = self.fam_head(pooled)
+        super_logits = self.super_head(pooled)
+        code_logits = self.code_head(pooled)
+
+        loss = None
+        if labels is not None:
+            fam_labels, super_labels, code_labels = labels
+            loss_fn = nn.CrossEntropyLoss()
+            loss = (
+                loss_fn(fam_logits, fam_labels) +
+                loss_fn(super_logits, super_labels) +
+                loss_fn(code_logits, code_labels)
+            )
+
         return {
-            "fam_logits": self.fam_head(pooled),
-            "super_logits": self.super_head(pooled),
-            "code_logits": self.code_head(pooled),
+            "loss": loss,
+            "fam_logits": fam_logits,
+            "super_logits": super_logits,
+            "code_logits": code_logits,
         }
 
 # === REGISTER MODEL ===