| """Field and relationship model over a page's text segments. |
| |
| Each segment is embedded from its characters, geometry and shape flags. A small Transformer with a |
| learned 2D spatial attention bias mixes context across the page, then three heads predict: |
| * a field/structure label per segment, |
| * pairwise relations (same_row, key_value), |
| * the document type (receipt / invoice / general). |
| """ |
|
|
| import torch |
| from torch import nn |
|
|
| from ..charset import EXTRACTOR_CLASSES |
| from ..features import NUM_FEATS |
| from ..schema import DOC_TYPES, FIELD_LABELS, REL_TYPES |
| from .common import TransformerBlock |
|
|
| NUM_PAIR_FEATS = 8 |
|
|
|
|
| def pair_features(rel_boxes): |
| """rel_boxes: (B, N, 4) boxes in median-text-height units -> (B, N, N, 8) pairwise geometry.""" |
| x0, y0, x1, y1 = rel_boxes.unbind(-1) |
| cx, cy = (x0 + x1) / 2, (y0 + y1) / 2 |
| h = (y1 - y0).clamp(min=1e-3) |
| w = (x1 - x0).clamp(min=1e-3) |
|
|
| def d(a, b): |
| return b[:, None, :] - a[:, :, None] |
|
|
| v_ov = (torch.minimum(y1[:, :, None], y1[:, None, :]) - torch.maximum(y0[:, :, None], y0[:, None, :])).clamp(min=0) |
| h_ov = (torch.minimum(x1[:, :, None], x1[:, None, :]) - torch.maximum(x0[:, :, None], x0[:, None, :])).clamp(min=0) |
| v_ov = v_ov / torch.minimum(h[:, :, None], h[:, None, :]) |
| h_ov = h_ov / torch.minimum(w[:, :, None], w[:, None, :]) |
| return torch.stack([ |
| d(cy, cy).clamp(-20, 20) / 20, |
| d(cx, cx).clamp(-40, 40) / 40, |
| v_ov, |
| h_ov, |
| d(x1, x0).clamp(-40, 40) / 40, |
| (v_ov > 0.5).to(rel_boxes.dtype), |
| d(x1, x1).abs().clamp(max=20) / 20, |
| d(x0, x0).abs().clamp(max=20) / 20, |
| ], dim=-1) |
|
|
|
|
| class CharEncoder(nn.Module): |
| def __init__(self, d, emb=48): |
| super().__init__() |
| self.emb = nn.Embedding(EXTRACTOR_CLASSES, emb, padding_idx=0) |
| self.conv1 = nn.Conv1d(emb, d, 3, padding=1) |
| self.conv2 = nn.Conv1d(d, d, 3, padding=1, groups=4) |
| self.act = nn.GELU() |
|
|
| def forward(self, ids): |
| """ids: (B, N, L) -> (B, N, d) via masked max-pool over characters.""" |
| B, N, L = ids.shape |
| x = self.emb(ids.reshape(B * N, L)).transpose(1, 2) |
| x = self.act(self.conv2(self.act(self.conv1(x)))) |
| mask = (ids.reshape(B * N, 1, L) == 0) |
| x = x.masked_fill(mask, -1e4).amax(-1) |
| empty = mask.all(-1) |
| x = torch.where(empty, torch.zeros_like(x), x) |
| return x.reshape(B, N, -1) |
|
|
|
|
| class Extractor(nn.Module): |
| def __init__(self, d=112, layers=3, heads=4, ff=256, rel_dim=48, dropout=0.1): |
| super().__init__() |
| self.heads = heads |
| self.chars = CharEncoder(d) |
| self.feat_proj = nn.Sequential(nn.Linear(NUM_FEATS, d), nn.GELU(), nn.Linear(d, d)) |
| self.ln_in = nn.LayerNorm(d) |
| self.attn_bias = nn.Sequential(nn.Linear(NUM_PAIR_FEATS, 32), nn.GELU(), nn.Linear(32, heads)) |
| self.blocks = nn.ModuleList(TransformerBlock(d, heads, ff, dropout) for _ in range(layers)) |
| self.ln = nn.LayerNorm(d) |
| self.field_head = nn.Linear(d, len(FIELD_LABELS)) |
| self.rel_q = nn.Linear(d, rel_dim * len(REL_TYPES)) |
| self.rel_k = nn.Linear(d, rel_dim * len(REL_TYPES)) |
| self.rel_bias = nn.Sequential(nn.Linear(NUM_PAIR_FEATS, 32), nn.GELU(), nn.Linear(32, len(REL_TYPES))) |
| self.doc_pool = nn.Linear(d, 1) |
| self.doc_head = nn.Linear(d, len(DOC_TYPES)) |
| self.rel_dim = rel_dim |
|
|
| def forward(self, char_ids, feats, rel_boxes, valid): |
| """char_ids (B,N,L) int64, feats (B,N,F), rel_boxes (B,N,4), valid (B,N) float 1/0. |
| |
| Returns field_logits (B,N,K), rel_logits (B,R,N,N), doc_logits (B,D). |
| """ |
| B, N, _ = feats.shape |
| x = self.ln_in(self.chars(char_ids) + self.feat_proj(feats)) |
| pf = pair_features(rel_boxes) |
| key_pad = (1.0 - valid)[:, None, None, :] * -1e4 |
| bias = self.attn_bias(pf).permute(0, 3, 1, 2) + key_pad |
| for blk in self.blocks: |
| x = blk(x, bias) |
| x = self.ln(x) |
|
|
| field_logits = self.field_head(x) |
| R, r = len(REL_TYPES), self.rel_dim |
| q = self.rel_q(x).reshape(B, N, R, r).permute(0, 2, 1, 3) |
| k = self.rel_k(x).reshape(B, N, R, r).permute(0, 2, 1, 3) |
| rel_logits = (q @ k.transpose(-2, -1)) / (r ** 0.5) + self.rel_bias(pf).permute(0, 3, 1, 2) |
|
|
| w = self.doc_pool(x).squeeze(-1) + (1.0 - valid) * -1e4 |
| pooled = (w.softmax(-1)[..., None] * x).sum(1) |
| return field_logits, rel_logits, self.doc_head(pooled) |
|
|
