File size: 2,733 Bytes
0c48771 | 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 | """Span representations: concat(start word embedding, end word embedding, [mean-pooled interior
embedding], width embedding).
Port of AllenNLP's `EndpointSpanExtractor(combination="x,y", bucket_widths=False)`, the only
configuration v1 used. v2 adds an optional mean-pooled block -- the average of every word
embedding inside [start, end] inclusive -- gated by `mean_pool` (see notes/v2_plan.md item 3,
`configs/medbit_span_pooling/`). Simplified for the batch_size=1 case (see dataset.py): no batch
dimension, no span masking -- every document has an exact, un-padded list of spans.
"""
import torch
from torch import nn
class EndpointSpanExtractor(nn.Module):
def __init__(self, input_dim: int, num_width_embeddings: int, span_width_embedding_dim: int,
mean_pool: bool = False):
super().__init__()
self.width_embedding = nn.Embedding(num_width_embeddings, span_width_embedding_dim)
self.mean_pool = mean_pool
self.output_dim = (3 if mean_pool else 2) * input_dim + span_width_embedding_dim
def get_output_dim(self) -> int:
return self.output_dim
def forward(self, word_embeddings: torch.Tensor, spans: torch.LongTensor) -> torch.Tensor:
"""word_embeddings: (num_words, input_dim). spans: (num_spans, 2) inclusive [start, end]."""
starts, ends = spans[:, 0], spans[:, 1]
start_emb = word_embeddings[starts]
end_emb = word_embeddings[ends]
width_emb = self.width_embedding(ends - starts) # 0-indexed width (widths 1..N -> 0..N-1)
if not self.mean_pool:
return torch.cat([start_emb, end_emb, width_emb], dim=-1)
mean_emb = self._interior_mean(word_embeddings, starts, ends)
return torch.cat([start_emb, end_emb, mean_emb, width_emb], dim=-1)
@staticmethod
def _interior_mean(word_embeddings: torch.Tensor, starts: torch.LongTensor,
ends: torch.LongTensor) -> torch.Tensor:
"""Mean of word_embeddings[start:end+1] per span, vectorized over the whole (bounded-
width) span list: gather a (num_spans, max_width) index grid and masked-average it. No
cumsum needed -- widths are capped by max_span_width (e.g. 12)."""
widths = ends - starts + 1 # (num_spans,)
max_width = int(widths.max().item())
offsets = torch.arange(max_width, device=word_embeddings.device).unsqueeze(0) # (1, W)
mask = offsets < widths.unsqueeze(1) # (num_spans, W)
idx = (starts.unsqueeze(1) + offsets).clamp(max=word_embeddings.size(0) - 1)
gathered = word_embeddings[idx] * mask.unsqueeze(-1) # (num_spans, W, input_dim)
return gathered.sum(dim=1) / widths.unsqueeze(1).to(word_embeddings.dtype)
|