switch automodel for automodelforseqclass and ensure correct broadcast in loss computation

modeling_emcoder.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from .rope_embeddings import RotaryEmbedding
-from transformers import PreTrainedModel, AutoConfig, AutoModel
+from transformers import PreTrainedModel, AutoConfig, AutoModelForSequenceClassification
 from transformers.modeling_outputs import SequenceClassifierOutput
 
 from .configuration_emcoder import EmCoderConfig
@@ -202,12 +202,12 @@ class EmCoder(PreTrainedModel):
         """
         return_dict = return_dict if return_dict is not None else True
 
-        x = input_ids if input_ids is not None else kwargs.get("x")
-        mask = attention_mask if attention_mask is not None else kwargs.get("mask")
-        
-        if x is None or mask is None:
-            raise ValueError("input_ids (x) and attention_mask (mask) must be provided")
+        x = input_ids if input_ids is not None else kwargs.get("input_ids")
+        mask = attention_mask if attention_mask is not None else kwargs.get("attention_mask")
 
+        if x is None or mask is None:
+            raise ValueError("input_ids and attention_mask must be provided")
+        
         if max_batch_size is None:
             max_batch_size = n_samples
 
@@ -241,7 +241,9 @@ class EmCoder(PreTrainedModel):
         loss = None
         if labels is not None:
             loss_fct = nn.BCEWithLogitsLoss()
-            loss = loss_fct(all_logits.mean(dim=0), labels.to(all_logits.dtype))
+            logits_mean = all_logits.mean(dim=0)  # (B, num_labels)
+            target_labels = labels.to(dtype=all_logits.dtype).view(logits_mean.shape)
+            loss = loss_fct(logits_mean, target_labels)
 
         if not return_dict:
             output = (all_logits,)
@@ -266,11 +268,11 @@ class EmCoder(PreTrainedModel):
         """Standard forward pass without MC Dropout."""
         return_dict = return_dict if return_dict is not None else True
 
-        x = input_ids if input_ids is not None else kwargs.get("x")
-        mask = attention_mask if attention_mask is not None else kwargs.get("mask")
+        x = input_ids if input_ids is not None else kwargs.get("input_ids")
+        mask = attention_mask if attention_mask is not None else kwargs.get("attention_mask")
         
         if x is None or mask is None:
-            raise ValueError("input_ids (x) and attention_mask (mask) must be provided")
+            raise ValueError("input_ids and attention_mask must be provided")
 
         features = self.encoder(x, mask)
 
@@ -281,7 +283,8 @@ class EmCoder(PreTrainedModel):
         loss = None
         if labels is not None:
             loss_fct = nn.BCEWithLogitsLoss()
-            loss = loss_fct(logits, labels.to(logits.dtype))
+            target_labels = labels.to(dtype=logits.dtype).view(logits.shape)
+            loss = loss_fct(logits, target_labels)
 
         if not return_dict:
             output = (logits,)
@@ -296,6 +299,6 @@ class EmCoder(PreTrainedModel):
 
 try:
     AutoConfig.register("emcoder", EmCoderConfig)
-    AutoModel.register(EmCoderConfig, EmCoder)
+    AutoModelForSequenceClassification.register(EmCoderConfig, EmCoder)
 except ValueError:
     pass