File size: 3,841 Bytes
fcf6574 8ea28d5 fcf6574 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | 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,
):
super().__init__(config)
self.num_labels = config.num_labels
self.config = config
self.backbone = AutoModel.from_config(config)
head_layers = getattr(config, "head_layers", 1)
head_dropout = getattr(config, "head_dropout", 0.1)
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,
)
|