File size: 2,884 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 50 51 | """Transformer-block classifier body: self-attention across the span dimension, used by
ner_head.py / relation_head.py in the `medbit_transformer` arm in place of the plain
`FeedForward` scorer (see notes/v2_plan.md item 1).
A standard `nn.TransformerEncoderLayer` already bundles self-attention with its own
position-wise FFN sublayer (residual + norm around both), so this fully replaces
`FeedForward` rather than wrapping it -- no separate downstream MLP survives; callers add
only a final `nn.Linear` to project to label logits, same shape as today.
Spans are the "sequence" here (length = number of candidate/pruned spans, not number of
words); there is no batch dimension and no padding mask, matching every other module's
batch_size=1 simplification (see pruner.py, span_extractor.py). A sinusoidal positional
encoding keyed on each span's start-word index gives attention a notion of span order/distance
-- the same choice already made in v2_plan.md item 1 for the (superseded) word-level
context_layer idea, just applied to span positions instead of word positions.
"""
import math
import torch
from torch import nn
def sinusoidal_position_encoding(positions: torch.Tensor, d_model: int) -> torch.Tensor:
"""positions: (N,) int64 span start indices. Returns (N, d_model)."""
div_term = torch.exp(torch.arange(0, d_model, 2, device=positions.device, dtype=torch.float32) *
(-math.log(10000.0) / d_model))
angles = positions.unsqueeze(1).to(torch.float32) * div_term.unsqueeze(0) # (N, d_model/2)
pe = torch.zeros(positions.size(0), d_model, device=positions.device)
pe[:, 0::2] = torch.sin(angles)
pe[:, 1::2] = torch.cos(angles[:, :pe[:, 1::2].size(1)])
return pe
class TransformerBlock(nn.Module):
def __init__(self, input_dim: int, d_model: int, nhead: int, num_layers: int,
dim_feedforward: int, dropout: float, activation: str = "gelu"):
super().__init__()
self.input_proj = nn.Linear(input_dim, d_model) if input_dim != d_model else nn.Identity()
layer = nn.TransformerEncoderLayer(
d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout,
activation=activation, norm_first=True, batch_first=True)
self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers, enable_nested_tensor=False)
self.d_model = d_model
self.output_dim = d_model
def forward(self, span_embeddings: torch.Tensor, start_positions: torch.Tensor) -> torch.Tensor:
"""span_embeddings: (num_spans, input_dim). start_positions: (num_spans,) int64.
No batch dim on the way in/out -- one fake batch of size 1 for nn.TransformerEncoder."""
x = self.input_proj(span_embeddings) + sinusoidal_position_encoding(start_positions, self.d_model)
return self.encoder(x.unsqueeze(0)).squeeze(0)
|