Upload LayoutDmFIDNetV3

config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "LayoutDmFIDNetV3"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_fidnet_v3.LayoutDmFIDNetV3Config",
+    "AutoModel": "modeling_fidnet_v3.LayoutDmFIDNetV3"
+  },
+  "d_model": 256,
+  "id2label": {
+    "0": "text",
+    "1": "title",
+    "2": "list",
+    "3": "table",
+    "4": "figure"
+  },
+  "label2id": {
+    "figure": 4,
+    "list": 2,
+    "table": 3,
+    "text": 0,
+    "title": 1
+  },
+  "max_bbox": 25,
+  "model_type": "layoutdm_fidnet_v3",
+  "nhead": 4,
+  "num_layers": 4,
+  "torch_dtype": "float32",
+  "transformers_version": "4.36.2"
+}

configuration_fidnet_v3.py

@@ -0,0 +1,19 @@
+from transformers.configuration_utils import PretrainedConfig
+
+
+class LayoutDmFIDNetV3Config(PretrainedConfig):
+    model_type = "layoutdm_fidnet_v3"
+
+    def __init__(
+        self,
+        d_model: int = 256,
+        nhead: int = 4,
+        num_layers: int = 4,
+        max_bbox: int = 50,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.d_model = d_model
+        self.nhead = nhead
+        self.num_layers = num_layers
+        self.max_bbox = max_bbox

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f8e2ed0173dbc3942b4f58b0664fa2917ec46bfad5e8b59e4301d2a2ea44c80
+size 11673377

modeling_fidnet_v3.py

@@ -0,0 +1,117 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_fidnet_v3 import LayoutDmFIDNetV3Config
+
+
+@dataclass
+class LayoutDmFIDNetV3Output(object):
+    logit_dict: torch.Tensor
+    logit_cls: torch.Tensor
+    bbox_pred: torch.Tensor
+
+
+class TransformerWithToken(nn.Module):
+    def __init__(self, d_model: int, nhead: int, dim_feedforward: int, num_layers: int):
+        super().__init__()
+
+        self.token = nn.Parameter(torch.randn(1, 1, d_model))
+        token_mask = torch.zeros(1, 1, dtype=torch.bool)
+        self.register_buffer("token_mask", token_mask)
+
+        self.core = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                d_model=d_model,
+                nhead=nhead,
+                dim_feedforward=dim_feedforward,
+            ),
+            num_layers=num_layers,
+        )
+
+    def forward(self, x, src_key_padding_mask):
+        # x: [N, B, E]
+        # padding_mask: [B, N]
+        #   `False` for valid values
+        #   `True` for padded values
+
+        B = x.size(1)
+
+        token = self.token.expand(-1, B, -1)
+        x = torch.cat([token, x], dim=0)
+
+        token_mask = self.token_mask.expand(B, -1)
+        padding_mask = torch.cat([token_mask, src_key_padding_mask], dim=1)
+
+        x = self.core(x, src_key_padding_mask=padding_mask)
+
+        return x
+
+
+class LayoutDmFIDNetV3(PreTrainedModel):
+    config_class = LayoutDmFIDNetV3Config
+
+    def __init__(self, config: LayoutDmFIDNetV3Config):
+        super().__init__(config)
+        self.config = config
+
+        # encoder
+        self.emb_label = nn.Embedding(config.num_labels, config.d_model)
+        self.fc_bbox = nn.Linear(4, config.d_model)
+        self.enc_fc_in = nn.Linear(config.d_model * 2, config.d_model)
+
+        self.enc_transformer = TransformerWithToken(
+            d_model=config.d_model,
+            dim_feedforward=config.d_model // 2,
+            nhead=config.nhead,
+            num_layers=config.num_layers,
+        )
+
+        self.fc_out_disc = nn.Linear(config.d_model, 1)
+
+        # decoder
+        self.pos_token = nn.Parameter(torch.rand(config.max_bbox, 1, config.d_model))
+        self.dec_fc_in = nn.Linear(config.d_model * 2, config.d_model)
+
+        te = nn.TransformerEncoderLayer(
+            d_model=config.d_model,
+            nhead=config.nhead,
+            dim_feedforward=config.d_model // 2,
+        )
+        self.dec_transformer = nn.TransformerEncoder(te, num_layers=config.num_layers)
+
+        self.fc_out_cls = nn.Linear(config.d_model, config.num_labels)
+        self.fc_out_bbox = nn.Linear(config.d_model, 4)
+
+    def extract_features(self, bbox, label, padding_mask):
+        b = self.fc_bbox(bbox)
+        l = self.emb_label(label)
+        x = self.enc_fc_in(torch.cat([b, l], dim=-1))
+        x = torch.relu(x).permute(1, 0, 2)
+        x = self.enc_transformer(x, padding_mask)
+        return x[0]
+
+    def forward(self, bbox, label, padding_mask):
+        B, N, _ = bbox.size()
+        x = self.extract_features(bbox, label, padding_mask)
+
+        logit_disc = self.fc_out_disc(x).squeeze(-1)
+
+        x = x.unsqueeze(0).expand(N, -1, -1)
+        t = self.pos_token[:N].expand(-1, B, -1)
+        x = torch.cat([x, t], dim=-1)
+        x = torch.relu(self.dec_fc_in(x))
+
+        x = self.dec_transformer(x, src_key_padding_mask=padding_mask)
+        # x = x.permute(1, 0, 2)[~padding_mask]
+        x = x.permute(1, 0, 2)
+
+        # logit_cls: [B, N, L]    bbox_pred: [B, N, 4]
+        logit_cls = self.fc_out_cls(x)
+        bbox_pred = torch.sigmoid(self.fc_out_bbox(x))
+
+        return LayoutDmFIDNetV3Output(
+            logit_disc=logit_disc, logit_cls=logit_cls, bbox_pred=bbox_pred
+        )