File size: 3,095 Bytes
5952424 | 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 | """Smoke test for the joint tagger+parser pipeline's tensor plumbing: a tiny, randomly
initialized encoder + JointModel, fed a synthetic mini-batch, produces correctly-shaped
tag/lemma/UPOS/arc/label predictions end to end. Does NOT test tagging/parsing quality
(real weights, real treebank) -- only that packing/pooling/scalar-mix/biaffine survive a
refactor without shape or index errors."""
import sys
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model.char_bert import CharBertConfig
from tagger.backbone import CharBertWithHidden
from tagger.edits import LabelVocab
from tagger.model import TaggerConfig
from parser.biaffine import ParserConfig
from parser.joint_model import JointModel
def _tiny_vocab():
return LabelVocab(
scripts=[("", "")],
xpos_alpha=[["-", "n", "v"] for _ in range(9)],
tags=["n-s---fa-", "v2spma---"],
upos=["NOUN", "VERB"],
lex_ft={}, lex_f={}, nongreek={},
)
def test_joint_model_forward_pass():
enc_cfg = CharBertConfig(d_model=32, n_heads=4, depth=2, char_window=8, attn_impl="sdpa")
encoder = CharBertWithHidden(enc_cfg)
encoder.return_layers = True
encoder.eval()
vocab = _tiny_vocab()
tcfg = TaggerConfig(pool="mean", use_cap=False, scalar_mix=True, w_flat=0.0)
pcfg = ParserConfig(d_arc=16, d_rel=8, dropout=0.0, n_labels=4)
W = 6 # max words per row
model = JointModel(encoder, vocab, tcfg, pcfg, W=W)
model.eval()
# synthetic mini-batch: 2 short "sentences" of 3 words each, packed into one row
T = 20
B = 1
input_ids = torch.randint(0, 24, (B, T))
boundary = torch.zeros(B, T, dtype=torch.long)
boundary[:, [4, 9, 14, 19]] = 1 # word ends
boundary[:, 19] = 2 # sentence end
dia = torch.zeros(B, T, dtype=torch.long)
punct = torch.zeros(B, T, dtype=torch.long)
seg_id = torch.zeros(B, T, dtype=torch.long)
# word_id: -1 for non-final char positions is not required by pool_words (mean over all
# positions sharing a word id); assign each 5-char span to one word slot 0..3
word_id = torch.tensor([[w for w in range(4) for _ in range(5)]], dtype=torch.long)
# one sentence (id 0) occupying word slots 0..3 of this single packed row
slots = [[(0, 0), (0, 1), (0, 2), (0, 3)]]
batch = dict(input_ids=input_ids, boundary=boundary, dia=dia, punct=punct,
seg_id=seg_id, word_id=word_id, slots=slots)
with torch.no_grad():
tag_out, arc_scores, rel_scores, word_mask, sent_ids = model(batch)
# tagger heads: factored XPOS (list of 9 per-position logit tensors), UPOS, lemma-script
assert len(tag_out["xpos"]) == 9
assert tag_out["xpos"][0].shape[:2] == (B, W)
assert tag_out["upos"].shape[:2] == (B, W)
assert tag_out["script"].shape[:2] == (B, W)
# biaffine arc/label scores: one sentence, 4 real words + 1 root column
assert arc_scores.shape[0] == 1
assert arc_scores.shape[1] == word_mask.shape[1]
assert rel_scores.shape[0] == 1
assert sent_ids == [0]
|