File size: 4,343 Bytes
7ed86c3 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """JointModel: one fine-tuned Stoicheia backbone -> word vectors -> ALL CoNLL-U columns.
Shares the tagger's pooled + scalar-mixed word representation between:
* the tagger heads (factored XPOS / lemma edit-script / UPOS / flat-tag) — row-level, and
* a Dozat-Manning biaffine head (HEAD + DEPREL) — regrouped per sentence.
The tagger packs several sentences per row (block-diagonal attention) and never splits a
sentence across rows, so each word slot's (sentence_index, token_index) in `batch['slots']`
lets us gather a sentence's word vectors back into a contiguous (n_sent, max_w, D) tensor for
the arc/label scorer. Word order within a sentence is token order over *encodable* tokens —
exactly parser.model.build_gold's indexing, so gold heads/labels line up slot-for-slot.
"""
from __future__ import annotations
import torch
import torch.nn as nn
from tagger.model import TaggerModel, pool_words
from parser.biaffine import BiaffineHead
class JointModel(nn.Module):
def __init__(self, encoder, vocab, tcfg, pcfg, W=384):
super().__init__()
self.tagger = TaggerModel(encoder, vocab, tcfg, W=W) # heads + scalar mix + LM-head freeze
self.biaffine = BiaffineHead(encoder.cfg.d_model, pcfg)
self.W = W
# ---- shared word representation (mirrors TaggerModel.forward's pooling) --------------
def _word_vectors(self, batch):
enc = self.tagger.encoder
out = enc(batch)
if self.tagger.tcfg.scalar_mix:
pooled = torch.stack(
[pool_words(h, batch["word_id"], self.W, self.tagger.tcfg.pool)
for h in [*out["layers"], out["hidden"]]]) # (L+1,B,W,D)
mix = torch.softmax(self.tagger.mix_w, 0)
w = torch.einsum("l,lbwd->bwd", mix.to(pooled.dtype), pooled)
else:
w = pool_words(out["hidden"], batch["word_id"], self.W, self.tagger.tcfg.pool)
return w
def _tag_heads(self, w):
tg = self.tagger
wt = tg.dropout(w)
r = dict(xpos=[hd(wt) for hd in tg.xpos_heads],
script=tg.head_script(wt),
upos=tg.head_upos(wt))
if tg.head_flat is not None:
r["flat"] = tg.head_flat(wt)
return r
@staticmethod
def _regroup(w, slots):
"""(B,W,D) row-packed word vectors -> (n_sent, max_w, D) per-sentence, using slots.
Returns (w_sent, word_mask, sent_ids) where sent_ids[local] is the global sentence
index and word position = running per-sentence counter (== build_gold order)."""
B, W, D = w.shape
sent_ids, sid_to_local, pos_counter = [], {}, {}
b_idx, k_idx, l_idx, p_idx = [], [], [], []
for b, rs in enumerate(slots):
for k, (si, _ti) in enumerate(rs):
if si not in sid_to_local:
sid_to_local[si] = len(sent_ids); sent_ids.append(si)
local = sid_to_local[si]
pos = pos_counter.get(local, 0); pos_counter[local] = pos + 1
b_idx.append(b); k_idx.append(k); l_idx.append(local); p_idx.append(pos)
n_sent = len(sent_ids)
max_w = max(pos_counter.values(), default=0)
w_sent = w.new_zeros(n_sent, max_w, D)
mask = torch.zeros(n_sent, max_w, dtype=torch.bool, device=w.device)
if b_idx:
dev = w.device
bs = torch.tensor(b_idx, device=dev); ks = torch.tensor(k_idx, device=dev)
ls = torch.tensor(l_idx, device=dev); ps = torch.tensor(p_idx, device=dev)
w_sent[ls, ps] = w[bs, ks] # differentiable advanced-index gather
mask[ls, ps] = True
return w_sent, mask, sent_ids
def forward(self, batch):
"""-> (tag_out, arc_scores, rel_scores, word_mask, sent_ids). Loss/gold in the trainer
(it holds the Sentence objects); all trainable submodules run inside this one forward
so DDP sees the whole graph each step."""
w = self._word_vectors(batch)
tag_out = self._tag_heads(w)
w_sent, word_mask, sent_ids = self._regroup(w, batch["slots"])
if w_sent.shape[1] == 0:
return tag_out, None, None, word_mask, sent_ids
arc_scores, rel_scores = self.biaffine(w_sent, word_mask)
return tag_out, arc_scores, rel_scores, word_mask, sent_ids
|