| import torch |
| import torch.nn as nn |
| from transformers import AutoConfig, AutoModel, PreTrainedModel |
| from transformers.modeling_outputs import SequenceClassifierOutput |
| from .configuration_baseline_enc import BaselineEncConfig |
|
|
| class Encoder(nn.Module): |
| def __init__(self, name, pooling="cls"): |
| super().__init__() |
| base_cfg = AutoConfig.from_pretrained(name) |
| self.backbone = AutoModel.from_config(base_cfg) |
| self.pooling = pooling |
| self.hidden_dim = self.backbone.config.hidden_size |
|
|
| def forward(self, ids, mask): |
| hidden = self.backbone(input_ids=ids, attention_mask=mask).last_hidden_state |
| if self.pooling == "cls": |
| return hidden[:, 0, :] |
| m = mask.unsqueeze(-1).float() |
| return (hidden * m).sum(1) / m.sum(1).clamp(min=1e-9) |
|
|
| class SarcasmHead(nn.Module): |
| def __init__(self, in_dim, hidden_dim, dropout=0.1): |
| super().__init__() |
| self.mlp = nn.Sequential( |
| nn.Linear(in_dim, hidden_dim), |
| nn.ReLU(), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_dim, 1), |
| ) |
|
|
| def forward(self, hidden): |
| return self.mlp(hidden).squeeze(-1) |
|
|
| class BaselineEncForSequenceClassification(PreTrainedModel): |
| config_class = BaselineEncConfig |
| base_model_prefix = "baseline_enc" |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.encoder = Encoder(config.encoder_name, config.pooling) |
| self.head = SarcasmHead( |
| self.encoder.hidden_dim, |
| config.mlp_hidden, |
| config.dropout, |
| ) |
| self.post_init() |
|
|
| def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs): |
| pooled = self.encoder(input_ids, attention_mask) |
| binary_logit = self.head(pooled) |
| logits = torch.stack([-binary_logit, binary_logit], dim=-1) |
|
|
| loss = None |
| if labels is not None: |
| loss = nn.CrossEntropyLoss()(logits, labels) |
|
|
| return SequenceClassifierOutput(loss=loss, logits=logits) |
|
|
