Update modeling.py

modeling.py

@@ -1,69 +1,20 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from transformers import PreTrainedModel, AutoModel, AutoConfig
 
-from transformers import (
-    PreTrainedModel,
-    PretrainedConfig,
-    AutoModel,
-    AutoConfig,
-)
-
-
-# ============================================================
-# CONFIG CLASS (NO **kwargs in __init__)
-# ============================================================
-class EmoAxisConfig(PretrainedConfig):
-    model_type = "emoaxis"
-
-    def __init__(self, num_labels=28, base_model_name="roberta-base"):
-        # Explicitly pass model_type to super() if removing **kwargs
-        super().__init__(model_type=self.model_type, num_labels=num_labels, base_model_name=base_model_name)
-        self.num_labels = num_labels
-        self.base_model_name = base_model_name
-
-    # CRITICAL FIX: Overload from_pretrained (must retain **kwargs here to work)
-    # NOTE: You MUST retain **kwargs in the from_pretrained signature for it to function correctly
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
-        kwargs.pop("return_unused_kwargs", None)
-        kwargs.pop("config", None)
-        
-        # Call the base PretrainedConfig's method directly
-        return PretrainedConfig.from_pretrained(
-            cls, 
-            pretrained_model_name_or_path, 
-            **kwargs
-        )
-
-
-# ============================================================
-# ENCODER MODULE
-# ============================================================
 class Encoder(nn.Module):
     def __init__(self, base_encoder):
         super().__init__()
         self.encoder = base_encoder
 
     def forward(self, inputs):
-        # NOTE: If the base encoder takes specific non-standard arguments, 
-        # this will break without **kwargs. We assume only input_ids/attention_mask are passed.
-        outputs = self.encoder(
-            input_ids=inputs["input_ids"], 
-            attention_mask=inputs["attention_mask"], 
-            output_hidden_states=True
-        )
+        outputs = self.encoder(**inputs, output_hidden_states=True)
         last_hidden = outputs.hidden_states[-1]
-
         mask = inputs["attention_mask"].unsqueeze(-1).float()
         pooled = (last_hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
-
         return F.normalize(pooled, p=2, dim=1)
 
-
-# ============================================================
-# CLASSIFIER MODULE
-# ============================================================
 class Classifier(nn.Module):
     def __init__(self, input_dim=768, num_classes=28):
         super().__init__()
@@ -78,28 +29,17 @@ class Classifier(nn.Module):
     def forward(self, x):
         return self.mlp(x)
 
-
-# ============================================================
-# MAIN MODEL (NO *args or **kwargs)
-# ============================================================
 class EmoAxis(PreTrainedModel):
-    config_class = EmoAxisConfig
-    
-    # Removed *args and **kwargs from signature
-    def __init__(self, config): 
+    config_class = AutoConfig
+
+    def __init__(self, config):
         super().__init__(config)
-        
-        # This line remains correct for loading saved weights
-        base_encoder = AutoModel.from_config(AutoConfig.from_pretrained(config.base_model_name))
-        
+        base_encoder = AutoModel.from_config(config)  # IMPORTANT: use from_config
         self.encoder = Encoder(base_encoder)
-        self.classifier = Classifier(
-            input_dim=base_encoder.config.hidden_size,
-            num_classes=config.num_labels
-        )
+        self.classifier = Classifier(input_dim=base_encoder.config.hidden_size,num_classes=config.num_labels)
 
-    # Removed **kwargs from signature
-    def forward(self, input_ids=None, attention_mask=None): 
+        
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
         inputs = {
             "input_ids": input_ids,
             "attention_mask": attention_mask,