This commit is contained in:
260
ppocr/modeling/backbones/vqa_layoutlm.py
Normal file
260
ppocr/modeling/backbones/vqa_layoutlm.py
Normal file
@@ -0,0 +1,260 @@
|
||||
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
from paddle import nn
|
||||
|
||||
from paddlenlp.transformers import (
|
||||
LayoutXLMModel,
|
||||
LayoutXLMForTokenClassification,
|
||||
LayoutXLMForRelationExtraction,
|
||||
)
|
||||
from paddlenlp.transformers import LayoutLMModel, LayoutLMForTokenClassification
|
||||
from paddlenlp.transformers import (
|
||||
LayoutLMv2Model,
|
||||
LayoutLMv2ForTokenClassification,
|
||||
LayoutLMv2ForRelationExtraction,
|
||||
)
|
||||
from paddlenlp.transformers import AutoModel
|
||||
|
||||
__all__ = ["LayoutXLMForSer", "LayoutLMForSer"]
|
||||
|
||||
pretrained_model_dict = {
|
||||
LayoutXLMModel: {
|
||||
"base": "layoutxlm-base-uncased",
|
||||
"vi": "vi-layoutxlm-base-uncased",
|
||||
},
|
||||
LayoutLMModel: {
|
||||
"base": "layoutlm-base-uncased",
|
||||
},
|
||||
LayoutLMv2Model: {
|
||||
"base": "layoutlmv2-base-uncased",
|
||||
"vi": "vi-layoutlmv2-base-uncased",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class NLPBaseModel(nn.Layer):
|
||||
def __init__(
|
||||
self,
|
||||
base_model_class,
|
||||
model_class,
|
||||
mode="base",
|
||||
type="ser",
|
||||
pretrained=True,
|
||||
checkpoints=None,
|
||||
**kwargs,
|
||||
):
|
||||
super(NLPBaseModel, self).__init__()
|
||||
if checkpoints is not None: # load the trained model
|
||||
self.model = model_class.from_pretrained(checkpoints)
|
||||
else: # load the pretrained-model
|
||||
pretrained_model_name = pretrained_model_dict[base_model_class][mode]
|
||||
if type == "ser":
|
||||
self.model = model_class.from_pretrained(
|
||||
pretrained_model_name, num_classes=kwargs["num_classes"], dropout=0
|
||||
)
|
||||
else:
|
||||
self.model = model_class.from_pretrained(
|
||||
pretrained_model_name, dropout=0
|
||||
)
|
||||
self.out_channels = 1
|
||||
self.use_visual_backbone = True
|
||||
|
||||
|
||||
class LayoutLMForSer(NLPBaseModel):
|
||||
def __init__(
|
||||
self, num_classes, pretrained=True, checkpoints=None, mode="base", **kwargs
|
||||
):
|
||||
super(LayoutLMForSer, self).__init__(
|
||||
LayoutLMModel,
|
||||
LayoutLMForTokenClassification,
|
||||
mode,
|
||||
"ser",
|
||||
pretrained,
|
||||
checkpoints,
|
||||
num_classes=num_classes,
|
||||
)
|
||||
self.use_visual_backbone = False
|
||||
|
||||
def forward(self, x):
|
||||
x = self.model(
|
||||
input_ids=x[0],
|
||||
bbox=x[1],
|
||||
attention_mask=x[2],
|
||||
token_type_ids=x[3],
|
||||
position_ids=None,
|
||||
output_hidden_states=False,
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
class LayoutLMv2ForSer(NLPBaseModel):
|
||||
def __init__(
|
||||
self, num_classes, pretrained=True, checkpoints=None, mode="base", **kwargs
|
||||
):
|
||||
super(LayoutLMv2ForSer, self).__init__(
|
||||
LayoutLMv2Model,
|
||||
LayoutLMv2ForTokenClassification,
|
||||
mode,
|
||||
"ser",
|
||||
pretrained,
|
||||
checkpoints,
|
||||
num_classes=num_classes,
|
||||
)
|
||||
if (
|
||||
hasattr(self.model.layoutlmv2, "use_visual_backbone")
|
||||
and self.model.layoutlmv2.use_visual_backbone is False
|
||||
):
|
||||
self.use_visual_backbone = False
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_visual_backbone is True:
|
||||
image = x[4]
|
||||
else:
|
||||
image = None
|
||||
x = self.model(
|
||||
input_ids=x[0],
|
||||
bbox=x[1],
|
||||
attention_mask=x[2],
|
||||
token_type_ids=x[3],
|
||||
image=image,
|
||||
position_ids=None,
|
||||
head_mask=None,
|
||||
labels=None,
|
||||
)
|
||||
if self.training:
|
||||
res = {"backbone_out": x[0]}
|
||||
res.update(x[1])
|
||||
return res
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class LayoutXLMForSer(NLPBaseModel):
|
||||
def __init__(
|
||||
self, num_classes, pretrained=True, checkpoints=None, mode="base", **kwargs
|
||||
):
|
||||
super(LayoutXLMForSer, self).__init__(
|
||||
LayoutXLMModel,
|
||||
LayoutXLMForTokenClassification,
|
||||
mode,
|
||||
"ser",
|
||||
pretrained,
|
||||
checkpoints,
|
||||
num_classes=num_classes,
|
||||
)
|
||||
if (
|
||||
hasattr(self.model.layoutxlm, "use_visual_backbone")
|
||||
and self.model.layoutxlm.use_visual_backbone is False
|
||||
):
|
||||
self.use_visual_backbone = False
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_visual_backbone is True:
|
||||
image = x[4]
|
||||
else:
|
||||
image = None
|
||||
x = self.model(
|
||||
input_ids=x[0],
|
||||
bbox=x[1],
|
||||
attention_mask=x[2],
|
||||
token_type_ids=x[3],
|
||||
image=image,
|
||||
position_ids=None,
|
||||
head_mask=None,
|
||||
labels=None,
|
||||
)
|
||||
if self.training:
|
||||
res = {"backbone_out": x[0]}
|
||||
res.update(x[1])
|
||||
return res
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class LayoutLMv2ForRe(NLPBaseModel):
|
||||
def __init__(self, pretrained=True, checkpoints=None, mode="base", **kwargs):
|
||||
super(LayoutLMv2ForRe, self).__init__(
|
||||
LayoutLMv2Model,
|
||||
LayoutLMv2ForRelationExtraction,
|
||||
mode,
|
||||
"re",
|
||||
pretrained,
|
||||
checkpoints,
|
||||
)
|
||||
if (
|
||||
hasattr(self.model.layoutlmv2, "use_visual_backbone")
|
||||
and self.model.layoutlmv2.use_visual_backbone is False
|
||||
):
|
||||
self.use_visual_backbone = False
|
||||
|
||||
def forward(self, x):
|
||||
x = self.model(
|
||||
input_ids=x[0],
|
||||
bbox=x[1],
|
||||
attention_mask=x[2],
|
||||
token_type_ids=x[3],
|
||||
image=x[4],
|
||||
position_ids=None,
|
||||
head_mask=None,
|
||||
labels=None,
|
||||
entities=x[5],
|
||||
relations=x[6],
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
class LayoutXLMForRe(NLPBaseModel):
|
||||
def __init__(self, pretrained=True, checkpoints=None, mode="base", **kwargs):
|
||||
super(LayoutXLMForRe, self).__init__(
|
||||
LayoutXLMModel,
|
||||
LayoutXLMForRelationExtraction,
|
||||
mode,
|
||||
"re",
|
||||
pretrained,
|
||||
checkpoints,
|
||||
)
|
||||
if (
|
||||
hasattr(self.model.layoutxlm, "use_visual_backbone")
|
||||
and self.model.layoutxlm.use_visual_backbone is False
|
||||
):
|
||||
self.use_visual_backbone = False
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_visual_backbone is True:
|
||||
image = x[4]
|
||||
entities = x[5]
|
||||
relations = x[6]
|
||||
else:
|
||||
image = None
|
||||
entities = x[4]
|
||||
relations = x[5]
|
||||
x = self.model(
|
||||
input_ids=x[0],
|
||||
bbox=x[1],
|
||||
attention_mask=x[2],
|
||||
token_type_ids=x[3],
|
||||
image=image,
|
||||
position_ids=None,
|
||||
head_mask=None,
|
||||
labels=None,
|
||||
entities=entities,
|
||||
relations=relations,
|
||||
)
|
||||
return x
|
||||
Reference in New Issue
Block a user