parser/biaffine.py

"""Deep biaffine dependency parser (Dozat & Manning 2017) over a frozen Stoicheia
encoder, combined via an ELMo-style learned scalar mix of all layers (same recipe the
tagger used for XPOS/lemma: scalar_mix + light head, here applied to arc/label MLPs
instead of tag heads).

Backbones are FROZEN (only the scalar-mix weights + biaffine head train) — this keeps a fair,
fast 3-way ablation (char / lemma / fused) without re-touching either finished pretraining run.
"""
from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


class ScalarMix(nn.Module):
    """Learned softmax-weighted sum over N layer outputs (+ a global scale), ELMo-style."""
    def __init__(self, n_layers):
        super().__init__()
        self.w = nn.Parameter(torch.zeros(n_layers))
        self.gamma = nn.Parameter(torch.ones(1))

    def forward(self, layers):           # layers: list of (B,T,D), len n_layers
        w = torch.softmax(self.w, 0)
        mixed = sum(wi * h for wi, h in zip(w, layers))
        return self.gamma * mixed


def pool_words(hidden, word_id, W, mode="mean"):
    B, T, D = hidden.shape
    flat = hidden.reshape(B * T, D)
    wid = word_id.reshape(B * T)
    valid = wid >= 0
    off = (torch.arange(B, device=hidden.device) * W).repeat_interleave(T)
    idx = (wid + off)[valid]
    out = hidden.new_zeros(B * W, D)
    out.index_add_(0, idx, flat[valid])
    cnt = hidden.new_zeros(B * W).index_add_(0, idx, torch.ones_like(idx, dtype=hidden.dtype))
    out = out / cnt.clamp(min=1).unsqueeze(-1)
    return out.reshape(B, W, D)


class MLP(nn.Module):
    def __init__(self, d_in, d_out, dropout=0.33):
        super().__init__()
        self.lin = nn.Linear(d_in, d_out)
        self.act = nn.LeakyReLU(0.1)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        return self.drop(self.act(self.lin(x)))


class Biaffine(nn.Module):
    """s(x,y) = [x;1]^T W [y;1] (per output channel). x:(B,Lx,Di) y:(B,Ly,Di) -> (B,n_out,Lx,Ly)."""
    def __init__(self, d_in, n_out=1, bias_x=True, bias_y=True):
        super().__init__()
        self.bias_x, self.bias_y = bias_x, bias_y
        self.W = nn.Parameter(torch.zeros(n_out, d_in + int(bias_x), d_in + int(bias_y)))
        nn.init.xavier_uniform_(self.W)

    def forward(self, x, y):
        if self.bias_x:
            x = torch.cat([x, torch.ones_like(x[..., :1])], -1)
        if self.bias_y:
            y = torch.cat([y, torch.ones_like(y[..., :1])], -1)
        s = torch.einsum("bxi,oij,byj->boxy", x, self.W, y)
        return s.squeeze(1) if s.shape[1] == 1 else s


@dataclass
class ParserConfig:
    d_arc: int = 500
    d_rel: int = 150
    dropout: float = 0.33
    n_labels: int = 40


class BiaffineHead(nn.Module):
    """Arc + label scoring over word vectors (B,W,D). A learnable ROOT vector is the head
    candidate for column 0 (D&M's pseudo-root). Arc loss / label loss are standard CE."""
    def __init__(self, d_in, cfg: ParserConfig):
        super().__init__()
        self.cfg = cfg
        self.root = nn.Parameter(torch.zeros(d_in))
        nn.init.normal_(self.root, std=0.02)
        self.arc_dep = MLP(d_in, cfg.d_arc, cfg.dropout)
        self.arc_head = MLP(d_in, cfg.d_arc, cfg.dropout)
        self.rel_dep = MLP(d_in, cfg.d_rel, cfg.dropout)
        self.rel_head = MLP(d_in, cfg.d_rel, cfg.dropout)
        self.arc_biaf = Biaffine(cfg.d_arc, n_out=1, bias_x=True, bias_y=False)
        self.rel_biaf = Biaffine(cfg.d_rel, n_out=cfg.n_labels, bias_x=True, bias_y=True)

    def forward(self, w, word_mask):
        """w: (B,W,D) word vectors. word_mask: (B,W) bool, True at real words.
        Returns arc_scores (B,W,W+1) [col0=root], rel_scores (B,W,W+1,n_labels)."""
        B, W, D = w.shape
        root = self.root.view(1, 1, D).expand(B, 1, D)
        heads_in = torch.cat([root, w], 1)                       # (B,W+1,D): col0=root
        h_dep_arc = self.arc_dep(w)                               # (B,W,d_arc)
        h_head_arc = self.arc_head(heads_in)                      # (B,W+1,d_arc)
        arc_scores = self.arc_biaf(h_dep_arc, h_head_arc)          # (B,W,W+1)
        # mask: dependent i cannot pick itself as head (col i+1), and padded cols get -inf
        pad_head = torch.cat([torch.ones(B, 1, dtype=torch.bool, device=w.device), word_mask], 1)
        arc_scores = arc_scores.masked_fill(~pad_head[:, None, :], float("-inf"))
        self_idx = torch.arange(W, device=w.device)
        arc_scores[:, self_idx, self_idx + 1] = float("-inf")

        h_dep_rel = self.rel_dep(w)                                # (B,W,d_rel)
        h_head_rel = self.rel_head(heads_in)                       # (B,W+1,d_rel)
        rel_scores = self.rel_biaf(h_dep_rel, h_head_rel)           # (B,n_labels,W,W+1)
        rel_scores = rel_scores.permute(0, 2, 3, 1)                 # (B,W,W+1,n_labels)
        return arc_scores, rel_scores

    def loss(self, arc_scores, rel_scores, heads, labels, word_mask):
        """heads: (B,W) gold head col index (0=root, else 1..W); labels: (B,W) gold label id;
        both -100 where padded/not-a-word."""
        m = heads != -100
        arc_loss = F.cross_entropy(arc_scores[m], heads[m])
        B, W = heads.shape
        bi = torch.arange(B, device=heads.device)[:, None].expand(B, W)
        wi = torch.arange(W, device=heads.device)[None, :].expand(B, W)
        gold_head = heads.clamp(min=0)
        sel = rel_scores[bi, wi, gold_head]                        # (B,W,n_labels)
        rel_loss = F.cross_entropy(sel[m], labels[m])
        return arc_loss + rel_loss, dict(arc=round(arc_loss.item(), 4), rel=round(rel_loss.item(), 4))

    @torch.no_grad()
    def decode(self, arc_scores, rel_scores, word_mask):
        """Greedy per-token argmax head (col 0=root) + label argmax at the chosen head.
        Not tree-constrained (no MST projection) — the official conll18 LAS/UAS scorer
        compares HEAD/DEPREL per token regardless of global tree-validity, so this is a
        correct, simple decode for that metric (the standard simplification vs. full MST)."""
        heads_out = arc_scores.argmax(-1)                          # (B,W) in [0..W], 0=root
        B, W = heads_out.shape
        bi = torch.arange(B, device=arc_scores.device)[:, None].expand(B, W)
        wi = torch.arange(W, device=arc_scores.device)[None, :].expand(B, W)
        labels_out = rel_scores[bi, wi, heads_out].argmax(-1)
        return heads_out.cpu(), labels_out.cpu()