Upload modeling.py

Added modeling file

modeling.py

@@ -0,0 +1,545 @@
+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
+
+
+
+
+
+