Update modeling_generanno.py

modeling_generanno.py

@@ -1140,6 +1140,7 @@ class GenerannoForSequenceClassification(GenerannoPreTrainedModel):
 
         self.model = GenerannoModel(config)
         self.feature_layer = getattr(config, "feature_layer", -1)
+        self.use_mean_pooling = getattr(config, "use_mean_pooling", True)
         self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
         if getattr(config, "use_mlp_classifier", False):
             self.score = nn.Sequential(
@@ -1152,6 +1153,23 @@ class GenerannoForSequenceClassification(GenerannoPreTrainedModel):
         # Initialize weights and apply final processing
         self.post_init()
 
+    def _apply_mean_pooling(self, hidden_states, attention_mask):
+        if attention_mask is None:
+            return torch.mean(hidden_states, dim=1)
+
+        # Expand attention mask to match hidden states dimensions
+        input_mask_expanded = attention_mask.unsqueeze(-1).expand(hidden_states.size()).float()
+        sum_embeddings = torch.sum(hidden_states * input_mask_expanded, dim=1)
+        
+        # Compute number of valid tokens per sequence
+        sum_mask = input_mask_expanded.sum(dim=1)
+        sum_mask = torch.clamp(sum_mask, min=1e-9)
+        
+        # Compute mean
+        pooled_output = sum_embeddings / sum_mask
+
+        return pooled_output
+
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,
@@ -1173,20 +1191,33 @@ class GenerannoForSequenceClassification(GenerannoPreTrainedModel):
             return_dict if return_dict is not None else self.config.use_return_dict
         )
 
-        output_hidden_states = True
-        outputs = self.model(
-            input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            inputs_embeds=inputs_embeds,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-        hidden_states = outputs["hidden_states"][
-            self.feature_layer if hasattr(self, "feature_layer") else -1
-        ]
-        pooled_hidden_states = hidden_states[:, 0]
+        if self.feature_layer == -1:
+            outputs = self.model(
+                input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                inputs_embeds=inputs_embeds,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+            hidden_states = outputs[0]
+        else:
+            outputs = self.model(
+                input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                inputs_embeds=inputs_embeds,
+                output_attentions=output_attentions,
+                output_hidden_states=True,
+                return_dict=return_dict,
+            )
+            hidden_states = outputs.hidden_states[self.feature_layer]
+
+        if self.use_mean_pooling:
+            pooled_hidden_states = self._apply_mean_pooling(hidden_states, attention_mask)
+        else:
+            pooled_hidden_states = hidden_states[:, 0]
         logits = self.score(pooled_hidden_states)
 
         loss = None