Upload modeling.py

modeling.py

@@ -0,0 +1,111 @@
+import torch
+import torch.nn as nn
+
+from transformers import AutoModel, PreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+
+class CustomSequenceClassificationModel(PreTrainedModel):
+    """Sequence classification model with a configurable GELU + dropout dense head."""
+
+    def __init__(
+        self,
+        config,
+        head_layers: int = 1,
+        head_dropout: float = 0.1,
+    ):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.backbone = AutoModel.from_config(config)
+
+        hidden_size = getattr(config, "hidden_size", None)
+        if hidden_size is None:
+            raise ValueError(
+                "Backbone config is missing 'hidden_size'; cannot build custom head."
+            )
+
+        self.classifier = self._build_head(
+            input_dim=hidden_size,
+            output_dim=self.num_labels,
+            head_layers=head_layers,
+            head_dropout=head_dropout,
+        )
+        self.post_init()
+
+    @staticmethod
+    def _build_head(
+        input_dim: int, output_dim: int, head_layers: int, head_dropout: float
+    ) -> nn.Module:
+        if head_layers < 1:
+            raise ValueError("head_layers must be >= 1")
+
+        # head_layers includes the final projection layer.
+        # For head_layers=1 this is equivalent to a single linear classifier.
+        if head_layers == 1:
+            return nn.Linear(input_dim, output_dim)
+
+        layers = []
+        for _ in range(head_layers - 1):
+            layers.extend(
+                [
+                    nn.Linear(input_dim, input_dim),
+                    nn.GELU(),
+                    nn.Dropout(head_dropout),
+                ]
+            )
+        layers.append(nn.Linear(input_dim, output_dim))
+        return nn.Sequential(*layers)
+
+    def _pooled_representation(self, outputs) -> torch.Tensor:
+        if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
+            return outputs.pooler_output
+        return outputs.last_hidden_state[:, 0]
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        labels=None,
+        **kwargs,
+    ):
+        backbone_inputs = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            **kwargs,
+        }
+        if token_type_ids is not None:
+            backbone_inputs["token_type_ids"] = token_type_ids
+
+        backbone_outputs = self.backbone(**backbone_inputs)
+
+        pooled = self._pooled_representation(backbone_outputs)
+        logits = self.classifier(pooled)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif torch.is_floating_point(labels):
+                    self.config.problem_type = "multi_label_classification"
+                else:
+                    self.config.problem_type = "single_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = nn.MSELoss()
+                loss = loss_fct(logits.squeeze(), labels.squeeze())
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = nn.CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            else:
+                loss_fct = nn.BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=backbone_outputs.hidden_states,
+            attentions=backbone_outputs.attentions,
+        )