File size: 4,600 Bytes
4dce182 e24bf15 4dce182 e24bf15 4dce182 | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """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): # b_j - a_i
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) # (B,N,N,P)
key_pad = (1.0 - valid)[:, None, None, :] * -1e4 # (B,1,1,N)
bias = self.attn_bias(pf).permute(0, 3, 1, 2) + key_pad # (B,H,N,N)
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)
|