diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py
index 1217a6eb..fac55716 100644
--- a/modelscope/metainfo.py
+++ b/modelscope/metainfo.py
@@ -127,6 +127,7 @@ class Models(object):
multi_stage_diffusion = 'multi-stage-diffusion-text-to-image-synthesis'
team = 'team-multi-modal-similarity'
video_clip = 'video-clip-multi-modal-embedding'
+ vldoc = 'vldoc'
hitea = 'hitea'
# science models
@@ -336,6 +337,7 @@ class Pipelines(object):
image_text_retrieval = 'image-text-retrieval'
ofa_ocr_recognition = 'ofa-ocr-recognition'
ofa_asr = 'ofa-asr'
+ document_vl_embedding = 'document-vl-embedding'
video_captioning = 'video-captioning'
video_question_answering = 'video-question-answering'
@@ -464,6 +466,7 @@ class Preprocessors(object):
ofa_tasks_preprocessor = 'ofa-tasks-preprocessor'
clip_preprocessor = 'clip-preprocessor'
mplug_tasks_preprocessor = 'mplug-tasks-preprocessor'
+ vldoc_preprocessor = 'vldoc-preprocessor'
hitea_tasks_preprocessor = 'hitea-tasks-preprocessor'
# science preprocessor
diff --git a/modelscope/models/multi_modal/__init__.py b/modelscope/models/multi_modal/__init__.py
index ba30f7b7..4edf6212 100644
--- a/modelscope/models/multi_modal/__init__.py
+++ b/modelscope/models/multi_modal/__init__.py
@@ -16,6 +16,7 @@ if TYPE_CHECKING:
OfaForTextToImageSynthesis
from .multi_stage_diffusion import \
MultiStageDiffusionForTextToImageSynthesis
+ from .vldoc import VLDocForDocVLEmbedding
else:
_import_structure = {
@@ -29,7 +30,8 @@ else:
'ofa_for_text_to_image_synthesis_model':
['OfaForTextToImageSynthesis'],
'multi_stage_diffusion':
- ['MultiStageDiffusionForTextToImageSynthesis']
+ ['MultiStageDiffusionForTextToImageSynthesis'],
+ 'vldoc': ['VLDocForDocVLEmbedding'],
}
import sys
diff --git a/modelscope/models/multi_modal/vldoc/__init__.py b/modelscope/models/multi_modal/vldoc/__init__.py
new file mode 100644
index 00000000..8a231402
--- /dev/null
+++ b/modelscope/models/multi_modal/vldoc/__init__.py
@@ -0,0 +1,3 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+from .model import VLDocForDocVLEmbedding
diff --git a/modelscope/models/multi_modal/vldoc/conv_fpn_trans.py b/modelscope/models/multi_modal/vldoc/conv_fpn_trans.py
new file mode 100644
index 00000000..65e27ab5
--- /dev/null
+++ b/modelscope/models/multi_modal/vldoc/conv_fpn_trans.py
@@ -0,0 +1,293 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import random
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import DropPath, trunc_normal_
+
+from modelscope.models.multi_modal.vldoc.convnext import convnext_tiny
+from modelscope.utils.logger import get_logger
+
+try:
+ import apex
+ import apex.normalization
+ LN = apex.normalization.FusedLayerNorm
+except ImportError:
+ LN = torch.nn.LayerNorm
+
+logging = get_logger()
+
+
+class ResidualAttentionBlock(nn.Module):
+
+ def __init__(self,
+ d_model: int,
+ n_head: int,
+ attn_mask: torch.Tensor = None,
+ expand_ratio=4.0,
+ init_values: float = None):
+ """
+ The implementation of the transformer block refers to:
+ https://github.com/openai/CLIP/blob/b46f5ac7587d2e1862f8b7b1573179d80dcdd620/clip/model.py
+ """
+ super().__init__()
+
+ self.attn = nn.MultiheadAttention(d_model, n_head)
+ self.ln_1 = LN(d_model)
+ self.mlp = nn.Sequential(
+ OrderedDict([('c_fc', nn.Linear(d_model, d_model * expand_ratio)),
+ ('gelu', QuickGELU()),
+ ('c_proj', nn.Linear(d_model * expand_ratio,
+ d_model))]))
+ self.ln_2 = LN(d_model)
+ self.attn_mask = attn_mask
+ if init_values is not None:
+ self.gamma_1 = nn.Parameter(
+ init_values * torch.ones((d_model)), requires_grad=True)
+ self.gamma_2 = nn.Parameter(
+ init_values * torch.ones((d_model)), requires_grad=True)
+ else:
+ self.gamma_1, self.gamma_2 = 1.0, 1.0
+
+ def attention(self, x: torch.Tensor):
+ self.attn_mask = self.attn_mask.to(
+ dtype=x.dtype,
+ device=x.device) if self.attn_mask is not None else None
+ return self.attn(
+ x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+ def forward(self, x: torch.Tensor):
+ x = x + self.gamma_1 * self.attention(self.ln_1(x))
+ x = x + self.gamma_2 * self.mlp(self.ln_2(x))
+ return x
+
+
+class QuickGELU(nn.Module):
+
+ def forward(self, x: torch.Tensor):
+ return x * torch.sigmoid(1.702 * x)
+
+
+def drop_grid(grid_map, drop_range=(0.3, 0.8), training=False):
+ """
+ only drop in the training phase.
+ grid_map: [N, D, T1, ...]
+ """
+ if training:
+ drop_ratio = random.random() * (drop_range[1]
+ - drop_range[0]) + drop_range[0]
+ # [N, T1, ...], True will be dropped
+ mask = (torch.rand_like(grid_map[:, 0]) < drop_ratio).bool()
+ grid_map = grid_map.masked_fill(mask.unsqueeze(1), 0.0)
+ return grid_map
+
+
+class GumbelSample(nn.Module):
+
+ def __init__(self, in_dim, num_keep):
+ super(GumbelSample, self).__init__()
+ self.keep_layer = nn.Sequential(
+ nn.Conv2d(
+ in_dim, 256, 3, stride=1, padding=3, dilation=3, bias=False),
+ nn.BatchNorm2d(256),
+ nn.ReLU(inplace=True),
+ nn.Conv2d(
+ 256, 256, 3, stride=1, padding=2, dilation=2, bias=False),
+ nn.BatchNorm2d(256),
+ nn.ReLU(inplace=True),
+ nn.Conv2d(256, 1, 3, stride=1, padding=1, dilation=1),
+ nn.Sigmoid(),
+ )
+ self.num_keep = num_keep
+ self.diffusion = nn.Conv2d(in_dim, in_dim, 3, padding=1)
+ self.dropout = nn.Dropout(0.1)
+
+ def forward(self, x, tau=1):
+ """
+ x: [N, C, H, W]
+ """
+ N = x.size(0)
+ keep_score = self.keep_layer(x)
+ keep_score = torch.clamp(keep_score, min=0.0, max=1.0)
+ keep_score = torch.cat([keep_score, 1 - keep_score],
+ dim=1) + 1e-5 # [N, 2, H, W]
+ gumbel_score = F.gumbel_softmax(
+ keep_score.log(), tau=tau, hard=False, dim=1)
+ # differentiable hard mode
+ index = gumbel_score.max(dim=1, keepdim=True)[1]
+ gumbel_hard = torch.zeros_like(
+ gumbel_score,
+ memory_format=torch.legacy_contiguous_format).scatter_(
+ 1, index, 1.0)
+ gumbel_hard = gumbel_hard - gumbel_score.detach() + gumbel_score
+ #
+ gumbel_score = gumbel_score[:, 0].contiguous().view(N,
+ -1) # [N, H x W]
+ gumbel_hard = gumbel_hard[:, 0].contiguous().view(N, -1)
+ # sort by score
+ idx_true = torch.topk(
+ gumbel_score, self.num_keep, dim=1)[1] # [N, num_keep]
+ topk_mask = torch.zeros_like(gumbel_score).bool().fill_(
+ False).scatter_(1, idx_true, True) # [N, H x W]
+ return topk_mask, gumbel_hard, keep_score[:, 0]
+
+ def sample(self, x, topk_mask, gumbel_hard):
+ N, D, H, W = x.size()
+ x = x.contiguous().view(N, D, -1) # [N, D, HxW]
+ x = x * gumbel_hard.unsqueeze(1)
+ x = x.transpose(1, 2) # [N, HxW, D]
+ x = x[topk_mask].contiguous().view(N, -1, D) # [N, num_keep, D]
+ x = drop_grid(
+ x.transpose(1, 2), drop_range=(0.0, 0.2),
+ training=self.training).transpose(1, 2)
+ return x
+
+ def random_sample(self, x):
+ N, D, H, W = x.size()
+ x = x.contiguous().view(N, D, -1) # [N, D, HxW]
+ x = x.transpose(1, 2) # [N, HxW, D]
+ # generate random mask
+ idx_true = torch.topk(
+ torch.rand_like(x[:, :, 0]), self.num_keep,
+ dim=1)[1] # [N, num_keep]
+ topk_mask = torch.zeros_like(x[:, :, 0]).bool().fill_(False).scatter_(
+ 1, idx_true, True) # [N, H x W]
+ # apply the mask
+ x = x[topk_mask].contiguous().view(N, -1, D) # [N, num_keep, D]
+ x = drop_grid(
+ x.transpose(1, 2), drop_range=(0.0, 0.2),
+ training=self.training).transpose(1, 2)
+ return x, topk_mask
+
+ def restore(self, x, topk_mask, src):
+ """
+ x: [N, D, H, W]
+ topk_mask: [N, HxW]
+ src: [N, num_keep, D]
+ """
+ N, D, H, W = x.size()
+ x = drop_grid(x, drop_range=(0.2, 0.8), training=self.training)
+ x = x.contiguous().view(N, D, -1).transpose(1, 2) # [N, HxW, D]
+ x = x.masked_scatter(topk_mask.unsqueeze(-1), src)
+ x = x.transpose(1, 2).contiguous().view(N, D, H, W)
+ x = self.dropout(self.diffusion(x))
+ return x
+
+
+class FPNTrans(nn.Module):
+
+ def __init__(self,
+ trans_layers=2,
+ inner_channels=256,
+ img_size=(896, 896),
+ inner_vit=False,
+ out_sampling=False):
+ super(FPNTrans, self).__init__()
+ self.cnn = convnext_tiny(pretrained=True, in_22k=True)
+ self.dims = self.cnn.dims
+ self.img_size = img_size
+ # FPN in DB
+ self.up5 = nn.Upsample(scale_factor=2, mode='nearest')
+ self.up4 = nn.Upsample(scale_factor=2, mode='nearest')
+ self.up3 = nn.Upsample(scale_factor=2, mode='nearest')
+
+ self.in5 = nn.Conv2d(self.dims[-1], inner_channels, 1, bias=False)
+ self.in4 = nn.Conv2d(self.dims[-2], inner_channels, 1, bias=False)
+ self.in3 = nn.Conv2d(self.dims[-3], inner_channels, 1, bias=False)
+ self.in2 = nn.Conv2d(self.dims[-4], inner_channels, 1, bias=False)
+
+ self.out5 = nn.Sequential(
+ nn.Conv2d(
+ inner_channels, inner_channels // 4, 3, padding=1, bias=False),
+ nn.Upsample(scale_factor=8, mode='nearest'))
+ self.out4 = nn.Sequential(
+ nn.Conv2d(
+ inner_channels, inner_channels // 4, 3, padding=1, bias=False),
+ nn.Upsample(scale_factor=4, mode='nearest'))
+ self.out3 = nn.Sequential(
+ nn.Conv2d(
+ inner_channels, inner_channels // 4, 3, padding=1, bias=False),
+ nn.Upsample(scale_factor=2, mode='nearest'))
+ self.out2 = nn.Conv2d(
+ inner_channels, inner_channels // 4, 3, padding=1, bias=False)
+ self.inner_vit = inner_vit
+ if inner_vit:
+ # mini vit
+ self.num_keep1 = (self.img_size[0] // 64)**2
+ self.gumble_sample1 = GumbelSample(
+ inner_channels, num_keep=self.num_keep1)
+ self.pos_emb1 = nn.Parameter(
+ torch.randn(inner_channels, self.img_size[0] // 32,
+ self.img_size[1] // 32))
+ trunc_normal_(self.pos_emb1, std=.02)
+ self.mini_vit = nn.Sequential(*[
+ ResidualAttentionBlock(
+ inner_channels, 4, expand_ratio=2, init_values=0.1)
+ for _ in range(trans_layers)
+ ])
+ self.dropout_pos = nn.Dropout(0.1)
+ if out_sampling:
+ # sample for co-attention
+ self.num_keep2 = (self.img_size[0] // 64)**2
+ self.gumble_sample2 = GumbelSample(
+ inner_channels, num_keep=self.num_keep2)
+ self.pos_emb2 = nn.Parameter(
+ torch.randn(inner_channels, self.img_size[0] // 4,
+ self.img_size[1] // 4))
+ trunc_normal_(self.pos_emb2, std=.02)
+ self.out_sampling = out_sampling
+
+ self.drop_path = DropPath(0.1)
+
+ def forward(self, x):
+ ms_features = self.cnn(x)
+ c2, c3, c4, c5 = ms_features
+ in5 = self.in5(c5)
+ in4 = self.in4(c4)
+ in3 = self.in3(c3)
+ in2 = self.in2(c2)
+ N, D5, H5, W5 = in5.size()
+ if self.inner_vit:
+ # random sample
+ keep_score = None
+ in5_pos = self.dropout_pos(in5 + self.pos_emb1.unsqueeze(0))
+ in5_pos_in, topk_mask = self.gumble_sample1.random_sample(in5_pos)
+ in5_pos_in = in5_pos_in.transpose(0, 1) # [num_keep1, N, D5]
+ in5_pos_out = self.mini_vit(in5_pos_in).permute(
+ 1, 2, 0) # [N, D5, num_keep1]
+ in5 = self.gumble_sample1.restore(in5, topk_mask,
+ in5_pos_out.transpose(1, 2))
+ else:
+ keep_score = None
+ # FPN for fused multi-scale visual feature
+ out4 = self.up5(in5) + self.drop_path(in4)
+ out3 = self.up4(out4) + self.drop_path(in3)
+ out2 = self.up3(out3) + self.drop_path(in2)
+ p5 = self.out5(in5)
+ p4 = self.out4(out4)
+ p3 = self.out3(out3)
+ p2 = self.out2(out2)
+ feat_ms = torch.cat((p5, p4, p3, p2), 1)
+ ret_dict = dict(
+ feat_ms=feat_ms,
+ keep_score=keep_score,
+ )
+ if self.out_sampling:
+ # gumbel sampling
+ topk_mask2, gumbel_hard2, keep_score2 = self.gumble_sample2(
+ feat_ms)
+ feat_ms_pos = self.dropout_pos(feat_ms
+ + self.pos_emb2.unsqueeze(0))
+ feat_ms_pos_sampled = self.gumble_sample2.sample(
+ feat_ms_pos, topk_mask2,
+ gumbel_hard2).transpose(0, 1) # [num_keep2, N, inner_c]
+ ret_dict_sup = dict(
+ feat_ms_pos_sampled=feat_ms_pos_sampled,
+ sampler=self.gumble_sample2,
+ keep_score2=keep_score2,
+ )
+ ret_dict.update(ret_dict_sup)
+ return ret_dict
diff --git a/modelscope/models/multi_modal/vldoc/convnext.py b/modelscope/models/multi_modal/vldoc/convnext.py
new file mode 100644
index 00000000..3cf7cd63
--- /dev/null
+++ b/modelscope/models/multi_modal/vldoc/convnext.py
@@ -0,0 +1,168 @@
+# The implementation is borrowed and partly modified from ConvNext,
+# made publicly available under the MIT License at https://github.com/facebookresearch/ConvNeXt.
+
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import DropPath, trunc_normal_
+from timm.models.registry import register_model
+
+
+class Block(nn.Module):
+ r""" ConvNeXt Block. There are two equivalent implementations:
+ (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
+ (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
+ We use (2) as we find it slightly faster in PyTorch
+
+ Args:
+ dim (int): Number of input channels.
+ drop_path (float): Stochastic depth rate. Default: 0.0
+ layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+ """
+
+ def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):
+ super().__init__()
+ self.dwconv = nn.Conv2d(
+ dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
+ self.norm = LayerNorm(dim, eps=1e-6)
+ self.pwconv1 = nn.Linear(
+ dim,
+ 4 * dim) # pointwise/1x1 convs, implemented with linear layers
+ self.act = nn.GELU()
+ self.pwconv2 = nn.Linear(4 * dim, dim)
+ self.gamma = nn.Parameter(
+ layer_scale_init_value * torch.ones((dim)),
+ requires_grad=True) if layer_scale_init_value > 0 else None
+ self.drop_path = DropPath(
+ drop_path) if drop_path > 0. else nn.Identity()
+
+ def forward(self, x):
+ input = x
+ x = self.dwconv(x)
+ x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+ x = self.norm(x)
+ x = self.pwconv1(x)
+ x = self.act(x)
+ x = self.pwconv2(x)
+ if self.gamma is not None:
+ x = self.gamma * x
+ x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
+
+ x = input + self.drop_path(x)
+ return x
+
+
+class ConvNeXt(nn.Module):
+ r""" ConvNeXt
+ A PyTorch impl of : `A ConvNet for the 2020s` -
+ https://arxiv.org/pdf/2201.03545.pdf
+ Args:
+ in_chans (int): Number of input image channels. Default: 3
+ num_classes (int): Number of classes for classification head. Default: 1000
+ depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+ dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+ drop_path_rate (float): Stochastic depth rate. Default: 0.
+ layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+ head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+ """
+
+ def __init__(
+ self,
+ in_chans=3,
+ depths=[3, 3, 9, 3],
+ dims=[96, 192, 384, 768],
+ drop_path_rate=0.,
+ layer_scale_init_value=1e-6,
+ ):
+ super().__init__()
+
+ self.downsample_layers = nn.ModuleList(
+ ) # stem and 3 intermediate downsampling conv layers
+ stem = nn.Sequential(
+ nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+ LayerNorm(dims[0], eps=1e-6, data_format='channels_first'))
+ self.downsample_layers.append(stem)
+ for i in range(3):
+ downsample_layer = nn.Sequential(
+ LayerNorm(dims[i], eps=1e-6, data_format='channels_first'),
+ nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+ )
+ self.downsample_layers.append(downsample_layer)
+
+ self.stages = nn.ModuleList(
+ ) # 4 feature resolution stages, each consisting of multiple residual blocks
+ dp_rates = [
+ x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
+ ]
+ cur = 0
+ for i in range(4):
+ stage = nn.Sequential(*[
+ Block(
+ dim=dims[i],
+ drop_path=dp_rates[cur + j],
+ layer_scale_init_value=layer_scale_init_value)
+ for j in range(depths[i])
+ ])
+ self.stages.append(stage)
+ cur += depths[i]
+
+ # self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer
+ self.dims = dims
+
+ self.apply(self._init_weights)
+
+ def _init_weights(self, m):
+ if isinstance(m, (nn.Conv2d, nn.Linear)):
+ trunc_normal_(m.weight, std=.02)
+ nn.init.constant_(m.bias, 0)
+
+ def forward(self, x):
+ xs = []
+ for i in range(4):
+ x = self.downsample_layers[i](x)
+ x = self.stages[i](x)
+ xs.append(x)
+ # x = x.permute(0, 2, 3, 1) # [N, H, W, C]
+ # x = self.norm(x)
+ # x = x.permute(0, 3, 1, 2) # [N, C, H, W]
+ return tuple(xs)
+
+
+class LayerNorm(nn.Module):
+ r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
+ The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
+ shape (batch_size, height, width, channels) while channels_first corresponds to inputs
+ with shape (batch_size, channels, height, width).
+ """
+
+ def __init__(self,
+ normalized_shape,
+ eps=1e-6,
+ data_format='channels_last'):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(normalized_shape))
+ self.bias = nn.Parameter(torch.zeros(normalized_shape))
+ self.eps = eps
+ self.data_format = data_format
+ if self.data_format not in ['channels_last', 'channels_first']:
+ raise NotImplementedError
+ self.normalized_shape = (normalized_shape, )
+
+ def forward(self, x):
+ if self.data_format == 'channels_last':
+ return F.layer_norm(x, self.normalized_shape, self.weight,
+ self.bias, self.eps)
+ elif self.data_format == 'channels_first':
+ u = x.mean(1, keepdim=True)
+ s = (x - u).pow(2).mean(1, keepdim=True)
+ x = (x - u) / torch.sqrt(s + self.eps)
+ x = self.weight[:, None, None] * x + self.bias[:, None, None]
+ return x
+
+
+@register_model
+def convnext_tiny(pretrained=False, in_22k=False, **kwargs):
+ model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
+ return model
diff --git a/modelscope/models/multi_modal/vldoc/model.py b/modelscope/models/multi_modal/vldoc/model.py
new file mode 100644
index 00000000..5b21bf10
--- /dev/null
+++ b/modelscope/models/multi_modal/vldoc/model.py
@@ -0,0 +1,433 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import copy
+import logging
+import math
+import os
+import re
+import sys
+
+import json
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torchvision.ops import roi_align
+
+from modelscope.metainfo import Models
+from modelscope.models import TorchModel
+from modelscope.models.builder import MODELS
+from modelscope.models.multi_modal.vldoc.conv_fpn_trans import FPNTrans
+from modelscope.models.multi_modal.vldoc.modeling_layout_roberta import (
+ LayoutRobertaModel, LayoutRobertaPreTrainedModel)
+from modelscope.models.multi_modal.vldoc.transformer_local import (
+ TransformerDecoder, TransformerDecoderLayer)
+from modelscope.utils.constant import ModeKeys, ModelFile, Tasks
+from modelscope.utils.logger import get_logger
+
+logger = get_logger()
+
+__all__ = ['VLDocForDocVLEmbedding']
+
+
+class GeoVLDocModelOutputs(object):
+
+ def __init__(
+ self,
+ text_features,
+ text_mm_features,
+ block_vis_features,
+ block_vis_mm_features,
+ image_mm_features,
+ ):
+ # [batch size, sequence length, hidden size]
+ self.text_features = text_features
+ # [batch size, sequence length, hidden size]
+ self.text_mm_features = text_mm_features
+ # [batch size, block num, hidden size]
+ self.block_vis_features = block_vis_features
+ # [batch size, block num, hidden size]
+ self.block_vis_mm_features = block_vis_mm_features
+ # [batch size, hidden size]
+ self.image_mm_features = image_mm_features
+
+
+class GeoVLDocModel(LayoutRobertaPreTrainedModel):
+
+ def __init__(self, config, hard_negtive_sampling=False):
+ super().__init__(config)
+ self.config = config
+ self.hard_negtive_sampling = hard_negtive_sampling
+
+ if getattr(self.config, 'architectures', None):
+ if self.config.architectures[0] == 'LayoutRobertaModel':
+ self.text_encoder = LayoutRobertaModel(config)
+ else:
+ self.text_encoder = LayoutRobertaModel(config)
+ else:
+ self.text_encoder = LayoutRobertaModel(config)
+ self.visual_encoder = FPNTrans(
+ img_size=self.config.image_size, inner_vit=False)
+ self.pool = nn.AdaptiveAvgPool2d([1, 1])
+ self.vis_linear = nn.Linear(256, self.config.hidden_size)
+
+ cross_modal_text_layer = TransformerDecoderLayer(
+ self.config.hidden_size,
+ self.config.num_attention_heads,
+ self.config.intermediate_size,
+ self_attn=True)
+ self.cross_modal_text = TransformerDecoder(cross_modal_text_layer, 1)
+
+ cross_modal_visual_layer = TransformerDecoderLayer(
+ self.config.hidden_size,
+ self.config.num_attention_heads,
+ self.config.intermediate_size,
+ self_attn=True)
+ self.cross_modal_visual = TransformerDecoder(cross_modal_visual_layer,
+ 1)
+
+ self.init_weights()
+
+ def from_pretrained(self, ckpt_path: str):
+ state_dict = torch.load(ckpt_path, map_location='cpu')
+ state_dict_new = {}
+ for k, v in state_dict.items():
+ k = k.replace('geo_vl_doc_model.', '')
+ state_dict_new[k] = v
+ self.load_state_dict(state_dict_new)
+
+ def forward(self,
+ input_ids=None,
+ image=None,
+ bbox=None,
+ bbox_4p_normalized=None,
+ attention_mask=None,
+ first_token_idxes=None,
+ first_token_idxes_mask=None,
+ token_type_ids=None,
+ position_ids=None,
+ head_mask=None,
+ inputs_embeds=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_values=None,
+ use_cache=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ **kwargs):
+
+ batch_size, seq_len = input_ids.shape
+
+ return_dict = (
+ return_dict
+ if return_dict is not None else self.config.use_return_dict)
+
+ kwargs['line_bbox'] = bbox
+ # ################ get text representation ################
+ if self.config.architectures[0] == 'LayoutRobertaModel':
+ outputs = self.text_encoder(
+ input_ids,
+ bbox=bbox_4p_normalized,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ **kwargs)
+ else:
+ outputs = self.text_encoder(
+ input_ids,
+ bbox=bbox_4p_normalized,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ **kwargs)
+
+ # sequence_output: [batch_size, seq_len, hidden_size]
+ # pooled_output: [batch_size, hidden_size]
+ sequence_output, pooled_output = outputs[:2]
+
+ # ################ get visual representation ################
+ _, num_first = first_token_idxes.shape
+ B_batch_dim = torch.arange(
+ 0, batch_size,
+ device=input_ids.device).reshape(batch_size,
+ 1).expand(batch_size, num_first)
+
+ feature_bbox = bbox[B_batch_dim, first_token_idxes]
+ _, block_num, _ = feature_bbox.shape
+
+ visual_out = self.visual_encoder(image)
+ batch_idxs = torch.arange(
+ 0, batch_size, device=sequence_output.device).reshape(
+ batch_size, 1).expand(batch_size, block_num).unsqueeze(-1)
+
+ # [batch_size*block_num, 5]
+ batch_idx_with_bbox = torch.cat(
+ (batch_idxs, feature_bbox),
+ 2).reshape(batch_size * block_num,
+ 5).to(dtype=visual_out['feat_ms'].dtype)
+
+ if visual_out['feat_ms'].dtype == torch.float16:
+ # [batch_size*block_num, 256, 1, 1]
+ blk_vis_features = roi_align(
+ visual_out['feat_ms'].to(torch.float32),
+ batch_idx_with_bbox.to(torch.float32),
+ 1,
+ spatial_scale=visual_out['feat_ms'].size(-1) / 1000.0)
+ blk_vis_features = blk_vis_features.to(
+ dtype=visual_out['feat_ms'].dtype)
+ else:
+ blk_vis_features = roi_align(
+ visual_out['feat_ms'],
+ batch_idx_with_bbox.to(torch.float32),
+ 1,
+ spatial_scale=visual_out['feat_ms'].size(-1) / 1000.0)
+
+ # [batch_size*block_num, 256]
+ blk_vis_features = blk_vis_features.squeeze(2).squeeze(2).reshape(
+ batch_size, block_num, 256)
+
+ # visual block features:
+ # blk_vis_features: [batch_size, block_num, hidden_size]
+ blk_vis_features = self.vis_linear(blk_vis_features)
+ blk_vis_features = blk_vis_features * first_token_idxes_mask.unsqueeze(
+ 2)
+ # [batch_size, 256]
+ full_img_features = self.pool(
+ visual_out['feat_ms']).squeeze(2).squeeze(2)
+ # [batch_size, hidden_size]
+ full_img_features = self.vis_linear(full_img_features).unsqueeze(1)
+
+ # ################ multi-modal fusion ################
+
+ # cross attention inputs
+ vis_inps = torch.cat((full_img_features, blk_vis_features), 1)
+
+ glb_feat_attn = torch.ones((batch_size, 1)).to(input_ids.device)
+
+ vis_mask = torch.cat((glb_feat_attn, first_token_idxes_mask), 1)
+
+ # When we use transformer in torch.nn, the input size is
+ # [seq_len, batch_size, hidden_size]
+ # In attention_mask, 1 denotes masked
+ new_attention_mask = (1 - attention_mask) > 0
+ new_vis_mask = (1 - vis_mask) > 0
+
+ text_mm_feat = self.cross_modal_text(
+ tgt=sequence_output.transpose(0, 1),
+ memory=vis_inps.transpose(0, 1),
+ tgt_key_padding_mask=new_attention_mask,
+ memory_key_padding_mask=new_vis_mask)
+
+ vis_mm_feat = self.cross_modal_visual(
+ tgt=vis_inps.transpose(0, 1),
+ memory=sequence_output.transpose(0, 1),
+ tgt_key_padding_mask=new_vis_mask,
+ memory_key_padding_mask=new_attention_mask,
+ )
+
+ # [batch_size, seq_len, hidden_size]
+ text_mm_feat = text_mm_feat.transpose(0, 1)
+ # [batch_size, 1+block_num, hidden_size]
+ vis_mm_feat = vis_mm_feat.transpose(0, 1)
+
+ # image_mm_features = vis_mm_feat[:, 0, :]
+ block_vis_mm_features = vis_mm_feat[:, 1:]
+
+ return GeoVLDocModelOutputs(
+ text_features=sequence_output,
+ text_mm_features=text_mm_feat,
+ block_vis_features=blk_vis_features,
+ block_vis_mm_features=block_vis_mm_features,
+ image_mm_features=vis_mm_feat,
+ )
+
+
+@MODELS.register_module(Tasks.document_vl_embedding, module_name=Models.vldoc)
+class VLDocForDocVLEmbedding(TorchModel):
+ """
+ Generate multi-modal document embeddings in segment-level and token-level.
+
+ Args:
+ model_dir:
+ the path in model hub, e.g., 'damo/multi-modal_convnext-roberta-base_vldoc-embedding'
+ """
+
+ def __init__(self, model_dir: str, *args, **kwargs):
+ super().__init__(model_dir=model_dir, *args, **kwargs)
+
+ # Initialize the model.
+ from modelscope.models.multi_modal.vldoc.modeling_layout_roberta import LayoutRobertaConfig
+ model_cfg_path = os.path.join(model_dir, 'config.json')
+ logger.info('Loading config file from {}'.format(model_cfg_path))
+ assert os.path.exists(model_cfg_path)
+ self.config = LayoutRobertaConfig.from_json_file(model_cfg_path)
+ self.doc_model = GeoVLDocModel(self.config)
+
+ # restore the pretrained weight
+ model_path = os.path.join(model_dir, ModelFile.TORCH_MODEL_FILE)
+ assert os.path.exists(model_path)
+ self.doc_model.from_pretrained(model_path)
+ logger.info('Loading model from {}'.format(model_path))
+
+ # Initialize the tokenizer.
+ from modelscope.models.multi_modal.vldoc.tokenization import VLDocXLMTokenizer
+ tokenizer_path = os.path.join(model_dir, ModelFile.TOKENIZER_FOLDER)
+ self.tokenizer = VLDocXLMTokenizer.from_pretrained(tokenizer_path)
+
+ # place the model
+ self.device = 'cuda:{}'.format(int(os.environ.get(
+ 'LOCAL_RANK', 0))) if torch.cuda.is_available() else 'cpu'
+ if torch.cuda.is_available():
+ self.doc_model.to(self.device)
+ logger.info('Use GPU {} for finetuning & inference'.format(
+ int(os.environ.get('LOCAL_RANK', 0))))
+ else:
+ self.doc_model.float()
+ logger.info('Use CPU for finetuning & inference')
+
+ def forward(self,
+ input_ids=None,
+ image=None,
+ bbox=None,
+ bbox_4p_normalized=None,
+ attention_mask=None,
+ first_token_idxes=None,
+ first_token_idxes_mask=None,
+ token_type_ids=None,
+ position_ids=None,
+ head_mask=None,
+ inputs_embeds=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_values=None,
+ use_cache=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ **kwargs):
+ """
+ Args:
+ - input_ids: :math:`(B, T, E)`, the input tokens, where B is the batch size,
+ T is the max token size, E is the embedding dimension.
+ - image: :math:`(B, C, H, W)`, normalized images.
+ - bbox: :math:`(B, T, 4)`, segment boxes denoted by top-left and bottom-right
+ vertexes whose values are normalized to [0, 1000).
+ - bbox_4p_normalized: :math:`(B, T, 8)`, word boxes denoted by 4 vertexes, whose
+ values are normalized to [0, 1).
+ - attention_mask: :math:`(B, T)`, mask for input tokens, where 0 means masked.
+ - first_token_idxes: :math:`(B, S)`, indexes of the corresponding first tokens
+ of all segments, where S is the max segment size.
+ - first_token_idxes_mask: :math:`(B, S)`, mask for segments, where 0 means masked.
+ Optional:
+ - line_rank_id: :math:`(B, T)`, orders of segments.
+ - line_rank_inner_id: :math:`(B, T)`, BIE-like tags.
+
+ To be more specific, please refer to the class `TextLayoutSerializer` in
+ `modelscope/models/multi_modal/vldoc/processing.py`.
+ """
+
+ vldoc_outputs = self.doc_model(
+ input_ids=input_ids,
+ image=image,
+ bbox=bbox,
+ bbox_4p_normalized=bbox_4p_normalized,
+ attention_mask=attention_mask,
+ first_token_idxes=first_token_idxes,
+ first_token_idxes_mask=first_token_idxes_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ head_mask=head_mask,
+ inputs_embeds=inputs_embeds,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ **kwargs)
+
+ return dict(
+ img_embedding=vldoc_outputs.image_mm_features,
+ text_embedding=vldoc_outputs.text_mm_features,
+ )
+
+
+def init_pretrained_weight(
+ model,
+ pretrained_model_path,
+ state_dict=None,
+ cache_dir=None,
+ init_backbone='roberta',
+):
+ if state_dict is None:
+ state_dict = torch.load(pretrained_model_path, map_location='cpu')
+
+ old_keys = []
+ new_keys = []
+ state_dict_keys = list(state_dict.keys())
+
+ if init_backbone == 'roberta':
+ for i in range(len(state_dict_keys)):
+ key = state_dict_keys[i]
+ new_key = None
+
+ if key.startswith('roberta.'):
+ new_key = key.replace('roberta.',
+ 'geo_vl_doc_model.text_encoder.')
+ key = copy.deepcopy(new_key)
+
+ if new_key:
+ old_keys.append(state_dict_keys[i])
+ new_keys.append(new_key)
+
+ for old_key, new_key in zip(old_keys, new_keys):
+ state_dict[new_key] = state_dict.pop(old_key)
+
+ missing_keys = []
+ unexpected_keys = []
+ error_msgs = []
+
+ # copy state_dict so _load_from_state_dict can modify it
+ metadata = getattr(state_dict, '_metadata', None)
+ state_dict = state_dict.copy()
+ if metadata is not None:
+ state_dict._metadata = metadata
+
+ def load(module, prefix=''):
+ local_metadata = {} if metadata is None else metadata.get(
+ prefix[:-1], {})
+ module._load_from_state_dict(state_dict, prefix, local_metadata, True,
+ missing_keys, unexpected_keys, error_msgs)
+ for name, child in module._modules.items():
+ if child is not None:
+ load(child, prefix + name + '.')
+
+ start_prefix = ''
+ if not hasattr(model, 'geo_vl_doc_model') and any(
+ s.startswith('geo_vl_doc_model.') for s in state_dict.keys()):
+ start_prefix = 'geo_vl_doc_model.'
+ load(model, prefix=start_prefix)
+ if len(missing_keys) > 0:
+ logger.info(
+ 'Weights of {} not initialized from pretrained model: {}'.format(
+ model.__class__.__name__, missing_keys))
+ if len(unexpected_keys) > 0:
+ logger.info('Weights from pretrained model not used in {}: {}'.format(
+ model.__class__.__name__, unexpected_keys))
+ if len(error_msgs) > 0:
+ raise RuntimeError(
+ 'Error(s) in loading state_dict for {}:\n\t{}'.format(
+ model.__class__.__name__, '\n\t'.join(error_msgs)))
+
+ return model
diff --git a/modelscope/models/multi_modal/vldoc/modeling_layout_roberta.py b/modelscope/models/multi_modal/vldoc/modeling_layout_roberta.py
new file mode 100644
index 00000000..5ae32a76
--- /dev/null
+++ b/modelscope/models/multi_modal/vldoc/modeling_layout_roberta.py
@@ -0,0 +1,1140 @@
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright 2021-2022 The Alibaba DAMO Duguang Team Authors. All rights reserved.
+
+import math
+import os
+
+import torch
+import torch.utils.checkpoint
+from packaging import version
+from torch import nn
+from transformers.activations import ACT2FN, gelu
+from transformers.configuration_utils import PretrainedConfig
+from transformers.file_utils import (add_code_sample_docstrings,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ replace_return_docstrings)
+from transformers.modeling_outputs import (
+ BaseModelOutputWithPastAndCrossAttentions,
+ BaseModelOutputWithPoolingAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions, MaskedLMOutput,
+ MultipleChoiceModelOutput, QuestionAnsweringModelOutput,
+ SequenceClassifierOutput, TokenClassifierOutput)
+from transformers.modeling_utils import (PreTrainedModel,
+ apply_chunking_to_forward,
+ find_pruneable_heads_and_indices,
+ prune_linear_layer)
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class LayoutRobertaConfig(PretrainedConfig):
+ model_type = 'layoutroberta'
+
+ def __init__(self,
+ vocab_size=30522,
+ hidden_size=768,
+ num_hidden_layers=12,
+ num_attention_heads=12,
+ intermediate_size=3072,
+ hidden_act='gelu',
+ hidden_dropout_prob=0.1,
+ attention_probs_dropout_prob=0.1,
+ max_position_embeddings=512,
+ type_vocab_size=2,
+ initializer_range=0.02,
+ layer_norm_eps=1e-12,
+ pad_token_id=1,
+ bos_token_id=0,
+ eos_token_id=2,
+ bbox_scale=100.0,
+ pe_type='crel',
+ position_embedding_type='absolute',
+ use_cache=True,
+ classifier_dropout=None,
+ **kwargs):
+ super().__init__(
+ vocab_size=vocab_size,
+ hidden_size=hidden_size,
+ num_hidden_layers=num_hidden_layers,
+ num_attention_heads=num_attention_heads,
+ intermediate_size=intermediate_size,
+ hidden_act=hidden_act,
+ hidden_dropout_prob=hidden_dropout_prob,
+ attention_probs_dropout_prob=attention_probs_dropout_prob,
+ max_position_embeddings=max_position_embeddings,
+ type_vocab_size=type_vocab_size,
+ initializer_range=initializer_range,
+ layer_norm_eps=layer_norm_eps,
+ pad_token_id=pad_token_id,
+ bos_token_id=bos_token_id,
+ eos_token_id=eos_token_id,
+ **kwargs,
+ )
+
+ self.bbox_scale = bbox_scale
+ self.pe_type = pe_type
+
+
+class PositionalEmbedding1D(nn.Module):
+ # Reference:
+ # https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/mem_transformer.py#L15
+
+ def __init__(self, demb):
+ super(PositionalEmbedding1D, self).__init__()
+
+ self.demb = demb
+
+ inv_freq = 1 / (10000**(torch.arange(0.0, demb, 2.0) / demb))
+ self.register_buffer('inv_freq', inv_freq)
+
+ def forward(self, pos_seq, bsz=None):
+ seq_size = pos_seq.size()
+
+ if len(seq_size) == 2:
+ b1, b2 = seq_size
+ sinusoid_inp = pos_seq.view(b1, b2, 1) * self.inv_freq.view(
+ 1, 1, self.demb // 2)
+ elif len(seq_size) == 3:
+ b1, b2, b3 = seq_size
+ sinusoid_inp = pos_seq.view(b1, b2, b3, 1) * self.inv_freq.view(
+ 1, 1, 1, self.demb // 2)
+ else:
+ raise ValueError(f'Invalid seq_size={len(seq_size)}')
+
+ pos_emb = torch.cat([sinusoid_inp.sin(), sinusoid_inp.cos()], dim=-1)
+
+ return pos_emb
+
+
+class PositionalEmbedding2D(nn.Module):
+
+ def __init__(self, demb, dim_bbox=8):
+ super(PositionalEmbedding2D, self).__init__()
+
+ self.demb = demb
+ self.dim_bbox = dim_bbox
+
+ self.x_pos_emb = PositionalEmbedding1D(demb // dim_bbox)
+ self.y_pos_emb = PositionalEmbedding1D(demb // dim_bbox)
+
+ inv_freq = 1 / (10000**(torch.arange(0.0, demb, 2.0) / demb))
+ self.register_buffer('inv_freq', inv_freq)
+
+ def forward(self, bbox):
+ # bbox: [seq_length, batch_size, dim_bbox]
+ stack = []
+ for i in range(self.dim_bbox):
+ if i % 2 == 0:
+ stack.append(self.x_pos_emb(bbox[..., i]))
+ else:
+ stack.append(self.y_pos_emb(bbox[..., i]))
+ bbox_pos_emb = torch.cat(stack, dim=-1)
+ return bbox_pos_emb
+
+
+class LayoutRobertaEmbeddings(nn.Module):
+ """
+ Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
+ """
+
+ # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.__init__
+ def __init__(self, config):
+ super().__init__()
+ self.word_embeddings = nn.Embedding(
+ config.vocab_size,
+ config.hidden_size,
+ padding_idx=config.pad_token_id)
+ self.position_embeddings = nn.Embedding(config.max_position_embeddings,
+ config.hidden_size)
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size,
+ config.hidden_size)
+
+ # layout-related embeddings
+ self.line_rank_embeddings = nn.Embedding(
+ config.max_position_embeddings, config.hidden_size)
+
+ self.line_rank_inner_embeddings = nn.Embedding(4, config.hidden_size)
+
+ self.x_position_embeddings = nn.Embedding(
+ config.max_2d_position_embeddings, config.coordinate_size)
+ self.y_position_embeddings = nn.Embedding(
+ config.max_2d_position_embeddings, config.coordinate_size)
+ self.h_position_embeddings = nn.Embedding(
+ config.max_2d_position_embeddings, config.shape_size)
+ self.w_position_embeddings = nn.Embedding(
+ config.max_2d_position_embeddings, config.shape_size)
+
+ self.LayerNorm = nn.LayerNorm(
+ config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.position_embedding_type = getattr(config,
+ 'position_embedding_type',
+ 'absolute')
+
+ self.register_buffer(
+ 'position_ids',
+ torch.arange(config.max_position_embeddings).expand((1, -1)))
+
+ if version.parse(torch.__version__) > version.parse('1.6.0'):
+ self.register_buffer(
+ 'token_type_ids',
+ torch.zeros(self.position_ids.size(), dtype=torch.long),
+ persistent=False,
+ )
+
+ if config.pe_type == 'pdpdq_ws':
+ dim_bbox_sinusoid_emb = config.hidden_size
+ dim_bbox_projection = config.hidden_size
+ elif config.pe_type == 'crel':
+ dim_bbox_sinusoid_emb = config.hidden_size // 4
+ dim_bbox_projection = config.hidden_size // config.num_attention_heads
+ else:
+ raise ValueError(f'Unknown config.pe_type={config.pe_type}')
+
+ self.bbox_sinusoid_emb = PositionalEmbedding2D(
+ dim_bbox_sinusoid_emb, dim_bbox=8)
+ self.bbox_projection = nn.Linear(
+ dim_bbox_sinusoid_emb, dim_bbox_projection, bias=False)
+
+ # End copy
+ self.padding_idx = config.pad_token_id
+ self.position_embeddings = nn.Embedding(
+ config.max_position_embeddings,
+ config.hidden_size,
+ padding_idx=self.padding_idx)
+
+ def _cal_spatial_position_embeddings(self, bbox):
+ try:
+ left_position_embeddings = self.x_position_embeddings(bbox[:, :,
+ 0])
+ upper_position_embeddings = self.y_position_embeddings(bbox[:, :,
+ 1])
+ right_position_embeddings = self.x_position_embeddings(bbox[:, :,
+ 2])
+ lower_position_embeddings = self.y_position_embeddings(bbox[:, :,
+ 3])
+ except IndexError as e:
+ raise IndexError(
+ 'The :obj:`bbox`coordinate values should be within 0-1000 range.'
+ ) from e
+
+ h_position_embeddings = self.h_position_embeddings(bbox[:, :, 3]
+ - bbox[:, :, 1])
+ w_position_embeddings = self.w_position_embeddings(bbox[:, :, 2]
+ - bbox[:, :, 0])
+
+ spatial_position_embeddings = torch.cat(
+ [
+ left_position_embeddings,
+ upper_position_embeddings,
+ right_position_embeddings,
+ lower_position_embeddings,
+ h_position_embeddings,
+ w_position_embeddings,
+ ],
+ dim=-1,
+ )
+ return spatial_position_embeddings
+
+ def forward(self,
+ input_ids=None,
+ token_type_ids=None,
+ position_ids=None,
+ inputs_embeds=None,
+ past_key_values_length=0,
+ **kwargs):
+ if position_ids is None:
+ if input_ids is not None:
+ # Create the position ids from the input token ids. Any padded tokens remain padded.
+ position_ids = create_position_ids_from_input_ids(
+ input_ids, self.padding_idx, past_key_values_length)
+ else:
+ position_ids = self.create_position_ids_from_inputs_embeds(
+ inputs_embeds)
+
+ if input_ids is not None:
+ input_shape = input_ids.size()
+ else:
+ input_shape = inputs_embeds.size()[:-1]
+
+ seq_length = input_shape[1]
+
+ # Setting the token_type_ids to the registered buffer in constructor
+ # where it is all zeros, which usually occurs
+ # when its auto-generated, registered buffer helps users when
+ # tracing the model without passing token_type_ids, solves
+ # issue #5664
+ if token_type_ids is None:
+ if hasattr(self, 'token_type_ids'):
+ buffered_token_type_ids = self.token_type_ids[:, :seq_length]
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(
+ input_shape[0], seq_length)
+ token_type_ids = buffered_token_type_ids_expanded
+ else:
+ token_type_ids = torch.zeros(
+ input_shape,
+ dtype=torch.long,
+ device=self.position_ids.device)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+ embeddings = inputs_embeds + token_type_embeddings
+ if self.position_embedding_type == 'absolute':
+ position_embeddings = self.position_embeddings(position_ids)
+ embeddings += position_embeddings
+
+ if 'line_bbox' in kwargs:
+ embeddings += self._cal_spatial_position_embeddings(
+ kwargs['line_bbox'])
+
+ if 'line_rank_id' in kwargs:
+ embeddings += self.line_rank_embeddings(kwargs['line_rank_id'])
+
+ if 'line_rank_inner_id' in kwargs:
+ embeddings += self.line_rank_inner_embeddings(
+ kwargs['line_rank_inner_id'])
+
+ embeddings = self.LayerNorm(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+ def calc_bbox_pos_emb(self, bbox, pe_type):
+ # bbox_t: [seq_length, batch_size, dim_bbox]
+ bbox_t = bbox.transpose(0, 1)
+
+ if pe_type == 'pdpdq_ws':
+ bbox_pos = bbox_t
+ elif pe_type == 'crel':
+ # bbox_pos: [seq_length, seq_length, batch_size, dim_bbox]
+ bbox_pos = bbox_t[None, :, :, :] - bbox_t[:, None, :, :]
+ else:
+ raise ValueError(f'Unknown pe_type={pe_type}')
+
+ bbox_pos_emb = self.bbox_sinusoid_emb(bbox_pos)
+ bbox_pos_emb = self.bbox_projection(bbox_pos_emb)
+
+ return bbox_pos_emb
+
+ def create_position_ids_from_inputs_embeds(self, inputs_embeds):
+ """
+ We are provided embeddings directly.
+ We cannot infer which are padded so just generate sequential position ids.
+
+ Args:
+ inputs_embeds: torch.Tensor
+
+ Returns: torch.Tensor
+ """
+ input_shape = inputs_embeds.size()[:-1]
+ sequence_length = input_shape[1]
+
+ position_ids = torch.arange(
+ self.padding_idx + 1,
+ sequence_length + self.padding_idx + 1,
+ dtype=torch.long,
+ device=inputs_embeds.device)
+ return position_ids.unsqueeze(0).expand(input_shape)
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->Roberta
+class LayoutRobertaSelfAttention(nn.Module):
+
+ def __init__(self, config, position_embedding_type=None):
+ super().__init__()
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
+ config, 'embedding_size'):
+ raise ValueError(
+ f'The hidden size ({config.hidden_size}) is not a multiple of the number of attention '
+ f'heads ({config.num_attention_heads})')
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size
+ / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+ self.position_embedding_type = position_embedding_type or getattr(
+ config, 'position_embedding_type', 'absolute')
+ if self.position_embedding_type == 'relative_key' or self.position_embedding_type == 'relative_key_query':
+ self.max_position_embeddings = config.max_position_embeddings
+ self.distance_embedding = nn.Embedding(
+ 2 * config.max_position_embeddings - 1,
+ self.attention_head_size)
+
+ self.is_decoder = config.is_decoder
+
+ self.pe_type = config.pe_type
+
+ def transpose_for_scores(self, x):
+ new_x_shape = x.size()[:-1] + (
+ self.num_attention_heads,
+ self.attention_head_size,
+ )
+ x = x.view(*new_x_shape)
+ return x.permute(0, 2, 1, 3)
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ head_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_value=None,
+ output_attentions=False,
+ bbox_pos_emb=None,
+ bbox_pos_mask=None,
+ ):
+ mixed_query_layer = self.query(hidden_states)
+
+ # If this is instantiated as a cross-attention module, the keys
+ # and values come from an encoder; the attention mask needs to be
+ # such that the encoder's padding tokens are not attended to.
+ is_cross_attention = encoder_hidden_states is not None
+
+ if is_cross_attention and past_key_value is not None:
+ # reuse k,v, cross_attentions
+ key_layer = past_key_value[0]
+ value_layer = past_key_value[1]
+ attention_mask = encoder_attention_mask
+ elif is_cross_attention:
+ key_layer = self.transpose_for_scores(
+ self.key(encoder_hidden_states))
+ value_layer = self.transpose_for_scores(
+ self.value(encoder_hidden_states))
+ attention_mask = encoder_attention_mask
+ elif past_key_value is not None:
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
+ key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
+ value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
+ else:
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
+
+ query_layer = self.transpose_for_scores(mixed_query_layer)
+
+ if self.is_decoder:
+ # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
+ # Further calls to cross_attention layer can then reuse all cross-attention
+ # key/value_states (first "if" case)
+ # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
+ # all previous decoder key/value_states. Further calls to uni-directional self-attention
+ # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
+ # if encoder bi-directional self-attention `past_key_value` is always `None`
+ past_key_value = (key_layer, value_layer)
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attention_scores = torch.matmul(query_layer,
+ key_layer.transpose(-1, -2))
+
+ if (self.position_embedding_type == 'relative_key'
+ or self.position_embedding_type == 'relative_key_query'):
+ seq_length = hidden_states.size()[1]
+ position_ids_l = torch.arange(
+ seq_length, dtype=torch.long,
+ device=hidden_states.device).view(-1, 1)
+ position_ids_r = torch.arange(
+ seq_length, dtype=torch.long,
+ device=hidden_states.device).view(1, -1)
+ distance = position_ids_l - position_ids_r
+ positional_embedding = self.distance_embedding(
+ distance + self.max_position_embeddings - 1)
+ positional_embedding = positional_embedding.to(
+ dtype=query_layer.dtype) # fp16 compatibility
+
+ if self.position_embedding_type == 'relative_key':
+ relative_position_scores = torch.einsum(
+ 'bhld,lrd->bhlr', query_layer, positional_embedding)
+ attention_scores = attention_scores + relative_position_scores
+ elif self.position_embedding_type == 'relative_key_query':
+ relative_position_scores_query = torch.einsum(
+ 'bhld,lrd->bhlr', query_layer, positional_embedding)
+ relative_position_scores_key = torch.einsum(
+ 'bhrd,lrd->bhlr', key_layer, positional_embedding)
+ attention_scores = (
+ attention_scores + relative_position_scores_query
+ + relative_position_scores_key)
+
+ # bbox positional encoding
+ batch_size, n_head, seq_length, d_head = query_layer.shape
+ if self.pe_type == 'pdpdq_ws':
+ head_q_pos = self.query(bbox_pos_emb)
+ head_k_pos = self.key(bbox_pos_emb)
+ head_q_pos = head_q_pos.view(seq_length, batch_size, n_head,
+ d_head)
+ head_k_pos = head_k_pos.view(seq_length, batch_size, n_head,
+ d_head)
+ head_q_pos = head_q_pos.permute([1, 2, 0, 3])
+ head_k_pos = head_k_pos.permute([1, 2, 0, 3])
+
+ bbox_pos_scores_1 = torch.einsum(
+ 'bnid,bnjd->bnij',
+ (torch.mul(query_layer, head_q_pos), head_k_pos))
+ bbox_pos_scores_2 = torch.einsum('bnid,bnjd->bnij',
+ (head_q_pos, head_k_pos))
+ bbox_pos_scores = bbox_pos_scores_1 + bbox_pos_scores_2
+ elif self.pe_type == 'crel':
+ bbox_pos_emb = bbox_pos_emb.view(seq_length, seq_length,
+ batch_size, d_head)
+ bbox_pos_emb = bbox_pos_emb.permute([2, 0, 1, 3])
+ bbox_pos_scores = torch.einsum('bnid,bijd->bnij',
+ (query_layer, bbox_pos_emb))
+ else:
+ raise ValueError(f'Unknown self.pe_type={self.pe_type}')
+
+ if bbox_pos_mask is not None:
+ # bbox_pos_mask is [batch_size, seq_length]
+ bbox_pos_mask = 1 - bbox_pos_mask
+ # [batch_size, 1, seq_length]
+ M1 = bbox_pos_mask.unsqueeze(1)
+ # [batch_size, seq_length, 1]
+ MT = M1.permute(0, 2, 1)
+ # [batch_size, seq_length, seq_length]
+ bbox_pos_mask_final = torch.matmul(
+ MT.to(bbox_pos_scores.dtype), M1.to(bbox_pos_scores.dtype))
+ else:
+ bbox_pos_mask_final = None
+
+ if bbox_pos_mask_final is not None:
+ bbox_pos_scores = torch.mul(bbox_pos_scores,
+ bbox_pos_mask_final.unsqueeze(1))
+
+ # [batch_size, d_head, seq_length, seq_length]
+ attention_scores = attention_scores + bbox_pos_scores
+
+ attention_scores = attention_scores / math.sqrt(
+ self.attention_head_size)
+ if attention_mask is not None:
+ # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function)
+ attention_scores = attention_scores + attention_mask
+
+ # Normalize the attention scores to probabilities.
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attention_probs = self.dropout(attention_probs)
+
+ # Mask heads if we want to
+ if head_mask is not None:
+ attention_probs = attention_probs * head_mask
+
+ context_layer = torch.matmul(attention_probs, value_layer)
+
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (
+ self.all_head_size, )
+ context_layer = context_layer.view(new_context_layer_shape)
+
+ outputs = (context_layer,
+ attention_probs) if output_attentions else (context_layer, )
+
+ if self.is_decoder:
+ outputs = outputs + (past_key_value, )
+ return outputs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
+class LayoutRobertaSelfOutput(nn.Module):
+
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(
+ config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states, input_tensor):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Roberta
+class LayoutRobertaAttention(nn.Module):
+
+ def __init__(self, config, position_embedding_type=None):
+ super().__init__()
+ self.self = LayoutRobertaSelfAttention(
+ config, position_embedding_type=position_embedding_type)
+ self.output = LayoutRobertaSelfOutput(config)
+ self.pruned_heads = set()
+
+ def prune_heads(self, heads):
+ if len(heads) == 0:
+ return
+ heads, index = find_pruneable_heads_and_indices(
+ heads, self.self.num_attention_heads,
+ self.self.attention_head_size, self.pruned_heads)
+
+ # Prune linear layers
+ self.self.query = prune_linear_layer(self.self.query, index)
+ self.self.key = prune_linear_layer(self.self.key, index)
+ self.self.value = prune_linear_layer(self.self.value, index)
+ self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+ # Update hyper params and store pruned heads
+ self.self.num_attention_heads = self.self.num_attention_heads - len(
+ heads)
+ self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
+ self.pruned_heads = self.pruned_heads.union(heads)
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ head_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_value=None,
+ output_attentions=False,
+ bbox_pos_emb=None,
+ bbox_pos_mask=None,
+ ):
+ self_outputs = self.self(
+ hidden_states,
+ attention_mask,
+ head_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ past_key_value,
+ output_attentions,
+ bbox_pos_emb=bbox_pos_emb,
+ bbox_pos_mask=bbox_pos_mask,
+ )
+ attention_output = self.output(self_outputs[0], hidden_states)
+ outputs = (attention_output,
+ ) + self_outputs[1:] # add attentions if we output them
+ return outputs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate
+class LayoutRobertaIntermediate(nn.Module):
+
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOutput
+class LayoutRobertaOutput(nn.Module):
+
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(
+ config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states, input_tensor):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Roberta
+class LayoutRobertaLayer(nn.Module):
+
+ def __init__(self, config):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = LayoutRobertaAttention(config)
+ self.is_decoder = config.is_decoder
+ self.add_cross_attention = config.add_cross_attention
+ if self.add_cross_attention:
+ if not self.is_decoder:
+ raise ValueError(
+ f'{self} should be used as a decoder model if cross attention is added'
+ )
+ self.crossattention = LayoutRobertaAttention(
+ config, position_embedding_type='absolute')
+ self.intermediate = LayoutRobertaIntermediate(config)
+ self.output = LayoutRobertaOutput(config)
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ bbox_pos_emb=None,
+ bbox_pos_mask=None,
+ head_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_value=None,
+ output_attentions=False,
+ ):
+ # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
+ self_attn_past_key_value = past_key_value[:
+ 2] if past_key_value is not None else None
+ self_attention_outputs = self.attention(
+ hidden_states,
+ attention_mask,
+ head_mask,
+ output_attentions=output_attentions,
+ past_key_value=self_attn_past_key_value,
+ bbox_pos_emb=bbox_pos_emb,
+ bbox_pos_mask=bbox_pos_mask,
+ )
+ attention_output = self_attention_outputs[0]
+
+ # if decoder, the last output is tuple of self-attn cache
+ if self.is_decoder:
+ outputs = self_attention_outputs[1:-1]
+ present_key_value = self_attention_outputs[-1]
+ else:
+ outputs = self_attention_outputs[
+ 1:] # add self attentions if we output attention weights
+
+ cross_attn_present_key_value = None
+ if self.is_decoder and encoder_hidden_states is not None:
+ if not hasattr(self, 'crossattention'):
+ raise ValueError(
+ f'If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers'
+ ' by setting `config.add_cross_attention=True`')
+
+ # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
+ cross_attn_past_key_value = past_key_value[
+ -2:] if past_key_value is not None else None
+ cross_attention_outputs = self.crossattention(
+ attention_output,
+ attention_mask,
+ head_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ cross_attn_past_key_value,
+ output_attentions,
+ )
+ attention_output = cross_attention_outputs[0]
+ outputs = outputs + cross_attention_outputs[
+ 1:-1] # add cross attentions if we output attention weights
+
+ # add cross-attn cache to positions 3,4 of present_key_value tuple
+ cross_attn_present_key_value = cross_attention_outputs[-1]
+ present_key_value = present_key_value + cross_attn_present_key_value
+
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk,
+ self.chunk_size_feed_forward,
+ self.seq_len_dim,
+ attention_output,
+ )
+ outputs = (layer_output, ) + outputs
+
+ # if decoder, return the attn key/values as the last output
+ if self.is_decoder:
+ outputs = outputs + (present_key_value, )
+
+ return outputs
+
+ def feed_forward_chunk(self, attention_output):
+ intermediate_output = self.intermediate(attention_output)
+ layer_output = self.output(intermediate_output, attention_output)
+ return layer_output
+
+
+# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Roberta
+class LayoutRobertaEncoder(nn.Module):
+
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([
+ LayoutRobertaLayer(config) for _ in range(config.num_hidden_layers)
+ ])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ head_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_values=None,
+ use_cache=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ bbox_pos_emb=None,
+ bbox_pos_mask=None,
+ ):
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attentions = () if output_attentions else None
+ all_cross_attentions = (
+ ) if output_attentions and self.config.add_cross_attention else None
+
+ next_decoder_cache = () if use_cache else None
+ for i, layer_module in enumerate(self.layer):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states, )
+
+ layer_head_mask = head_mask[i] if head_mask is not None else None
+ past_key_value = past_key_values[
+ i] if past_key_values is not None else None
+
+ if self.gradient_checkpointing and self.training:
+
+ if use_cache:
+ logger.warning(
+ '`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...'
+ )
+ use_cache = False
+
+ def create_custom_forward(module):
+
+ def custom_forward(*inputs):
+ return module(*inputs, past_key_value,
+ output_attentions)
+
+ return custom_forward
+
+ layer_outputs = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(layer_module),
+ hidden_states,
+ attention_mask,
+ bbox_pos_emb,
+ bbox_pos_mask,
+ layer_head_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ )
+ else:
+ layer_outputs = layer_module(
+ hidden_states,
+ attention_mask,
+ bbox_pos_emb,
+ bbox_pos_mask,
+ layer_head_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ past_key_value,
+ output_attentions,
+ )
+
+ hidden_states = layer_outputs[0]
+ if use_cache:
+ next_decoder_cache += (layer_outputs[-1], )
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (
+ layer_outputs[1], )
+ if self.config.add_cross_attention:
+ all_cross_attentions = all_cross_attentions + (
+ layer_outputs[2], )
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states, )
+
+ if not return_dict:
+ return tuple(v for v in [
+ hidden_states,
+ next_decoder_cache,
+ all_hidden_states,
+ all_self_attentions,
+ all_cross_attentions,
+ ] if v is not None)
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=next_decoder_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ cross_attentions=all_cross_attentions,
+ )
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPooler
+class LayoutRobertaPooler(nn.Module):
+
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = nn.Tanh()
+
+ def forward(self, hidden_states):
+ # We "pool" the model by simply taking the hidden state corresponding
+ # to the first token.
+ first_token_tensor = hidden_states[:, 0]
+ pooled_output = self.dense(first_token_tensor)
+ pooled_output = self.activation(pooled_output)
+ return pooled_output
+
+
+class LayoutRobertaPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = LayoutRobertaConfig
+ base_model_prefix = 'layoutroberta'
+ supports_gradient_checkpointing = True
+ _keys_to_ignore_on_load_missing = [r'position_ids']
+
+ # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ if isinstance(module, nn.Linear):
+ # Slightly different from the TF version which uses truncated_normal for initialization
+ # cf https://github.com/pytorch/pytorch/pull/5617
+ module.weight.data.normal_(
+ mean=0.0, std=self.config.initializer_range)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(
+ mean=0.0, std=self.config.initializer_range)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+ elif isinstance(module, nn.LayerNorm):
+ module.bias.data.zero_()
+ module.weight.data.fill_(1.0)
+
+ def _set_gradient_checkpointing(self, module, value=False):
+ if isinstance(module, LayoutRobertaEncoder):
+ module.gradient_checkpointing = value
+
+ def update_keys_to_ignore(self, config, del_keys_to_ignore):
+ """Remove some keys from ignore list"""
+ if not config.tie_word_embeddings:
+ # must make a new list, or the class variable gets modified!
+ self._keys_to_ignore_on_save = [
+ k for k in self._keys_to_ignore_on_save
+ if k not in del_keys_to_ignore
+ ]
+ self._keys_to_ignore_on_load_missing = [
+ k for k in self._keys_to_ignore_on_load_missing
+ if k not in del_keys_to_ignore
+ ]
+
+
+class LayoutRobertaModel(LayoutRobertaPreTrainedModel):
+ """
+
+ BROS + Roberta
+
+ """
+
+ # Copied from transformers.models.bert.modeling_bert.BertModel.__init__ with Bert->Roberta
+ def __init__(self, config, add_pooling_layer=True):
+ super().__init__(config)
+ self.config = config
+
+ self.embeddings = LayoutRobertaEmbeddings(config)
+ self.encoder = LayoutRobertaEncoder(config)
+
+ self.pooler = LayoutRobertaPooler(
+ config) if add_pooling_layer else None
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value):
+ self.embeddings.word_embeddings = value
+
+ def _prune_heads(self, heads_to_prune):
+ """
+ Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+ class PreTrainedModel
+ """
+ for layer, heads in heads_to_prune.items():
+ self.encoder.layer[layer].attention.prune_heads(heads)
+
+ # Copied from transformers.models.bert.modeling_bert.BertModel.forward
+ def forward(self,
+ input_ids=None,
+ bbox=None,
+ bbox_mask=None,
+ attention_mask=None,
+ token_type_ids=None,
+ position_ids=None,
+ head_mask=None,
+ inputs_embeds=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_values=None,
+ use_cache=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ **kwargs):
+ r"""
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+ the model is configured as a decoder.
+ encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+ the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple
+ having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+ Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+
+ If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
+ don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
+ `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+ `past_key_values`).
+ """
+ output_attentions = (
+ output_attentions if output_attentions is not None else
+ self.config.output_attentions)
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else
+ self.config.output_hidden_states)
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ if self.config.is_decoder:
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+ else:
+ use_cache = False
+
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError(
+ 'You cannot specify both input_ids and inputs_embeds at the same time'
+ )
+ elif input_ids is not None:
+ input_shape = input_ids.size()
+ elif inputs_embeds is not None:
+ input_shape = inputs_embeds.size()[:-1]
+ else:
+ raise ValueError(
+ 'You have to specify either input_ids or inputs_embeds')
+
+ batch_size, seq_length = input_shape
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+ # past_key_values_length
+ past_key_values_length = past_key_values[0][0].shape[
+ 2] if past_key_values is not None else 0
+
+ if attention_mask is None:
+ attention_mask = torch.ones(
+ ((batch_size, seq_length + past_key_values_length)),
+ device=device)
+
+ if token_type_ids is None:
+ if hasattr(self.embeddings, 'token_type_ids'):
+ buffered_token_type_ids = self.embeddings.token_type_ids[:, :
+ seq_length]
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(
+ batch_size, seq_length)
+ token_type_ids = buffered_token_type_ids_expanded
+ else:
+ token_type_ids = torch.zeros(
+ input_shape, dtype=torch.long, device=device)
+
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+ # ourselves in which case we just need to make it broadcastable to all heads.
+ extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(
+ attention_mask, input_shape, device)
+
+ # If a 2D or 3D attention mask is provided for the cross-attention
+ # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+ if self.config.is_decoder and encoder_hidden_states is not None:
+ encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size(
+ )
+ encoder_hidden_shape = (encoder_batch_size,
+ encoder_sequence_length)
+ if encoder_attention_mask is None:
+ encoder_attention_mask = torch.ones(
+ encoder_hidden_shape, device=device)
+ encoder_extended_attention_mask = self.invert_attention_mask(
+ encoder_attention_mask)
+ else:
+ encoder_extended_attention_mask = None
+
+ # Prepare head mask if needed
+ # 1.0 in head_mask indicate we keep the head
+ # attention_probs has shape bsz x n_heads x N x N
+ # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+ # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+ head_mask = self.get_head_mask(head_mask,
+ self.config.num_hidden_layers)
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ past_key_values_length=past_key_values_length,
+ **kwargs)
+
+ scaled_bbox = bbox * self.config.bbox_scale
+ bbox_pos_emb = self.embeddings.calc_bbox_pos_emb(
+ scaled_bbox, self.config.pe_type)
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=extended_attention_mask,
+ head_mask=head_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_extended_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ bbox_pos_emb=bbox_pos_emb,
+ bbox_pos_mask=bbox_mask,
+ )
+ sequence_output = encoder_outputs[0]
+ pooled_output = self.pooler(
+ sequence_output) if self.pooler is not None else None
+
+ if not return_dict:
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPoolingAndCrossAttentions(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ past_key_values=encoder_outputs.past_key_values,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ cross_attentions=encoder_outputs.cross_attentions,
+ )
+
+
+def create_position_ids_from_input_ids(input_ids,
+ padding_idx,
+ past_key_values_length=0):
+ """
+ Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
+ are ignored. This is modified from fairseq's `utils.make_positions`.
+
+ Args:
+ x: torch.Tensor x:
+
+ Returns: torch.Tensor
+ """
+ # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
+ mask = input_ids.ne(padding_idx).int()
+ incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask)
+ + past_key_values_length) * mask
+ return incremental_indices.long() + padding_idx
diff --git a/modelscope/models/multi_modal/vldoc/processing.py b/modelscope/models/multi_modal/vldoc/processing.py
new file mode 100644
index 00000000..afef8bdb
--- /dev/null
+++ b/modelscope/models/multi_modal/vldoc/processing.py
@@ -0,0 +1,538 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+"""
+Processor class for GeoLayoutLM.
+"""
+
+from collections import defaultdict
+from typing import Dict, Iterable, List, Union
+
+import cv2
+import numpy as np
+import PIL
+import torch
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from torchvision import transforms
+
+from modelscope.preprocessors.image import LoadImage
+
+
+def custom_tokenize(tokenizer, text):
+ toks = tokenizer.tokenize('pad ' + text)[1:]
+ toks2 = toks[1:] if len(toks) > 0 and toks[0] == '▁' else toks
+ return toks2
+
+
+class ImageProcessor(object):
+ r"""
+ Construct a GeoLayoutLM image processor
+ Args:
+ do_preprocess (`bool`): whether to do preprocess to unify the image format,
+ resize and convert to tensor.
+ do_rescale: only works when we disable do_preprocess.
+ """
+
+ def __init__(self,
+ do_preprocess: bool = True,
+ do_resize: bool = False,
+ image_size: Dict[str, int] = None,
+ do_rescale: bool = False,
+ rescale_factor: float = 1. / 255,
+ do_normalize: bool = True,
+ image_mean: Union[float, Iterable[float]] = None,
+ image_std: Union[float, Iterable[float]] = None,
+ apply_ocr: bool = True,
+ **kwargs) -> None:
+ self.do_preprocess = do_preprocess
+ self.do_resize = do_resize
+ self.size = image_size if image_size is not None else {
+ 'height': 768,
+ 'width': 768
+ }
+ self.do_rescale = do_rescale and (not do_preprocess)
+ self.rescale_factor = rescale_factor
+ self.do_normalize = do_normalize
+ image_mean = IMAGENET_DEFAULT_MEAN if image_mean is None else image_mean
+ image_std = IMAGENET_DEFAULT_STD if image_std is None else image_std
+ self.image_mean = (image_mean, image_mean, image_mean) if isinstance(
+ image_mean, float) else image_mean
+ self.image_std = (image_std, image_std, image_std) if isinstance(
+ image_std, float) else image_std
+ self.apply_ocr = apply_ocr
+ self.kwargs = kwargs
+
+ self.totensor = transforms.ToTensor()
+
+ def preprocess(self, image: Union[np.ndarray, PIL.Image.Image]):
+ """ unify the image format, resize and convert to tensor.
+ """
+ image = LoadImage.convert_to_ndarray(image)[:, :, ::-1]
+ size_raw = image.shape[:2]
+ if self.do_resize:
+ image = cv2.resize(image,
+ (self.size['width'], self.size['height']))
+ # convert to pytorch tensor
+ image_pt = self.totensor(image)
+ return image_pt, size_raw
+
+ def __call__(self, images: Union[list, np.ndarray, PIL.Image.Image, str]):
+ """
+ Args:
+ images: list of np.ndarrays, PIL images or image tensors.
+ """
+ if not isinstance(images, list):
+ images = [images]
+ sizes_raw = []
+ if self.do_preprocess:
+ for i in range(len(images)):
+ images[i], size_raw = self.preprocess(images[i])
+ sizes_raw.append(size_raw)
+ images_pt = torch.stack(images, dim=0) # [b, c, h, w]
+ if self.do_rescale:
+ images_pt = images_pt * self.rescale_factor
+ if self.do_normalize:
+ mu = torch.tensor(self.image_mean).view(1, 3, 1, 1)
+ std = torch.tensor(self.image_std).view(1, 3, 1, 1)
+ images_pt = (images_pt - mu) / (std + 1e-8)
+
+ # TODO: apply OCR
+ ocr_infos = None
+ if self.apply_ocr:
+ raise NotImplementedError('OCR service is not available yet!')
+ if len(sizes_raw) == 0:
+ sizes_raw = None
+ data = {
+ 'images': images_pt,
+ 'ocr_infos': ocr_infos,
+ 'sizes_raw': sizes_raw
+ }
+ return data
+
+
+class OCRUtils(object):
+
+ def __init__(self):
+ self.version = 'v0'
+
+ def __call__(self, ocr_infos):
+ """
+ sort boxes, filtering or other preprocesses
+ should return sorted ocr_infos
+ """
+ raise NotImplementedError
+
+
+def bound_box(box, height, width):
+ # box: [x_tl, y_tl, x_br, y_br] or ...
+ assert len(box) == 4 or len(box) == 8
+ for i in range(len(box)):
+ if i & 1:
+ box[i] = max(0, min(box[i], height))
+ else:
+ box[i] = max(0, min(box[i], width))
+ return box
+
+
+def bbox2pto4p(box2p):
+ box4p = [
+ box2p[0], box2p[1], box2p[2], box2p[1], box2p[2], box2p[3], box2p[0],
+ box2p[3]
+ ]
+ return box4p
+
+
+def bbox4pto2p(box4p):
+ box2p = [
+ min(box4p[0], box4p[2], box4p[4], box4p[6]),
+ min(box4p[1], box4p[3], box4p[5], box4p[7]),
+ max(box4p[0], box4p[2], box4p[4], box4p[6]),
+ max(box4p[1], box4p[3], box4p[5], box4p[7]),
+ ]
+ return box2p
+
+
+def stack_tensor_dict(tensor_dicts: List[Dict[str, torch.Tensor]]):
+ one_dict = defaultdict(list)
+ for td in tensor_dicts:
+ for k, v in td.items():
+ one_dict[k].append(v)
+ res_dict = {}
+ for k, v in one_dict.items():
+ res_dict[k] = torch.stack(v, dim=0)
+ return res_dict
+
+
+class TextLayoutSerializer(object):
+
+ def __init__(self,
+ max_seq_length: int,
+ max_block_num: int,
+ tokenizer,
+ width=768,
+ height=768,
+ use_roberta_tokenizer: bool = True,
+ ocr_utils: OCRUtils = None):
+ self.version = 'v0'
+
+ self.max_seq_length = max_seq_length
+ self.max_block_num = max_block_num
+ self.tokenizer = tokenizer
+ self.width = width
+ self.height = height
+ self.use_roberta_tokenizer = use_roberta_tokenizer
+ self.ocr_utils = ocr_utils
+
+ self.pad_token_id = tokenizer.pad_token_id
+ self.cls_token_id = tokenizer.bos_token_id
+ self.sep_token_id = tokenizer.eos_token_id
+ self.unk_token_id = tokenizer.unk_token_id
+ self.cls_bbs_word = [0.0] * 8
+ self.cls_bbs_line = [0] * 4
+
+ def label2seq(self, ocr_info: list, label_info: list):
+ raise NotImplementedError
+
+ def serialize_single(
+ self,
+ ocr_info: list = None,
+ input_ids: list = None,
+ bbox_line: List[List] = None,
+ bbox_word: List[List] = None,
+ width: int = 768,
+ height: int = 768,
+ ):
+ r"""
+ Either ocr_info or (input_ids, bbox_line, bbox_word)
+ should be provided.
+ If (input_ids, bbox_line, bbox_word) is provided,
+ convinient plug into the serialization (customization)
+ is offered. The tokens must be organised by blocks and words.
+ Else, ocr_info must be provided, to be parsed
+ to sequences directly (the simplest way).
+ Args:
+ ocr_info: [
+ {"text": "xx", "box": [a,b,c,d],
+ "words": [{"text": "x", "box": [e,f,g,h]}, ...]},
+ ...
+ ]
+ bbox_line: the coordinate value should match the original image
+ (i.e., not be normalized).
+ """
+ if input_ids is not None:
+ assert len(input_ids) == len(bbox_line)
+ assert len(input_ids) == len(bbox_word)
+ input_ids, bbs_word, bbs_line, first_token_idxes, \
+ line_rank_ids, line_rank_inner_ids, word_rank_ids = \
+ self.halfseq2seq(input_ids, bbox_line, bbox_word, width, height)
+ else:
+ assert ocr_info is not None
+ input_ids, bbs_word, bbs_line, first_token_idxes, \
+ line_rank_ids, line_rank_inner_ids, word_rank_ids = \
+ self.ocr_info2seq(ocr_info, width, height)
+
+ token_seq = {}
+ token_seq['input_ids'] = torch.ones(
+ self.max_seq_length, dtype=torch.int64) * self.pad_token_id
+ token_seq['attention_mask'] = torch.zeros(
+ self.max_seq_length, dtype=torch.int64)
+ token_seq['first_token_idxes'] = torch.zeros(
+ self.max_block_num, dtype=torch.int64)
+ token_seq['first_token_idxes_mask'] = torch.zeros(
+ self.max_block_num, dtype=torch.int64)
+ token_seq['bbox_4p_normalized'] = torch.zeros(
+ self.max_seq_length, 8, dtype=torch.float32)
+ token_seq['bbox'] = torch.zeros(
+ self.max_seq_length, 4, dtype=torch.float32)
+ token_seq['line_rank_id'] = torch.zeros(
+ self.max_seq_length, dtype=torch.int64) # start from 1
+ token_seq['line_rank_inner_id'] = torch.ones(
+ self.max_seq_length, dtype=torch.int64) # 1 2 2 3
+ token_seq['word_rank_id'] = torch.zeros(
+ self.max_seq_length, dtype=torch.int64) # start from 1
+
+ # expand using cls and sep tokens
+ sep_bbs_word = [width, height] * 4
+ sep_bbs_line = [width, height] * 2
+ input_ids = [self.cls_token_id] + input_ids + [self.sep_token_id]
+ bbs_line = [self.cls_bbs_line] + bbs_line + [sep_bbs_line]
+ bbs_word = [self.cls_bbs_word] + bbs_word + [sep_bbs_word]
+
+ # assign
+ len_tokens = len(input_ids)
+ len_lines = len(first_token_idxes)
+ token_seq['input_ids'][:len_tokens] = torch.tensor(input_ids)
+ token_seq['attention_mask'][:len_tokens] = 1
+ token_seq['first_token_idxes'][:len_lines] = torch.tensor(
+ first_token_idxes)
+ token_seq['first_token_idxes_mask'][:len_lines] = 1
+ token_seq['line_rank_id'][1:len_tokens
+ - 1] = torch.tensor(line_rank_ids)
+ token_seq['line_rank_inner_id'][1:len_tokens - 1] = torch.tensor(
+ line_rank_inner_ids)
+ token_seq['line_rank_inner_id'] = token_seq[
+ 'line_rank_inner_id'] * token_seq['attention_mask']
+ token_seq['word_rank_id'][1:len_tokens
+ - 1] = torch.tensor(word_rank_ids)
+
+ token_seq['bbox_4p_normalized'][:len_tokens, :] = torch.tensor(
+ bbs_word)
+ # word bbox normalization -> [0, 1]
+ token_seq['bbox_4p_normalized'][:, [0, 2, 4, 6]] = \
+ token_seq['bbox_4p_normalized'][:, [0, 2, 4, 6]] / width
+ token_seq['bbox_4p_normalized'][:, [1, 3, 5, 7]] = \
+ token_seq['bbox_4p_normalized'][:, [1, 3, 5, 7]] / height
+
+ token_seq['bbox'][:len_tokens, :] = torch.tensor(bbs_line)
+ # line bbox -> [0, 1000)
+ token_seq['bbox'][:,
+ [0, 2]] = token_seq['bbox'][:, [0, 2]] / width * 1000
+ token_seq['bbox'][:,
+ [1, 3]] = token_seq['bbox'][:,
+ [1, 3]] / height * 1000
+ token_seq['bbox'] = token_seq['bbox'].long()
+
+ return token_seq
+
+ def ocr_info2seq(self, ocr_info: list, width: int, height: int):
+ input_ids = []
+ bbs_word = []
+ bbs_line = []
+ first_token_idxes = []
+ line_rank_ids = []
+ line_rank_inner_ids = []
+ word_rank_ids = []
+
+ early_stop = False
+ for line_idx, line in enumerate(ocr_info):
+ if line_idx == self.max_block_num:
+ early_stop = True
+ if early_stop:
+ break
+ lbox = line['box']
+ lbox = bound_box(lbox, height, width)
+ is_first_word = True
+ for word_id, word_info in enumerate(line['words']):
+ wtext = word_info['text']
+ wbox = word_info['box']
+ wbox = bound_box(wbox, height, width)
+ wbox4p = bbox2pto4p(wbox)
+ if self.use_roberta_tokenizer:
+ wtokens = custom_tokenize(self.tokenizer, wtext)
+ else:
+ wtokens = self.tokenizer.tokenize(wtext)
+ wtoken_ids = self.tokenizer.convert_tokens_to_ids(wtokens)
+ if len(wtoken_ids) == 0:
+ wtoken_ids.append(self.unk_token_id)
+ n_tokens = len(wtoken_ids)
+ # reserve for cls and sep
+ if len(input_ids) + n_tokens > self.max_seq_length - 2:
+ early_stop = True
+ break # chunking early for long documents
+ if is_first_word:
+ first_token_idxes.append(len(input_ids) + 1)
+ input_ids.extend(wtoken_ids)
+ bbs_word.extend([wbox4p] * n_tokens)
+ bbs_line.extend([lbox] * n_tokens)
+ word_rank_ids.extend([word_id + 1] * n_tokens)
+ line_rank_ids.extend([line_idx + 1] * n_tokens)
+ if is_first_word:
+ if len(line_rank_inner_ids
+ ) > 0 and line_rank_inner_ids[-1] == 2:
+ line_rank_inner_ids[-1] = 3
+ line_rank_inner_ids.extend([1] + (n_tokens - 1) * [2])
+ is_first_word = False
+ else:
+ line_rank_inner_ids.extend(n_tokens * [2])
+ if len(line_rank_inner_ids) > 0 and line_rank_inner_ids[-1] == 2:
+ line_rank_inner_ids[-1] = 3
+
+ return input_ids, bbs_word, bbs_line, first_token_idxes, line_rank_ids, \
+ line_rank_inner_ids, word_rank_ids
+
+ def halfseq2seq(self, input_ids: list, bbox_line: List[List],
+ bbox_word: List[List], width: int, height: int):
+ """
+ for convinient plug into the serialization, given the 3 customized sequences.
+ They should not contain special tokens like [CLS] or [SEP].
+ """
+ bbs_word = []
+ bbs_line = []
+ first_token_idxes = []
+ line_rank_ids = []
+ line_rank_inner_ids = []
+ word_rank_ids = []
+
+ n_real_tokens = len(input_ids)
+ lb_prev, wb_prev = None, None
+ line_id = 0
+ word_id = 1
+ for i in range(n_real_tokens):
+ lb_now = bbox_line[i]
+ wb_now = bbox_word[i]
+ line_start = lb_prev is None or lb_now != lb_prev
+ word_start = wb_prev is None or wb_now != wb_prev
+ lb_prev, wb_prev = lb_now, wb_now
+
+ if len(lb_now) == 8:
+ lb_now = bbox4pto2p(lb_now)
+ assert len(lb_now) == 4
+ lb_now = bound_box(lb_now, height, width)
+ if len(wb_now) == 4:
+ wb_now = bbox2pto4p(wb_now)
+ assert len(wb_now) == 8
+ wb_now = bound_box(wb_now, height, width)
+
+ bbs_word.append(wb_now)
+ bbs_line.append(lb_now)
+
+ if word_start:
+ word_id += 1
+ if line_start:
+ line_id += 1
+ first_token_idxes.append(i + 1)
+ if len(line_rank_inner_ids
+ ) > 0 and line_rank_inner_ids[-1] == 2:
+ line_rank_inner_ids[-1] = 3
+ line_rank_inner_ids.append(1)
+ word_id = 1
+ else:
+ line_rank_inner_ids.append(2)
+ line_rank_ids.append(line_id)
+ word_rank_ids.append(word_id)
+
+ if len(line_rank_inner_ids) > 0 and line_rank_inner_ids[-1] == 2:
+ line_rank_inner_ids[-1] = 3
+
+ return input_ids, bbs_word, bbs_line, first_token_idxes, \
+ line_rank_ids, line_rank_inner_ids, word_rank_ids
+
+ def __call__(
+ self,
+ ocr_infos: List[List] = None,
+ input_ids: list = None,
+ bboxes_line: List[List] = None,
+ bboxes_word: List[List] = None,
+ sizes_raw: list = None,
+ **kwargs,
+ ):
+ n_samples = len(ocr_infos) if ocr_infos is not None else len(input_ids)
+ if sizes_raw is None:
+ sizes_raw = [(self.height, self.width)] * n_samples
+ seqs = []
+ if input_ids is not None:
+ assert len(input_ids) == len(bboxes_line)
+ assert len(input_ids) == len(bboxes_word)
+ for input_id, bbox_line, bbox_word, size_raw in zip(
+ input_ids, bboxes_line, bboxes_word, sizes_raw):
+ height, width = size_raw
+ token_seq = self.serialize_single(None, input_id, bbox_line,
+ bbox_word, width, height)
+ seqs.append(token_seq)
+ else:
+ assert ocr_infos is not None, 'For serialization, ocr_infos must not be NoneType!'
+ if self.ocr_utils is not None:
+ ocr_infos = self.ocr_utils(ocr_infos)
+ for ocr_info, size_raw in zip(ocr_infos, sizes_raw):
+ height, width = size_raw
+ token_seq = self.serialize_single(
+ ocr_info, width=width, height=height)
+ seqs.append(token_seq)
+ pt_seqs = stack_tensor_dict(seqs)
+ return pt_seqs
+
+
+class Processor(object):
+ r"""Construct a GeoLayoutLM processor.
+
+ Args:
+ max_seq_length: max length for token
+ max_block_num: max number of text lines (blocks or segments)
+ img_processor: type of ImageProcessor.
+ tokenizer: to tokenize strings.
+ use_roberta_tokenizer: Whether the tokenizer is originated from RoBerta tokenizer
+ (True by default).
+ ocr_utils: a tool to preprocess ocr_infos.
+ width: default width. It can be used only when all the images are of the same shape.
+ height: default height. It can be used only when all the images are of the same shape.
+
+ In `serialize_from_tokens`, the 3 sequences (i.e., `input_ids`, `bboxes_line`, `bboxes_word`)
+ must not contain special tokens like [CLS] or [SEP].
+ The boxes in `bboxes_line` and `bboxes_word` can be presented by either 2 points or 4 points.
+ The value in boxes should keep original.
+ Here is an example of the 3 arguments:
+ ```
+ input_ids ->
+ [[6, 2391, 6, 31833, 6, 10132, 6, 2283, 6, 17730, 6, 2698, 152]]
+ bboxes_line ->
+ [[[230, 1, 353, 38], [230, 1, 353, 38], [230, 1, 353, 38], [230, 1, 353, 38],
+ [230, 1, 353, 38], [230, 1, 353, 38], [230, 1, 353, 38], [230, 1, 353, 38],
+ [257, 155, 338, 191], [257, 155, 338, 191], [257, 155, 338, 191], [257, 155, 338, 191],
+ [257, 155, 338, 191]]]
+ bboxes_word ->
+ [[[231, 2, 267, 2, 267, 38, 231, 38], [231, 2, 267, 2, 267, 38, 231, 38],
+ [264, 7, 298, 7, 298, 36, 264, 36], [264, 7, 298, 7, 298, 36, 264, 36],
+ [293, 3, 329, 3, 329, 41, 293, 41], [293, 3, 329, 3, 329, 41, 293, 41],
+ [330, 4, 354, 4, 354, 39, 330, 39], [330, 4, 354, 4, 354, 39, 330, 39],
+ [258, 156, 289, 156, 289, 193, 258, 193], [258, 156, 289, 156, 289, 193, 258, 193],
+ [288, 158, 321, 158, 321, 192, 288, 192], [288, 158, 321, 158, 321, 192, 288, 192],
+ [321, 156, 336, 156, 336, 190, 321, 190]]]
+ ```
+
+ """
+
+ def __init__(self,
+ max_seq_length,
+ max_block_num,
+ img_processor: ImageProcessor,
+ tokenizer=None,
+ use_roberta_tokenizer: bool = True,
+ ocr_utils: OCRUtils = None,
+ width=768,
+ height=768,
+ **kwargs):
+ self.img_processor = img_processor
+ self.tokenizer = tokenizer
+ self.kwargs = kwargs
+
+ self.serializer = TextLayoutSerializer(
+ max_seq_length,
+ max_block_num,
+ tokenizer,
+ width,
+ height,
+ use_roberta_tokenizer=use_roberta_tokenizer,
+ ocr_utils=ocr_utils)
+
+ def __call__(
+ self,
+ images: Union[list, np.ndarray, PIL.Image.Image, str],
+ ocr_infos: List[List] = None,
+ token_seqs: dict = None,
+ sizes_raw: list = None,
+ ):
+ img_data = self.img_processor(images)
+ images = img_data['images']
+ ocr_infos = img_data['ocr_infos'] if ocr_infos is None else ocr_infos
+ sizes_raw = img_data['sizes_raw'] if sizes_raw is None else sizes_raw
+ if token_seqs is None:
+ token_seqs = self.serializer(ocr_infos, sizes_raw=sizes_raw)
+ else:
+ token_seqs = self.serializer(
+ None, sizes_raw=sizes_raw, **token_seqs)
+ assert token_seqs is not None, 'token_seqs must not be NoneType!'
+ batch = {}
+ batch['image'] = images
+ for k, v in token_seqs.items():
+ batch[k] = token_seqs[k]
+ return batch
+
+ def serialize_from_tokens(self,
+ images,
+ input_ids,
+ bboxes_line,
+ bboxes_word,
+ sizes_raw=None):
+ half_batch = {}
+ half_batch['input_ids'] = input_ids
+ half_batch['bboxes_line'] = bboxes_line
+ half_batch['bboxes_word'] = bboxes_word
+ return self(images, None, half_batch, sizes_raw)
diff --git a/modelscope/models/multi_modal/vldoc/tokenization.py b/modelscope/models/multi_modal/vldoc/tokenization.py
new file mode 100644
index 00000000..a16c5849
--- /dev/null
+++ b/modelscope/models/multi_modal/vldoc/tokenization.py
@@ -0,0 +1,89 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import os
+
+from transformers import XLMRobertaTokenizer
+
+SPIECE_UNDERLINE = '▁'
+
+
+class VLDocXLMTokenizer(XLMRobertaTokenizer):
+ """
+ Adapted from [`RobertaTokenizer`] and [`XLNetTokenizer`]. Based on
+ [SentencePiece](https://github.com/google/sentencepiece).
+
+ This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+ this superclass for more information regarding those methods.
+
+ Args:
+ vocab_file (`str`):
+ Path to the vocabulary file.
+ bos_token (`str`, *optional*, defaults to `""`):
+ The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
+
+
+
+ When building a sequence using special tokens, this is not the token that is used for the beginning of
+ sequence. The token used is the `cls_token`.
+
+
+
+ eos_token (`str`, *optional*, defaults to `""`):
+ The end of sequence token.
+
+
+
+ When building a sequence using special tokens, this is not the token that is used for the end of sequence.
+ The token used is the `sep_token`.
+
+
+
+ sep_token (`str`, *optional*, defaults to `""`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ cls_token (`str`, *optional*, defaults to `""`):
+ The classifier token which is used when doing sequence classification (classification of the whole sequence
+ instead of per-token classification). It is the first token of the sequence when built with special tokens.
+ unk_token (`str`, *optional*, defaults to `""`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ pad_token (`str`, *optional*, defaults to `""`):
+ The token used for padding, for example when batching sequences of different lengths.
+ mask_token (`str`, *optional*, defaults to `""`):
+ The token used for masking values. This is the token used when training this model with masked language
+ modeling. This is the token which the model will try to predict.
+ cls_token_box (`List[int]`, *optional*, defaults to `[0, 0, 0, 0]`):
+ The bounding box to use for the special [CLS] token.
+ sep_token_box (`List[int]`, *optional*, defaults to `[1000, 1000, 1000, 1000]`):
+ The bounding box to use for the special [SEP] token.
+ pad_token_box (`List[int]`, *optional*, defaults to `[0, 0, 0, 0]`):
+ The bounding box to use for the special [PAD] token.
+ pad_token_label (`int`, *optional*, defaults to -100):
+ The label to use for padding tokens. Defaults to -100, which is the `ignore_index` of PyTorch's
+ CrossEntropyLoss.
+ only_label_first_subword (`bool`, *optional*, defaults to `True`):
+ Whether or not to only label the first subword, in case word labels are provided.
+ additional_special_tokens (`List[str]`, *optional*, defaults to `["NOTUSED", "NOTUSED"]`):
+ Additional special tokens used by the tokenizer.
+ sp_model_kwargs (`dict`, *optional*):
+ Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
+ SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
+ to set:
+
+ - `enable_sampling`: Enable subword regularization.
+ - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
+
+ - `nbest_size = {0,1}`: No sampling is performed.
+ - `nbest_size > 1`: samples from the nbest_size results.
+ - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
+ using forward-filtering-and-backward-sampling algorithm.
+
+ - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
+ BPE-dropout.
+
+ Attributes:
+ sp_model (`SentencePieceProcessor`):
+ The *SentencePiece* processor that is used for every conversion (string, tokens and IDs).
+ """
+ model_input_names = ['input_ids', 'attention_mask']
diff --git a/modelscope/models/multi_modal/vldoc/transformer_local.py b/modelscope/models/multi_modal/vldoc/transformer_local.py
new file mode 100644
index 00000000..4c0dd55d
--- /dev/null
+++ b/modelscope/models/multi_modal/vldoc/transformer_local.py
@@ -0,0 +1,204 @@
+# The implementation is borrowed and modified from the official PyTorch website and ABINet:
+# https://pytorch.org/docs/stable/_modules/torch/nn/modules/transformer.html
+# https://github.com/FangShancheng/ABINet/blob/main/modules/transformer.py
+
+import copy
+
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import Dropout, LayerNorm, Linear, Module, ModuleList
+from torch.nn import functional as F
+
+
+class TransformerDecoder(Module):
+ r"""TransformerDecoder is a stack of N decoder layers
+
+ Args:
+ decoder_layer: an instance of the TransformerDecoderLayer() class (required).
+ num_layers: the number of sub-decoder-layers in the decoder (required).
+ norm: the layer normalization component (optional).
+
+ Examples::
+ >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
+ >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
+ >>> memory = torch.rand(10, 32, 512)
+ >>> tgt = torch.rand(20, 32, 512)
+ >>> out = transformer_decoder(tgt, memory)
+ """
+ __constants__ = ['norm']
+
+ def __init__(self, decoder_layer, num_layers, norm=None):
+ super(TransformerDecoder, self).__init__()
+ self.layers = _get_clones(decoder_layer, num_layers)
+ self.num_layers = num_layers
+ self.norm = norm
+
+ def forward(self,
+ tgt,
+ memory,
+ memory2=None,
+ tgt_mask=None,
+ memory_mask=None,
+ memory_mask2=None,
+ tgt_key_padding_mask=None,
+ memory_key_padding_mask=None,
+ memory_key_padding_mask2=None):
+ r"""Pass the inputs (and mask) through the decoder layer in turn.
+
+ Args:
+ tgt: the sequence to the decoder (required).
+ memory: the sequence from the last layer of the encoder (required).
+ tgt_mask: the mask for the tgt sequence (optional).
+ memory_mask: the mask for the memory sequence (optional).
+ tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+ memory_key_padding_mask: the mask for the memory keys per batch (optional).
+
+ Shape:
+ see the docs in Transformer class.
+ """
+ output = tgt
+
+ for mod in self.layers:
+ output = mod(
+ output,
+ memory,
+ memory2=memory2,
+ tgt_mask=tgt_mask,
+ memory_mask=memory_mask,
+ memory_mask2=memory_mask2,
+ tgt_key_padding_mask=tgt_key_padding_mask,
+ memory_key_padding_mask=memory_key_padding_mask,
+ memory_key_padding_mask2=memory_key_padding_mask2)
+
+ if self.norm is not None:
+ output = self.norm(output)
+
+ return output
+
+
+class TransformerDecoderLayer(Module):
+ r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
+
+ Args:
+ d_model: the number of expected features in the input (required).
+ nhead: the number of heads in the multiheadattention models (required).
+ dim_feedforward: the dimension of the feedforward network model (default=2048).
+ dropout: the dropout value (default=0.1).
+ activation: the activation function of intermediate layer, relu or gelu (default=relu).
+ """
+
+ def __init__(self,
+ d_model,
+ nhead,
+ dim_feedforward=2048,
+ dropout=0.1,
+ activation='relu',
+ self_attn=True,
+ siamese=False,
+ debug=False):
+ super(TransformerDecoderLayer, self).__init__()
+ self.has_self_attn, self.siamese = self_attn, siamese
+ self.debug = debug
+ if self.has_self_attn:
+ self.self_attn = nn.MultiheadAttention(
+ d_model, nhead, dropout=dropout)
+ self.norm1 = LayerNorm(d_model)
+ self.dropout1 = Dropout(dropout)
+ self.multihead_attn = nn.MultiheadAttention(
+ d_model, nhead, dropout=dropout)
+ # Implementation of Feedforward model
+ self.linear1 = Linear(d_model, dim_feedforward)
+ self.dropout = Dropout(dropout)
+ self.linear2 = Linear(dim_feedforward, d_model)
+
+ self.norm2 = LayerNorm(d_model)
+ self.norm3 = LayerNorm(d_model)
+ self.dropout2 = Dropout(dropout)
+ self.dropout3 = Dropout(dropout)
+ if self.siamese:
+ self.multihead_attn2 = nn.MultiheadAttention(
+ d_model, nhead, dropout=dropout)
+
+ self.activation = _get_activation_fn(activation)
+
+ def __setstate__(self, state):
+ if 'activation' not in state:
+ state['activation'] = F.relu
+ super(TransformerDecoderLayer, self).__setstate__(state)
+
+ def forward(self,
+ tgt,
+ memory,
+ tgt_mask=None,
+ memory_mask=None,
+ tgt_key_padding_mask=None,
+ memory_key_padding_mask=None,
+ memory2=None,
+ memory_mask2=None,
+ memory_key_padding_mask2=None):
+ r"""Pass the inputs (and mask) through the decoder layer.
+
+ Args:
+ tgt: the sequence to the decoder layer (required).
+ memory: the sequence from the last layer of the encoder (required).
+ tgt_mask: the mask for the tgt sequence (optional).
+ memory_mask: the mask for the memory sequence (optional).
+ tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+ memory_key_padding_mask: the mask for the memory keys per batch (optional).
+
+ Shape:
+ see the docs in Transformer class.
+ """
+ if self.has_self_attn:
+ tgt2, attn = self.self_attn(
+ tgt,
+ tgt,
+ tgt,
+ attn_mask=tgt_mask,
+ key_padding_mask=tgt_key_padding_mask,
+ )
+ tgt = tgt + self.dropout1(tgt2)
+ tgt = self.norm1(tgt)
+ if self.debug:
+ self.attn = attn
+ tgt2, attn2 = self.multihead_attn(
+ tgt,
+ memory,
+ memory,
+ attn_mask=memory_mask,
+ key_padding_mask=memory_key_padding_mask)
+ if self.debug:
+ self.attn2 = attn2
+
+ if self.siamese:
+ tgt3, attn3 = self.multihead_attn2(
+ tgt,
+ memory2,
+ memory2,
+ attn_mask=memory_mask2,
+ key_padding_mask=memory_key_padding_mask2)
+ tgt = tgt + self.dropout2(tgt3)
+ if self.debug:
+ self.attn3 = attn3
+
+ tgt = tgt + self.dropout2(tgt2)
+ tgt = self.norm2(tgt)
+ tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+ tgt = tgt + self.dropout3(tgt2)
+ tgt = self.norm3(tgt)
+
+ return tgt
+
+
+def _get_clones(module, N):
+ return ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation):
+ if activation == 'relu':
+ return F.relu
+ elif activation == 'gelu':
+ return F.gelu
+
+ raise RuntimeError(
+ 'activation should be relu/gelu, not {}'.format(activation))
diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py
index 784dbf71..3fbaba68 100644
--- a/modelscope/outputs/outputs.py
+++ b/modelscope/outputs/outputs.py
@@ -74,6 +74,14 @@ TASK_OUTPUTS = {
# }
Tasks.ocr_recognition: [OutputKeys.TEXT],
+ # document vl embedding for single sample
+ # {
+ # "img_embedding": np.array with shape [M, D],
+ # "text_embedding": np.array with shape [N, D]
+ # }
+ Tasks.document_vl_embedding:
+ [OutputKeys.IMG_EMBEDDING, OutputKeys.TEXT_EMBEDDING],
+
# face 2d keypoint result for single sample
# {
# "keypoints": [
@@ -346,8 +354,9 @@ TASK_OUTPUTS = {
# "output_video": "path_to_rendered_video" , this is optional
# and is only avaialbe when the "render" option is enabled.
# }
- Tasks.body_3d_keypoints:
- [OutputKeys.KEYPOINTS, OutputKeys.TIMESTAMPS, OutputKeys.OUTPUT_VIDEO],
+ Tasks.body_3d_keypoints: [
+ OutputKeys.KEYPOINTS, OutputKeys.TIMESTAMPS, OutputKeys.OUTPUT_VIDEO
+ ],
# 2D hand keypoints result for single sample
# {
diff --git a/modelscope/pipelines/builder.py b/modelscope/pipelines/builder.py
index 58916066..fe2f1c1d 100644
--- a/modelscope/pipelines/builder.py
+++ b/modelscope/pipelines/builder.py
@@ -91,6 +91,9 @@ DEFAULT_MODEL_FOR_PIPELINE = {
Tasks.table_recognition:
(Pipelines.table_recognition,
'damo/cv_dla34_table-structure-recognition_cycle-centernet'),
+ Tasks.document_vl_embedding:
+ (Pipelines.document_vl_embedding,
+ 'damo/multi-modal_convnext-roberta-base_vldoc-embedding'),
Tasks.license_plate_detection:
(Pipelines.license_plate_detection,
'damo/cv_resnet18_license-plate-detection_damo'),
@@ -245,9 +248,9 @@ DEFAULT_MODEL_FOR_PIPELINE = {
Tasks.video_object_segmentation:
(Pipelines.video_object_segmentation,
'damo/cv_rdevos_video-object-segmentation'),
- Tasks.image_multi_view_depth_estimation:
- (Pipelines.image_multi_view_depth_estimation,
- 'damo/cv_casmvs_multi-view-depth-estimation_general'),
+ Tasks.image_multi_view_depth_estimation: (
+ Pipelines.image_multi_view_depth_estimation,
+ 'damo/cv_casmvs_multi-view-depth-estimation_general'),
}
diff --git a/modelscope/pipelines/multi_modal/__init__.py b/modelscope/pipelines/multi_modal/__init__.py
index b16eb360..07e27e14 100644
--- a/modelscope/pipelines/multi_modal/__init__.py
+++ b/modelscope/pipelines/multi_modal/__init__.py
@@ -14,6 +14,7 @@ if TYPE_CHECKING:
VideoMultiModalEmbeddingPipeline
from .visual_question_answering_pipeline import VisualQuestionAnsweringPipeline
from .asr_pipeline import AutomaticSpeechRecognitionPipeline
+ from .document_vl_embedding_pipeline import DocumentVLEmbeddingPipeline
from .video_captioning_pipeline import VideoCaptioningPipeline
from .video_question_answering_pipeline import VideoQuestionAnsweringPipeline
else:
@@ -30,6 +31,7 @@ else:
'generative_multi_modal_embedding_pipeline':
['GEMMMultiModalEmbeddingPipeline'],
'asr_pipeline': ['AutomaticSpeechRecognitionPipeline'],
+ 'document_vl_embedding_pipeline': ['DocumentVLEmbeddingPipeline'],
'video_captioning_pipeline': ['VideoCaptioningPipeline'],
'video_question_answering_pipeline':
['VideoQuestionAnsweringPipeline']
diff --git a/modelscope/pipelines/multi_modal/document_vl_embedding_pipeline.py b/modelscope/pipelines/multi_modal/document_vl_embedding_pipeline.py
new file mode 100644
index 00000000..754d2d2b
--- /dev/null
+++ b/modelscope/pipelines/multi_modal/document_vl_embedding_pipeline.py
@@ -0,0 +1,62 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import Any, Dict, Optional, Union
+
+import torch
+
+from modelscope.metainfo import Pipelines
+from modelscope.models.multi_modal.vldoc.model import VLDocForDocVLEmbedding
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines.base import Input, Model, Pipeline
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.preprocessors.multi_modal import (Preprocessor,
+ VLDocPreprocessor)
+from modelscope.utils.constant import ModelFile, Tasks
+from modelscope.utils.logger import get_logger
+
+logger = get_logger()
+
+
+@PIPELINES.register_module(
+ Tasks.document_vl_embedding, module_name=Pipelines.document_vl_embedding)
+class DocumentVLEmbeddingPipeline(Pipeline):
+
+ def __init__(self,
+ model: Union[Model, str],
+ preprocessor: Optional[Preprocessor] = None,
+ **kwargs):
+ """ The pipeline for multi-modal document embedding generation.
+
+ Args:
+ model: model id on modelscope hub.
+ preprocessor: type `Preprocessor`. If None, `VLDocPreprocessor` is used.
+
+ Example:
+ ```python
+ >>> from modelscope.models import Model
+ >>> from modelscope.pipelines import pipeline
+ >>> model = Model.from_pretrained(
+ 'damo/multi-modal_convnext-roberta-base_vldoc-embedding')
+ >>> doc_VL_emb_pipeline = pipeline(task='document-vl-embedding', model=model)
+ >>> inp = {
+ 'images': ['data/demo.png'],
+ 'ocr_info_paths': ['data/demo.json']
+ }
+ >>> result = doc_VL_emb_pipeline(inp)
+ ```
+ """
+
+ super().__init__(model=model, preprocessor=preprocessor, **kwargs)
+ self.model.eval()
+ if preprocessor is None:
+ if isinstance(self.model, VLDocForDocVLEmbedding):
+ self.preprocessor = VLDocPreprocessor(self.model.model_dir)
+ else:
+ raise NotImplementedError
+
+ def forward(self, encodings: Dict[str, Any]) -> Dict[str, Any]:
+ for k, v in encodings.items():
+ encodings[k] = encodings[k].to(self.device)
+ return self.model(**encodings)
+
+ def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+ return inputs
diff --git a/modelscope/preprocessors/multi_modal.py b/modelscope/preprocessors/multi_modal.py
index 85ef4cdd..4cdca57b 100644
--- a/modelscope/preprocessors/multi_modal.py
+++ b/modelscope/preprocessors/multi_modal.py
@@ -390,6 +390,78 @@ class MPlugPreprocessor(Preprocessor):
return output
+@PREPROCESSORS.register_module(
+ Fields.multi_modal, module_name=Preprocessors.vldoc_preprocessor)
+class VLDocPreprocessor(Preprocessor):
+
+ def __init__(self,
+ model_dir: str,
+ mode: str = ModeKeys.INFERENCE,
+ *args,
+ **kwargs):
+ """Preprocess data for the model `VLDocForDocVLEmbedding`.
+
+ Args:
+ model_dir (str): model path in model hub.
+ mode (str): model mode, in ('train', 'eval', 'inference').
+ """
+ super().__init__(*args, **kwargs)
+
+ self.model_dir = model_dir
+ self.mode = mode
+
+ model_cfg_path = osp.join(model_dir, 'config.json')
+ with open(model_cfg_path, 'r', encoding='utf-8') as f:
+ model_cfg = json.load(f)
+
+ from modelscope.models.multi_modal.vldoc.tokenization import VLDocXLMTokenizer
+ tokenizer_path = osp.join(model_dir, ModelFile.TOKENIZER_FOLDER)
+ self.tokenizer = VLDocXLMTokenizer.from_pretrained(tokenizer_path)
+
+ from modelscope.models.multi_modal.vldoc.processing import Processor, ImageProcessor
+ self.img_proc = ImageProcessor(
+ do_preprocess=True,
+ do_resize=True,
+ image_size={
+ 'height': model_cfg['image_size'][0],
+ 'width': model_cfg['image_size'][1],
+ },
+ do_normalize=True,
+ apply_ocr=False)
+ self.proc = Processor(
+ max_seq_length=model_cfg['max_seq_length'],
+ max_block_num=model_cfg['max_block_num'],
+ img_processor=self.img_proc,
+ tokenizer=self.tokenizer,
+ width=model_cfg['image_size'][1],
+ height=model_cfg['image_size'][0],
+ )
+
+ def __call__(self, input: Dict[str, Any], *args,
+ **kwargs) -> Dict[str, Any]:
+ """
+ Args:
+ input: {
+ 'images': ['img_path1', 'img_path2', ...],
+ 'ocr_info_paths': ['json_path1', 'json_path2', ...]
+ }
+ Return:
+ encodings: Dict[str, Tensor]
+ """
+
+ ocr_infos = []
+ for one_ocr_info_path in input['ocr_info_paths']:
+ with open(one_ocr_info_path, 'r') as f:
+ ocr_info = json.load(f)
+ ocr_info = ocr_info['form']
+ ocr_infos.append(ocr_info)
+
+ proc_input = {'images': input['images'], 'ocr_infos': ocr_infos}
+ encodings = self.proc(**proc_input)
+
+ return encodings
+
+
@PREPROCESSORS.register_module(
Fields.multi_modal, module_name=Preprocessors.hitea_tasks_preprocessor)
class HiTeAPreprocessor(Preprocessor):
diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py
index 1c615d3e..069ddc8b 100644
--- a/modelscope/utils/constant.py
+++ b/modelscope/utils/constant.py
@@ -176,6 +176,7 @@ class MultiModalTasks(object):
visual_entailment = 'visual-entailment'
video_multi_modal_embedding = 'video-multi-modal-embedding'
image_text_retrieval = 'image-text-retrieval'
+ document_vl_embedding = 'document-vl-embedding'
video_captioning = 'video-captioning'
video_question_answering = 'video-question-answering'
@@ -314,6 +315,7 @@ class ModelFile(object):
LABEL_MAPPING = 'label_mapping.json'
TRAIN_OUTPUT_DIR = 'output'
TS_MODEL_FILE = 'model.ts'
+ TOKENIZER_FOLDER = 'tokenizer'
class Invoke(object):
diff --git a/tests/pipelines/test_document_vl_embedding.py b/tests/pipelines/test_document_vl_embedding.py
new file mode 100644
index 00000000..f8d2d5a3
--- /dev/null
+++ b/tests/pipelines/test_document_vl_embedding.py
@@ -0,0 +1,60 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import os.path as osp
+import unittest
+
+import json
+
+from modelscope.hub.snapshot_download import snapshot_download
+from modelscope.models import Model
+from modelscope.pipelines import pipeline
+from modelscope.pipelines.base import Pipeline
+from modelscope.utils.constant import Tasks
+from modelscope.utils.demo_utils import DemoCompatibilityCheck
+from modelscope.utils.test_utils import test_level
+
+
+class DocumentVLEmbeddingTest(unittest.TestCase, DemoCompatibilityCheck):
+
+ def setUp(self) -> None:
+ self.model_id = 'damo/multi-modal_convnext-roberta-base_vldoc-embedding'
+ cache_path = snapshot_download(self.model_id)
+ self.test_image = osp.join(cache_path, 'data/demo.png')
+ self.test_json = osp.join(cache_path, 'data/demo.json')
+ self.task = Tasks.document_vl_embedding
+
+ def pipeline_inference(self, pipe: Pipeline):
+ inp = {'images': [self.test_image], 'ocr_info_paths': [self.test_json]}
+ result = pipe(inp)
+
+ print('Results of VLDoc: ')
+ for k, v in result.items():
+ print(f'{k}: {v.size()}')
+
+ @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+ def test_run_with_model_name(self):
+ doc_VL_emb_pipeline = pipeline(task=self.task, model=self.model_id)
+ self.pipeline_inference(doc_VL_emb_pipeline)
+
+ @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
+ def test_run_with_model_from_modelhub(self):
+ print('test_run_with_model_from_modelhub')
+ model = Model.from_pretrained(self.model_id)
+
+ doc_VL_emb_pipeline = pipeline(task=self.task, model=model)
+ self.pipeline_inference(doc_VL_emb_pipeline)
+
+ @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
+ def test_run_modelhub_default_model(self):
+ print('test_run_modelhub_default_model')
+ # default model: VLDoc
+ vldoc_doc_VL_emb_pipeline = pipeline(self.task)
+ self.pipeline_inference(vldoc_doc_VL_emb_pipeline)
+
+ @unittest.skip('demo compatibility test is only enabled on a needed-basis')
+ def test_demo_compatibility(self):
+ self.compatibility_check()
+
+
+if __name__ == '__main__':
+ unittest.main()