Spaces:
Running
Running
| """Optional pooling heads for the research/ablation track. | |
| DEPLOYMENT NOTE | |
| --------------- | |
| The production model MUST remain a stock `DebertaV2ForSequenceClassification` | |
| so the HF Inference API can serve it through the standard `text-classification` | |
| pipeline (no `trust_remote_code` on the Inference API). Everything in this file | |
| is for the ablation table in the paper, not for the deployed checkpoint. | |
| Motivation: DeBERTa's default `ContextPooler` reads only the [CLS] token. That is | |
| reasonable for a tweet, where the emotional payload is one clause. A verse of | |
| lyrics spreads the signal across the whole passage — "and at once I knew / I was | |
| not magnificent" carries its weight in the second line, not the first token. Mean | |
| and attention pooling recover that. | |
| Use `--pooling attention` in train_v2.py to run the ablation. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from transformers import AutoConfig, AutoModel | |
| from transformers.modeling_outputs import SequenceClassifierOutput | |
| class AttentionPooling(nn.Module): | |
| """Learned additive attention over token hidden states. | |
| Computes a scalar score per token, softmaxes over the (masked) sequence, and | |
| returns the weighted sum. Unlike [CLS] pooling this lets the model learn that | |
| e.g. "hollow", "still", "again" carry the emotional load in a lyric line. | |
| """ | |
| def __init__(self, hidden_size: int, dropout: float = 0.1): | |
| super().__init__() | |
| self.attn = nn.Sequential( | |
| nn.Linear(hidden_size, hidden_size), | |
| nn.Tanh(), | |
| nn.Dropout(dropout), | |
| nn.Linear(hidden_size, 1), | |
| ) | |
| def forward(self, hidden: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: | |
| # hidden: (B, T, H), mask: (B, T) | |
| scores = self.attn(hidden).squeeze(-1) # (B, T) | |
| scores = scores.masked_fill(mask == 0, torch.finfo(scores.dtype).min) | |
| weights = F.softmax(scores, dim=-1).unsqueeze(-1) # (B, T, 1) | |
| return (hidden * weights).sum(dim=1) # (B, H) | |
| class MeanMaxAttentionPool(nn.Module): | |
| """Concatenation of masked-mean, masked-max, and attention pooling. | |
| The three views are complementary: mean captures overall tone, max captures | |
| the single most charged token, attention learns a task-specific weighting. | |
| Output dimension is 3 * hidden_size. | |
| """ | |
| def __init__(self, hidden_size: int, dropout: float = 0.1): | |
| super().__init__() | |
| self.attention = AttentionPooling(hidden_size, dropout) | |
| self.output_size = hidden_size * 3 | |
| def forward(self, hidden: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: | |
| mask_f = mask.unsqueeze(-1).float() | |
| summed = (hidden * mask_f).sum(dim=1) | |
| counts = mask_f.sum(dim=1).clamp(min=1e-9) | |
| mean_pooled = summed / counts | |
| max_pooled = hidden.masked_fill(mask.unsqueeze(-1) == 0, torch.finfo(hidden.dtype).min) | |
| max_pooled = max_pooled.max(dim=1).values | |
| attn_pooled = self.attention(hidden, mask) | |
| return torch.cat([mean_pooled, max_pooled, attn_pooled], dim=-1) | |
| class MultiSampleDropoutHead(nn.Module): | |
| """Classifier that averages logits over several independent dropout masks. | |
| Cheap variance reduction: one forward pass through the backbone, N cheap | |
| passes through the linear head. Converges faster and generalizes better than | |
| a single-mask head, which matters when the tail classes have few examples. | |
| Reference: Inoue 2019, "Multi-Sample Dropout for Accelerated Training and | |
| Better Generalization" (arXiv:1905.09788) | |
| """ | |
| def __init__(self, input_size: int, num_labels: int, num_samples: int = 5, p: float = 0.3): | |
| super().__init__() | |
| self.dropouts = nn.ModuleList([nn.Dropout(p) for _ in range(num_samples)]) | |
| self.classifier = nn.Linear(input_size, num_labels) | |
| def forward(self, features: torch.Tensor) -> torch.Tensor: | |
| return torch.stack( | |
| [self.classifier(d(features)) for d in self.dropouts], dim=0 | |
| ).mean(dim=0) | |
| class EumoraClassifier(nn.Module): | |
| """DeBERTa backbone + rich pooling + multi-sample dropout head. | |
| Research-track model. Same backbone weights as the deployed checkpoint, so a | |
| TAPT-adapted encoder can be dropped into either path for a clean ablation. | |
| """ | |
| def __init__( | |
| self, | |
| model_name: str, | |
| num_labels: int, | |
| pooling: str = "attention", | |
| dropout: float = 0.3, | |
| num_dropout_samples: int = 5, | |
| ): | |
| super().__init__() | |
| self.config = AutoConfig.from_pretrained(model_name, num_labels=num_labels) | |
| self.backbone = AutoModel.from_pretrained(model_name, config=self.config) | |
| hidden = self.config.hidden_size | |
| self.pooling_type = pooling | |
| if pooling == "attention": | |
| self.pool = MeanMaxAttentionPool(hidden, dropout=0.1) | |
| head_input = self.pool.output_size | |
| elif pooling == "mean": | |
| self.pool = None | |
| head_input = hidden | |
| else: # "cls" | |
| self.pool = None | |
| head_input = hidden | |
| self.head = MultiSampleDropoutHead( | |
| head_input, num_labels, num_samples=num_dropout_samples, p=dropout | |
| ) | |
| def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs): | |
| hidden = self.backbone( | |
| input_ids=input_ids, attention_mask=attention_mask | |
| ).last_hidden_state | |
| if self.pooling_type == "attention": | |
| features = self.pool(hidden, attention_mask) | |
| elif self.pooling_type == "mean": | |
| mask_f = attention_mask.unsqueeze(-1).float() | |
| features = (hidden * mask_f).sum(dim=1) / mask_f.sum(dim=1).clamp(min=1e-9) | |
| else: | |
| features = hidden[:, 0] | |
| logits = self.head(features) | |
| return SequenceClassifierOutput(logits=logits) | |