Spaces:
Sleeping
Sleeping
| """Stage 2: encoder over `ctx [BOP] span [EOP] ctx` -> technique (plan D3, D7). | |
| Representation = [CLS ; h(BOP) ; h(EOP)], a common entity-marker pooling. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| class SpanClassifier(nn.Module): | |
| kind = "tc" | |
| def __init__(self, encoder: nn.Module, num_labels: int, dropout: float = 0.1, | |
| class_weights: torch.Tensor | None = None): | |
| super().__init__() | |
| self.encoder = encoder | |
| h = encoder.config.hidden_size | |
| self.head = nn.Sequential( | |
| nn.Dropout(dropout), nn.Linear(3 * h, h), nn.GELU(), | |
| nn.Dropout(dropout), nn.Linear(h, num_labels), | |
| ) | |
| self.register_buffer("class_weights", class_weights, persistent=False) | |
| def forward(self, input_ids, attention_mask, bop_pos, eop_pos, labels=None) -> dict: | |
| hidden = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state | |
| rows = torch.arange(hidden.size(0), device=hidden.device) | |
| rep = torch.cat([hidden[:, 0], hidden[rows, bop_pos], hidden[rows, eop_pos]], dim=-1) | |
| logits = self.head(rep).float() | |
| out = {"logits": logits} | |
| if labels is not None: | |
| out["loss"] = F.cross_entropy(logits, labels, weight=self.class_weights) | |
| return out | |