import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
from einops import rearrange
from torch import Tensor
class PositionalEncoding(nn.Module):
def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
self.max_len = max_len
self.d_model = d_model
position = torch.arange(max_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
pe = torch.zeros(1, max_len, d_model)
pe[0, :, 0::2] = torch.sin(position * div_term)
pe[0, :, 1::2] = torch.cos(position * div_term)
self.register_buffer("pe", pe)
def forward(self) -> Tensor:
x = self.pe[0, : self.max_len]
return self.dropout(x).unsqueeze(0)
class ResNetFeatureExtractor(nn.Module):
def __init__(self, hidden_dim = 512):
super().__init__()
# Making the resnet 50 model, which was used in the docformer for the purpose of visual feature extraction
resnet50 = models.resnet50(pretrained=False)
modules = list(resnet50.children())[:-2]
self.resnet50 = nn.Sequential(*modules)
# Applying convolution and linear layer
self.conv1 = nn.Conv2d(2048, 768, 1)
self.relu1 = F.relu
self.linear1 = nn.Linear(192, hidden_dim)
def forward(self, x):
x = self.resnet50(x)
x = self.conv1(x)
x = self.relu1(x)
x = rearrange(x, "b e w h -> b e (w h)") # b -> batch, e -> embedding dim, w -> width, h -> height
x = self.linear1(x)
x = rearrange(x, "b e s -> b s e") # b -> batch, e -> embedding dim, s -> sequence length
return x
class DocFormerEmbeddings(nn.Module):
"""Construct the embeddings from word, position and token_type embeddings."""
def __init__(self, config):
super(DocFormerEmbeddings, self).__init__()
self.config = config
self.position_embeddings_v = PositionalEncoding(
d_model=config["hidden_size"],
dropout=0.1,
max_len=config["max_position_embeddings"],
)
self.x_topleft_position_embeddings_v = nn.Embedding(config["max_2d_position_embeddings"], config["coordinate_size"])
self.x_bottomright_position_embeddings_v = nn.Embedding(config["max_2d_position_embeddings"], config["coordinate_size"])
self.w_position_embeddings_v = nn.Embedding(config["max_2d_position_embeddings"], config["shape_size"])
self.x_topleft_distance_to_prev_embeddings_v = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.x_bottomleft_distance_to_prev_embeddings_v = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.x_topright_distance_to_prev_embeddings_v = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.x_bottomright_distance_to_prev_embeddings_v = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.x_centroid_distance_to_prev_embeddings_v = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.y_topleft_position_embeddings_v = nn.Embedding(config["max_2d_position_embeddings"], config["coordinate_size"])
self.y_bottomright_position_embeddings_v = nn.Embedding(config["max_2d_position_embeddings"], config["coordinate_size"])
self.h_position_embeddings_v = nn.Embedding(config["max_2d_position_embeddings"], config["shape_size"])
self.y_topleft_distance_to_prev_embeddings_v = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.y_bottomleft_distance_to_prev_embeddings_v = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.y_topright_distance_to_prev_embeddings_v = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.y_bottomright_distance_to_prev_embeddings_v = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.y_centroid_distance_to_prev_embeddings_v = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.position_embeddings_t = PositionalEncoding(
d_model=config["hidden_size"],
dropout=0.1,
max_len=config["max_position_embeddings"],
)
self.x_topleft_position_embeddings_t = nn.Embedding(config["max_2d_position_embeddings"], config["coordinate_size"])
self.x_bottomright_position_embeddings_t = nn.Embedding(config["max_2d_position_embeddings"], config["coordinate_size"])
self.w_position_embeddings_t = nn.Embedding(config["max_2d_position_embeddings"], config["shape_size"])
self.x_topleft_distance_to_prev_embeddings_t = nn.Embedding(2*config["max_2d_position_embeddings"]+1, config["shape_size"])
self.x_bottomleft_distance_to_prev_embeddings_t = nn.Embedding(2*config["max_2d_position_embeddings"]+1, config["shape_size"])
self.x_topright_distance_to_prev_embeddings_t = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.x_bottomright_distance_to_prev_embeddings_t = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.x_centroid_distance_to_prev_embeddings_t = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.y_topleft_position_embeddings_t = nn.Embedding(config["max_2d_position_embeddings"], config["coordinate_size"])
self.y_bottomright_position_embeddings_t = nn.Embedding(config["max_2d_position_embeddings"], config["coordinate_size"])
self.h_position_embeddings_t = nn.Embedding(config["max_2d_position_embeddings"], config["shape_size"])
self.y_topleft_distance_to_prev_embeddings_t = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.y_bottomleft_distance_to_prev_embeddings_t = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.y_topright_distance_to_prev_embeddings_t = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.y_bottomright_distance_to_prev_embeddings_t = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.y_centroid_distance_to_prev_embeddings_t = nn.Embedding(2*config["max_2d_position_embeddings"] + 1, config["shape_size"])
self.LayerNorm = nn.LayerNorm(config["hidden_size"], eps=config["layer_norm_eps"])
self.dropout = nn.Dropout(config["hidden_dropout_prob"])
def forward(self, x_feature, y_feature):
"""
Arguments:
x_features of shape, (batch size, seq_len, 8)
y_features of shape, (batch size, seq_len, 8)
Outputs:
(V-bar-s, T-bar-s) of shape (batch size, 512,768),(batch size, 512,768)
What are the features:
0 -> top left x/y
1 -> bottom right x/y
2 -> width/height
3 -> diff top left x/y
4 -> diff bottom left x/y
5 -> diff top right x/y
6 -> diff bottom right x/y
7 -> centroids diff x/y
"""
batch, seq_len = x_feature.shape[:-1]
hidden_size = self.config["hidden_size"]
num_feat = x_feature.shape[-1]
sub_dim = hidden_size // num_feat
# Clamping and adding a bias for handling negative values
x_feature[:,:,3:] = torch.clamp(x_feature[:,:,3:],-self.config["max_2d_position_embeddings"],self.config["max_2d_position_embeddings"])
x_feature[:,:,3:]+= self.config["max_2d_position_embeddings"]
y_feature[:,:,3:] = torch.clamp(y_feature[:,:,3:],-self.config["max_2d_position_embeddings"],self.config["max_2d_position_embeddings"])
y_feature[:,:,3:]+= self.config["max_2d_position_embeddings"]
x_topleft_position_embeddings_v = self.x_topleft_position_embeddings_v(x_feature[:,:,0])
x_bottomright_position_embeddings_v = self.x_bottomright_position_embeddings_v(x_feature[:,:,1])
w_position_embeddings_v = self.w_position_embeddings_v(x_feature[:,:,2])
x_topleft_distance_to_prev_embeddings_v = self.x_topleft_distance_to_prev_embeddings_v(x_feature[:,:,3])
x_bottomleft_distance_to_prev_embeddings_v = self.x_bottomleft_distance_to_prev_embeddings_v(x_feature[:,:,4])
x_topright_distance_to_prev_embeddings_v = self.x_topright_distance_to_prev_embeddings_v(x_feature[:,:,5])
x_bottomright_distance_to_prev_embeddings_v = self.x_bottomright_distance_to_prev_embeddings_v(x_feature[:,:,6])
x_centroid_distance_to_prev_embeddings_v = self.x_centroid_distance_to_prev_embeddings_v(x_feature[:,:,7])
x_calculated_embedding_v = torch.cat(
[
x_topleft_position_embeddings_v,
x_bottomright_position_embeddings_v,
w_position_embeddings_v,
x_topleft_distance_to_prev_embeddings_v,
x_bottomleft_distance_to_prev_embeddings_v,
x_topright_distance_to_prev_embeddings_v,
x_bottomright_distance_to_prev_embeddings_v ,
x_centroid_distance_to_prev_embeddings_v
],
dim = -1
)
y_topleft_position_embeddings_v = self.y_topleft_position_embeddings_v(y_feature[:,:,0])
y_bottomright_position_embeddings_v = self.y_bottomright_position_embeddings_v(y_feature[:,:,1])
h_position_embeddings_v = self.h_position_embeddings_v(y_feature[:,:,2])
y_topleft_distance_to_prev_embeddings_v = self.y_topleft_distance_to_prev_embeddings_v(y_feature[:,:,3])
y_bottomleft_distance_to_prev_embeddings_v = self.y_bottomleft_distance_to_prev_embeddings_v(y_feature[:,:,4])
y_topright_distance_to_prev_embeddings_v = self.y_topright_distance_to_prev_embeddings_v(y_feature[:,:,5])
y_bottomright_distance_to_prev_embeddings_v = self.y_bottomright_distance_to_prev_embeddings_v(y_feature[:,:,6])
y_centroid_distance_to_prev_embeddings_v = self.y_centroid_distance_to_prev_embeddings_v(y_feature[:,:,7])
x_calculated_embedding_v = torch.cat(
[
x_topleft_position_embeddings_v,
x_bottomright_position_embeddings_v,
w_position_embeddings_v,
x_topleft_distance_to_prev_embeddings_v,
x_bottomleft_distance_to_prev_embeddings_v,
x_topright_distance_to_prev_embeddings_v,
x_bottomright_distance_to_prev_embeddings_v ,
x_centroid_distance_to_prev_embeddings_v
],
dim = -1
)
y_calculated_embedding_v = torch.cat(
[
y_topleft_position_embeddings_v,
y_bottomright_position_embeddings_v,
h_position_embeddings_v,
y_topleft_distance_to_prev_embeddings_v,
y_bottomleft_distance_to_prev_embeddings_v,
y_topright_distance_to_prev_embeddings_v,
y_bottomright_distance_to_prev_embeddings_v ,
y_centroid_distance_to_prev_embeddings_v
],
dim = -1
)
v_bar_s = x_calculated_embedding_v + y_calculated_embedding_v + self.position_embeddings_v()
x_topleft_position_embeddings_t = self.x_topleft_position_embeddings_t(x_feature[:,:,0])
x_bottomright_position_embeddings_t = self.x_bottomright_position_embeddings_t(x_feature[:,:,1])
w_position_embeddings_t = self.w_position_embeddings_t(x_feature[:,:,2])
x_topleft_distance_to_prev_embeddings_t = self.x_topleft_distance_to_prev_embeddings_t(x_feature[:,:,3])
x_bottomleft_distance_to_prev_embeddings_t = self.x_bottomleft_distance_to_prev_embeddings_t(x_feature[:,:,4])
x_topright_distance_to_prev_embeddings_t = self.x_topright_distance_to_prev_embeddings_t(x_feature[:,:,5])
x_bottomright_distance_to_prev_embeddings_t = self.x_bottomright_distance_to_prev_embeddings_t(x_feature[:,:,6])
x_centroid_distance_to_prev_embeddings_t = self.x_centroid_distance_to_prev_embeddings_t(x_feature[:,:,7])
x_calculated_embedding_t = torch.cat(
[
x_topleft_position_embeddings_t,
x_bottomright_position_embeddings_t,
w_position_embeddings_t,
x_topleft_distance_to_prev_embeddings_t,
x_bottomleft_distance_to_prev_embeddings_t,
x_topright_distance_to_prev_embeddings_t,
x_bottomright_distance_to_prev_embeddings_t ,
x_centroid_distance_to_prev_embeddings_t
],
dim = -1
)
y_topleft_position_embeddings_t = self.y_topleft_position_embeddings_t(y_feature[:,:,0])
y_bottomright_position_embeddings_t = self.y_bottomright_position_embeddings_t(y_feature[:,:,1])
h_position_embeddings_t = self.h_position_embeddings_t(y_feature[:,:,2])
y_topleft_distance_to_prev_embeddings_t = self.y_topleft_distance_to_prev_embeddings_t(y_feature[:,:,3])
y_bottomleft_distance_to_prev_embeddings_t = self.y_bottomleft_distance_to_prev_embeddings_t(y_feature[:,:,4])
y_topright_distance_to_prev_embeddings_t = self.y_topright_distance_to_prev_embeddings_t(y_feature[:,:,5])
y_bottomright_distance_to_prev_embeddings_t = self.y_bottomright_distance_to_prev_embeddings_t(y_feature[:,:,6])
y_centroid_distance_to_prev_embeddings_t = self.y_centroid_distance_to_prev_embeddings_t(y_feature[:,:,7])
x_calculated_embedding_t = torch.cat(
[
x_topleft_position_embeddings_t,
x_bottomright_position_embeddings_t,
w_position_embeddings_t,
x_topleft_distance_to_prev_embeddings_t,
x_bottomleft_distance_to_prev_embeddings_t,
x_topright_distance_to_prev_embeddings_t,
x_bottomright_distance_to_prev_embeddings_t ,
x_centroid_distance_to_prev_embeddings_t
],
dim = -1
)
y_calculated_embedding_t = torch.cat(
[
y_topleft_position_embeddings_t,
y_bottomright_position_embeddings_t,
h_position_embeddings_t,
y_topleft_distance_to_prev_embeddings_t,
y_bottomleft_distance_to_prev_embeddings_t,
y_topright_distance_to_prev_embeddings_t,
y_bottomright_distance_to_prev_embeddings_t ,
y_centroid_distance_to_prev_embeddings_t
],
dim = -1
)
t_bar_s = x_calculated_embedding_t + y_calculated_embedding_t + self.position_embeddings_t()
return v_bar_s, t_bar_s
# fmt: off
class PreNorm(nn.Module):
def __init__(self, dim, fn):
# Fig 1: http://proceedings.mlr.press/v119/xiong20b/xiong20b.pdf
super().__init__()
self.norm = nn.LayerNorm(dim)
self.fn = fn
def forward(self, x, **kwargs):
return self.fn(self.norm(x), **kwargs)
class PreNormAttn(nn.Module):
def __init__(self, dim, fn):
# Fig 1: http://proceedings.mlr.press/v119/xiong20b/xiong20b.pdf
super().__init__()
self.norm_t_bar = nn.LayerNorm(dim)
self.norm_v_bar = nn.LayerNorm(dim)
self.norm_t_bar_s = nn.LayerNorm(dim)
self.norm_v_bar_s = nn.LayerNorm(dim)
self.fn = fn
def forward(self, t_bar, v_bar, t_bar_s, v_bar_s, **kwargs):
return self.fn(self.norm_t_bar(t_bar),
self.norm_v_bar(v_bar),
self.norm_t_bar_s(t_bar_s),
self.norm_v_bar_s(v_bar_s), **kwargs)
class FeedForward(nn.Module):
def __init__(self, dim, hidden_dim, dropout=0.):
super().__init__()
self.net = nn.Sequential(
nn.Linear(dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, dim),
nn.Dropout(dropout)
)
def forward(self, x):
return self.net(x)
class RelativePosition(nn.Module):
def __init__(self, num_units, max_relative_position, max_seq_length):
super().__init__()
self.num_units = num_units
self.max_relative_position = max_relative_position
self.embeddings_table = nn.Parameter(torch.Tensor(max_relative_position * 2 + 1, num_units))
self.max_length = max_seq_length
range_vec_q = torch.arange(max_seq_length)
range_vec_k = torch.arange(max_seq_length)
distance_mat = range_vec_k[None, :] - range_vec_q[:, None]
distance_mat_clipped = torch.clamp(distance_mat, -self.max_relative_position, self.max_relative_position)
final_mat = distance_mat_clipped + self.max_relative_position
self.final_mat = torch.LongTensor(final_mat)
nn.init.xavier_uniform_(self.embeddings_table)
def forward(self, length_q, length_k):
embeddings = self.embeddings_table[self.final_mat[:length_q, :length_k]]
return embeddings
class MultiModalAttentionLayer(nn.Module):
def __init__(self, embed_dim, n_heads, max_relative_position, max_seq_length, dropout):
super().__init__()
assert embed_dim % n_heads == 0
self.embed_dim = embed_dim
self.n_heads = n_heads
self.head_dim = embed_dim // n_heads
self.relative_positions_text = RelativePosition(self.head_dim, max_relative_position, max_seq_length)
self.relative_positions_img = RelativePosition(self.head_dim, max_relative_position, max_seq_length)
# text qkv embeddings
self.fc_k_text = nn.Linear(embed_dim, embed_dim)
self.fc_q_text = nn.Linear(embed_dim, embed_dim)
self.fc_v_text = nn.Linear(embed_dim, embed_dim)
# image qkv embeddings
self.fc_k_img = nn.Linear(embed_dim, embed_dim)
self.fc_q_img = nn.Linear(embed_dim, embed_dim)
self.fc_v_img = nn.Linear(embed_dim, embed_dim)
# spatial qk embeddings (shared for visual and text)
self.fc_k_spatial = nn.Linear(embed_dim, embed_dim)
self.fc_q_spatial = nn.Linear(embed_dim, embed_dim)
self.dropout = nn.Dropout(dropout)
self.to_out = nn.Sequential(
nn.Linear(embed_dim, embed_dim),
nn.Dropout(dropout)
)
self.scale = embed_dim**0.5
def forward(self, text_feat, img_feat, text_spatial_feat, img_spatial_feat):
text_feat = text_feat
img_feat = img_feat
text_spatial_feat = text_spatial_feat
img_spatial_feat = img_spatial_feat
seq_length = text_feat.shape[1]
# self attention of text
# b -> batch, t -> time steps (l -> length has same meaning), head -> # of heads, k -> head dim.
key_text_nh = rearrange(self.fc_k_text(text_feat), 'b t (head k) -> head b t k', head=self.n_heads)
query_text_nh = rearrange(self.fc_q_text(text_feat), 'b l (head k) -> head b l k', head=self.n_heads)
value_text_nh = rearrange(self.fc_v_text(text_feat), 'b t (head k) -> head b t k', head=self.n_heads)
dots_text = torch.einsum('hblk,hbtk->hblt', query_text_nh, key_text_nh)
dots_text = dots_text/ self.scale
# 1D relative positions (query, key)
rel_pos_embed_text = self.relative_positions_text(seq_length, seq_length)
rel_pos_key_text = torch.einsum('bhrd,lrd->bhlr', key_text_nh, rel_pos_embed_text)
rel_pos_query_text = torch.einsum('bhld,lrd->bhlr', query_text_nh, rel_pos_embed_text)
# shared spatial <-> text hidden features
key_spatial_text = self.fc_k_spatial(text_spatial_feat)
query_spatial_text = self.fc_q_spatial(text_spatial_feat)
key_spatial_text_nh = rearrange(key_spatial_text, 'b t (head k) -> head b t k', head=self.n_heads)
query_spatial_text_nh = rearrange(query_spatial_text, 'b l (head k) -> head b l k', head=self.n_heads)
dots_text_spatial = torch.einsum('hblk,hbtk->hblt', query_spatial_text_nh, key_spatial_text_nh)
dots_text_spatial = dots_text_spatial/ self.scale
# Line 38 of pseudo-code
text_attn_scores = dots_text + rel_pos_key_text + rel_pos_query_text + dots_text_spatial
# self-attention of image
key_img_nh = rearrange(self.fc_k_img(img_feat), 'b t (head k) -> head b t k', head=self.n_heads)
query_img_nh = rearrange(self.fc_q_img(img_feat), 'b l (head k) -> head b l k', head=self.n_heads)
value_img_nh = rearrange(self.fc_v_img(img_feat), 'b t (head k) -> head b t k', head=self.n_heads)
dots_img = torch.einsum('hblk,hbtk->hblt', query_img_nh, key_img_nh)
dots_img = dots_img/ self.scale
# 1D relative positions (query, key)
rel_pos_embed_img = self.relative_positions_img(seq_length, seq_length)
rel_pos_key_img = torch.einsum('bhrd,lrd->bhlr', key_img_nh, rel_pos_embed_text)
rel_pos_query_img = torch.einsum('bhld,lrd->bhlr', query_img_nh, rel_pos_embed_text)
# shared spatial <-> image features
key_spatial_img = self.fc_k_spatial(img_spatial_feat)
query_spatial_img = self.fc_q_spatial(img_spatial_feat)
key_spatial_img_nh = rearrange(key_spatial_img, 'b t (head k) -> head b t k', head=self.n_heads)
query_spatial_img_nh = rearrange(query_spatial_img, 'b l (head k) -> head b l k', head=self.n_heads)
dots_img_spatial = torch.einsum('hblk,hbtk->hblt', query_spatial_img_nh, key_spatial_img_nh)
dots_img_spatial = dots_img_spatial/ self.scale
# Line 59 of pseudo-code
img_attn_scores = dots_img + rel_pos_key_img + rel_pos_query_img + dots_img_spatial
text_attn_probs = self.dropout(torch.softmax(text_attn_scores, dim=-1))
img_attn_probs = self.dropout(torch.softmax(img_attn_scores, dim=-1))
text_context = torch.einsum('hblt,hbtv->hblv', text_attn_probs, value_text_nh)
img_context = torch.einsum('hblt,hbtv->hblv', img_attn_probs, value_img_nh)
context = text_context + img_context
embeddings = rearrange(context, 'head b t d -> b t (head d)')
return self.to_out(embeddings)
class DocFormerEncoder(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.layers = nn.ModuleList([])
for _ in range(config['num_hidden_layers']):
encoder_block = nn.ModuleList([
PreNormAttn(config['hidden_size'],
MultiModalAttentionLayer(config['hidden_size'],
config['num_attention_heads'],
config['max_relative_positions'],
config['max_position_embeddings'],
config['hidden_dropout_prob'],
)
),
PreNorm(config['hidden_size'],
FeedForward(config['hidden_size'],
config['hidden_size'] * config['intermediate_ff_size_factor'],
dropout=config['hidden_dropout_prob']))
])
self.layers.append(encoder_block)
def forward(
self,
text_feat, # text feat or output from last encoder block
img_feat,
text_spatial_feat,
img_spatial_feat,
):
# Fig 1 encoder part (skip conn for both attn & FF): https://arxiv.org/abs/1706.03762
# TODO: ensure 1st skip conn (var "skip") in such a multimodal setting makes sense (most likely does)
for attn, ff in self.layers:
skip = text_feat + img_feat + text_spatial_feat + img_spatial_feat
x = attn(text_feat, img_feat, text_spatial_feat, img_spatial_feat) + skip
x = ff(x) + x
text_feat = x
return x
class LanguageFeatureExtractor(nn.Module):
def __init__(self):
super().__init__()
from transformers import LayoutLMForTokenClassification
layoutlm_dummy = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlm-base-uncased", num_labels=1)
self.embedding_vector = nn.Embedding.from_pretrained(layoutlm_dummy.layoutlm.embeddings.word_embeddings.weight)
def forward(self, x):
return self.embedding_vector(x)
class ExtractFeatures(nn.Module):
'''
Inputs: dictionary
Output: v_bar, t_bar, v_bar_s, t_bar_s
'''
def __init__(self, config):
super().__init__()
self.visual_feature = ResNetFeatureExtractor(hidden_dim = config['max_position_embeddings'])
self.language_feature = LanguageFeatureExtractor()
self.spatial_feature = DocFormerEmbeddings(config)
def forward(self, encoding):
image = encoding['resized_scaled_img']
language = encoding['input_ids']
x_feature = encoding['x_features']
y_feature = encoding['y_features']
v_bar = self.visual_feature(image)
t_bar = self.language_feature(language)
v_bar_s, t_bar_s = self.spatial_feature(x_feature, y_feature)
return v_bar, t_bar, v_bar_s, t_bar_s
class DocFormer(nn.Module):
'''
Easy boiler plate, because this model will just take as an input, the dictionary which is obtained from create_features function
'''
def __init__(self, config):
super().__init__()
self.config = config
self.extract_feature = ExtractFeatures(config)
self.encoder = DocFormerEncoder(config)
self.dropout = nn.Dropout(config['hidden_dropout_prob'])
def forward(self, x ,use_tdi=False):
v_bar, t_bar, v_bar_s, t_bar_s = self.extract_feature(x,use_tdi)
features = {'v_bar': v_bar, 't_bar': t_bar, 'v_bar_s': v_bar_s, 't_bar_s': t_bar_s}
output = self.encoder(features['t_bar'], features['v_bar'], features['t_bar_s'], features['v_bar_s'])
output = self.dropout(output)
return output