| """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) |
| self.biaffine = BiaffineHead(encoder.cfg.d_model, pcfg) |
| self.W = W |
|
|
| |
| 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"]]]) |
| 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] |
| 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 |
|
|
