File size: 2,339 Bytes
443fc75 | 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 | """Generic sequence classification head: AutoModel backbone + mean pooling + linear.
Auto-generated by the uv-scripts `train-classifier.py` recipe. Loaded via
`AutoModelForSequenceClassification.from_pretrained(repo, trust_remote_code=True)`;
the backbone class is resolved from this repo's own `auto_map`/code files.
"""
import torch
from torch import nn
from transformers import AutoModel, PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
class EncoderForSequenceClassification(PreTrainedModel):
base_model_prefix = "model"
supports_gradient_checkpointing = True
def __init__(self, config):
super().__init__(config)
self.num_labels = config.num_labels
self.model = AutoModel.from_config(config, trust_remote_code=True)
dropout = getattr(config, "classifier_dropout", None)
self.dropout = nn.Dropout(0.1 if dropout is None else dropout)
self.classifier = nn.Linear(config.hidden_size, config.num_labels)
self.post_init()
def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
hidden = outputs.last_hidden_state
if attention_mask is None:
pooled = hidden.mean(dim=1)
else:
mask = attention_mask.unsqueeze(-1).to(hidden.dtype)
pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
logits = self.classifier(self.dropout(pooled))
loss = None
if labels is not None:
if self.config.problem_type == "multi_label_classification":
loss = nn.functional.binary_cross_entropy_with_logits(
logits, labels.to(logits.dtype)
)
else:
loss = nn.functional.cross_entropy(logits, labels.view(-1))
return SequenceClassifierOutput(loss=loss, logits=logits)
# AutoModelForSequenceClassification.from_pretrained registers this class against the
# config class, and that requires config_class to be set (transformers v5 crashes on None).
try:
from transformers import Lfm2Config
EncoderForSequenceClassification.config_class = Lfm2Config
except ImportError: # flat import during training; the trainer sets config_class itself
pass
|