| from __future__ import annotations |
| import math |
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
|
|
| class PointerDecisionHead(nn.Module): |
| def __init__(self, hidden_size: int, head_dim: int = 512, normalize: bool = True): |
| super().__init__() |
| self.normalize = normalize |
| self.q_proj = nn.Linear(hidden_size, head_dim, bias=False) |
| self.k_proj = nn.Linear(hidden_size, head_dim, bias=False) |
| self.scale = math.sqrt(head_dim) |
|
|
| def forward(self, decide_h: torch.Tensor, option_h: torch.Tensor, option_mask: torch.Tensor) -> torch.Tensor: |
| decide_h = decide_h.float() |
| option_h = option_h.float() |
| q = self.q_proj(decide_h) |
| k = self.k_proj(option_h) |
| if self.normalize: |
| q = F.normalize(q, dim=-1) |
| k = F.normalize(k, dim=-1) |
| logits = torch.einsum("bd,bnd->bn", q, k) * 10.0 |
| else: |
| logits = torch.einsum("bd,bnd->bn", q, k) / self.scale |
| return logits.masked_fill(~option_mask, -1e9) |
|
|
|
|
| class SpanDecisionHead(nn.Module): |
| """Scores the mean hidden state across each option's semantic token span.""" |
| def __init__(self, hidden_size: int, head_dim: int = 384): |
| super().__init__() |
| self.q_proj = nn.Linear(hidden_size, head_dim, bias=False) |
| self.k_proj = nn.Linear(hidden_size, head_dim, bias=False) |
| self.mlp = nn.Sequential( |
| nn.Linear(head_dim * 4, head_dim), |
| nn.GELU(), |
| nn.Linear(head_dim, 1), |
| ) |
|
|
| def forward(self, decide_h: torch.Tensor, option_mean_h: torch.Tensor, option_mask: torch.Tensor) -> torch.Tensor: |
| q = F.normalize(self.q_proj(decide_h.float()), dim=-1) |
| k = F.normalize(self.k_proj(option_mean_h.float()), dim=-1) |
| qe = q[:, None, :].expand_as(k) |
| x = torch.cat([qe, k, qe * k, (qe - k).abs()], dim=-1) |
| logits = self.mlp(x).squeeze(-1) |
| return logits.masked_fill(~option_mask, -1e9) |
|
|
|
|
| class HybridDecisionHead(nn.Module): |
| """Pointer baseline plus a learned residual semantic-span scorer.""" |
| def __init__(self, hidden_size: int, head_dim: int = 384, pointer_dim: int = 384, normalize: bool = True): |
| super().__init__() |
| self.pointer = PointerDecisionHead(hidden_size, pointer_dim, normalize=normalize) |
| self.span = SpanDecisionHead(hidden_size, head_dim) |
| self.logit_scale = nn.Parameter(torch.tensor(0.0)) |
|
|
| def forward( |
| self, |
| decide_h: torch.Tensor, |
| option_last_h: torch.Tensor, |
| option_mean_h: torch.Tensor, |
| option_mask: torch.Tensor, |
| ) -> torch.Tensor: |
| base = self.pointer(decide_h, option_last_h, option_mask) |
| residual = self.span(decide_h, option_mean_h, option_mask) |
| return base + self.logit_scale.exp().clamp(max=4.0) * residual |
|
|
|
|
| def brier_loss(probs: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: |
| diff = (probs - targets).pow(2) * mask.float() |
| return diff.sum(dim=-1).mean() |
|
|
