File size: 1,747 Bytes
81e8ada | 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 | from __future__ import annotations
import torch
from torch import nn
from transformers import AutoConfig, AutoModel
class EncoderClassifier(nn.Module):
"""Small encoder + sentence pooling + a linear classification head."""
def __init__(
self,
base_model: str,
num_labels: int,
dropout: float = 0.1,
pooling: str = "mean",
pretrained: bool = True,
encoder_config: dict | None = None,
) -> None:
super().__init__()
if pretrained:
self.encoder = AutoModel.from_pretrained(base_model)
else:
if encoder_config is None:
raise ValueError("encoder_config is required when pretrained=False")
raw_config = dict(encoder_config)
model_type = raw_config.pop("model_type")
config = AutoConfig.for_model(model_type, **raw_config)
self.encoder = AutoModel.from_config(config)
self.dropout = nn.Dropout(dropout)
self.pooling = pooling
self.classifier = nn.Linear(self.encoder.config.hidden_size, num_labels)
def forward(
self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
) -> torch.Tensor:
hidden = self.encoder(
input_ids=input_ids,
attention_mask=attention_mask,
).last_hidden_state
if self.pooling == "cls":
pooled = hidden[:, 0]
elif self.pooling == "mean":
mask = attention_mask.unsqueeze(-1).to(hidden.dtype)
pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
else:
raise ValueError(f"Unsupported pooling: {self.pooling}")
return self.classifier(self.dropout(pooled))
|