chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 13:24:13 +08:00
commit 1037506f2e
6050 changed files with 1731598 additions and 0 deletions
+59
View File
@@ -0,0 +1,59 @@
from __future__ import absolute_import, division, print_function, unicode_literals
import logging
from transformers import BertConfig, RobertaConfig
from s2s_ft.configuration_unilm import UnilmConfig
# from s2s_ft.modeling import LayoutlmConfig
logger = logging.getLogger(__name__)
class BertForSeq2SeqConfig(BertConfig):
def __init__(self, label_smoothing=0.1, source_type_id=0, target_type_id=1, **kwargs):
super(BertForSeq2SeqConfig, self).__init__(**kwargs)
self.label_smoothing = label_smoothing
self.source_type_id = source_type_id
self.target_type_id = target_type_id
@classmethod
def from_exist_config(cls, config, label_smoothing=0.1, max_position_embeddings=None, max_source_length=None,
base_model_type='bert', layoutlm_only_layout_flag=False):
required_keys = [
"vocab_size", "hidden_size", "num_hidden_layers", "num_attention_heads",
"hidden_act", "intermediate_size", "hidden_dropout_prob", "attention_probs_dropout_prob",
"max_position_embeddings", "type_vocab_size", "initializer_range", "layer_norm_eps"]
kwargs = {}
for key in required_keys:
assert hasattr(config, key)
kwargs[key] = getattr(config, key)
kwargs["vocab_size_or_config_json_file"] = kwargs["vocab_size"]
if isinstance(config, RobertaConfig):
kwargs["type_vocab_size"] = 0
kwargs["max_position_embeddings"] = kwargs["max_position_embeddings"] - 2
additional_keys = [
"source_type_id", "target_type_id"
]
for key in additional_keys:
if hasattr(config, key):
kwargs[key] = getattr(config, key)
# if isinstance(config, LayoutlmConfig):
if hasattr(config, 'max_2d_position_embeddings'):
layoutlm_special_keys = ['max_2d_position_embeddings',]
for key in layoutlm_special_keys:
kwargs[key] = getattr(config, key)
kwargs['base_model_type'] = base_model_type
kwargs['layoutlm_only_layout'] = layoutlm_only_layout_flag
if max_position_embeddings is not None and max_position_embeddings > config.max_position_embeddings:
kwargs["max_position_embeddings"] = max_position_embeddings
logger.info(" ** Change max position embeddings to %d ** " % max_position_embeddings)
if max_source_length is not None:
kwargs['max_source_length'] = max_source_length
return cls(label_smoothing=label_smoothing, **kwargs)
+110
View File
@@ -0,0 +1,110 @@
# coding=utf-8
# The MIT License (MIT)
# Copyright (c) Microsoft Corporation
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
""" MiniLM model configuration """
from __future__ import absolute_import, division, print_function, unicode_literals
import json
import logging
import sys
from io import open
from transformers.configuration_utils import PretrainedConfig
logger = logging.getLogger(__name__)
MINILM_PRETRAINED_CONFIG_ARCHIVE_MAP = {
'minilm-l12-h384-uncased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/minilm-l12-h384-uncased-config.json?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
}
class MinilmConfig(PretrainedConfig):
r"""
:class:`~transformers.MinilmConfig` is the configuration class to store the configuration of a
`MinilmModel`.
Arguments:
vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `MiniLMModel`.
hidden_size: Size of the encoder layers and the pooler layer.
num_hidden_layers: Number of hidden layers in the Transformer encoder.
num_attention_heads: Number of attention heads for each attention layer in
the Transformer encoder.
intermediate_size: The size of the "intermediate" (i.e., feed-forward)
layer in the Transformer encoder.
hidden_act: The non-linear activation function (function or string) in the
encoder and pooler. If string, "gelu", "relu", "swish" and "gelu_new" are supported.
hidden_dropout_prob: The dropout probabilitiy for all fully connected
layers in the embeddings, encoder, and pooler.
attention_probs_dropout_prob: The dropout ratio for the attention
probabilities.
max_position_embeddings: The maximum sequence length that this model might
ever be used with. Typically set this to something large just in case
(e.g., 512 or 1024 or 2048).
type_vocab_size: The vocabulary size of the `token_type_ids` passed into
`MiniLMModel`.
initializer_range: The sttdev of the truncated_normal_initializer for
initializing all weight matrices.
layer_norm_eps: The epsilon used by LayerNorm.
"""
pretrained_config_archive_map = MINILM_PRETRAINED_CONFIG_ARCHIVE_MAP
def __init__(self,
vocab_size=28996,
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=6,
initializer_range=0.02,
layer_norm_eps=1e-12,
source_type_id=0,
target_type_id=1,
**kwargs):
super(MinilmConfig, self).__init__(**kwargs)
if isinstance(vocab_size, str) or (sys.version_info[0] == 2
and isinstance(vocab_size, unicode)):
with open(vocab_size, "r", encoding='utf-8') as reader:
json_config = json.loads(reader.read())
for key, value in json_config.items():
self.__dict__[key] = value
elif isinstance(vocab_size, int):
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
self.hidden_act = hidden_act
self.intermediate_size = intermediate_size
self.hidden_dropout_prob = hidden_dropout_prob
self.attention_probs_dropout_prob = attention_probs_dropout_prob
self.max_position_embeddings = max_position_embeddings
self.type_vocab_size = type_vocab_size
self.initializer_range = initializer_range
self.layer_norm_eps = layer_norm_eps
self.source_type_id = source_type_id
self.target_type_id = target_type_id
else:
raise ValueError("First argument must be either a vocabulary size (int)"
" or the path to a pretrained model config file (str)")
+114
View File
@@ -0,0 +1,114 @@
# coding=utf-8
# The MIT License (MIT)
# Copyright (c) Microsoft Corporation
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
""" UniLM model configuration """
from __future__ import absolute_import, division, print_function, unicode_literals
import json
import logging
import sys
from io import open
from transformers.configuration_utils import PretrainedConfig
logger = logging.getLogger(__name__)
UNILM_PRETRAINED_CONFIG_ARCHIVE_MAP = {
'unilm-large-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm-large-cased-config.json?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm-base-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm-base-cased-config.json?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm1-large-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm1-large-cased-config.json?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm1-base-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm1-base-cased-config.json?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm1.2-base-uncased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm1.2-base-uncased-config.json?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
}
class UnilmConfig(PretrainedConfig):
r"""
:class:`~transformers.UnilmConfig` is the configuration class to store the configuration of a
`UnilmModel`.
Arguments:
vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `UnilmModel`.
hidden_size: Size of the encoder layers and the pooler layer.
num_hidden_layers: Number of hidden layers in the Transformer encoder.
num_attention_heads: Number of attention heads for each attention layer in
the Transformer encoder.
intermediate_size: The size of the "intermediate" (i.e., feed-forward)
layer in the Transformer encoder.
hidden_act: The non-linear activation function (function or string) in the
encoder and pooler. If string, "gelu", "relu", "swish" and "gelu_new" are supported.
hidden_dropout_prob: The dropout probabilitiy for all fully connected
layers in the embeddings, encoder, and pooler.
attention_probs_dropout_prob: The dropout ratio for the attention
probabilities.
max_position_embeddings: The maximum sequence length that this model might
ever be used with. Typically set this to something large just in case
(e.g., 512 or 1024 or 2048).
type_vocab_size: The vocabulary size of the `token_type_ids` passed into
`UnilmModel`.
initializer_range: The sttdev of the truncated_normal_initializer for
initializing all weight matrices.
layer_norm_eps: The epsilon used by LayerNorm.
"""
pretrained_config_archive_map = UNILM_PRETRAINED_CONFIG_ARCHIVE_MAP
def __init__(self,
vocab_size=28996,
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=6,
initializer_range=0.02,
layer_norm_eps=1e-12,
source_type_id=0,
target_type_id=1,
**kwargs):
super(UnilmConfig, self).__init__(**kwargs)
if isinstance(vocab_size, str) or (sys.version_info[0] == 2
and isinstance(vocab_size, unicode)):
with open(vocab_size, "r", encoding='utf-8') as reader:
json_config = json.loads(reader.read())
for key, value in json_config.items():
self.__dict__[key] = value
elif isinstance(vocab_size, int):
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
self.hidden_act = hidden_act
self.intermediate_size = intermediate_size
self.hidden_dropout_prob = hidden_dropout_prob
self.attention_probs_dropout_prob = attention_probs_dropout_prob
self.max_position_embeddings = max_position_embeddings
self.type_vocab_size = type_vocab_size
self.initializer_range = initializer_range
self.layer_norm_eps = layer_norm_eps
self.source_type_id = source_type_id
self.target_type_id = target_type_id
else:
raise ValueError("First argument must be either a vocabulary size (int)"
" or the path to a pretrained model config file (str)")
+127
View File
@@ -0,0 +1,127 @@
import torch
import logging
from transformers.modeling_utils import cached_path, WEIGHTS_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME
logger = logging.getLogger(__name__)
def get_checkpoint_from_transformer_cache(
archive_file, pretrained_model_name_or_path, pretrained_model_archive_map,
cache_dir, force_download, proxies, resume_download,
):
try:
resolved_archive_file = cached_path(archive_file, cache_dir=cache_dir, force_download=force_download,
proxies=proxies, resume_download=resume_download)
except EnvironmentError:
if pretrained_model_name_or_path in pretrained_model_archive_map:
msg = "Couldn't reach server at '{}' to download pretrained weights.".format(
archive_file)
else:
msg = "Model name '{}' was not found in model name list ({}). " \
"We assumed '{}' was a path or url to model weight files named one of {} but " \
"couldn't find any such file at this path or url.".format(
pretrained_model_name_or_path,
', '.join(pretrained_model_archive_map.keys()),
archive_file,
[WEIGHTS_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME])
raise EnvironmentError(msg)
if resolved_archive_file == archive_file:
logger.info("loading weights file {}".format(archive_file))
else:
logger.info("loading weights file {} from cache at {}".format(
archive_file, resolved_archive_file))
return torch.load(resolved_archive_file, map_location='cpu')
def hf_roberta_to_hf_bert(state_dict):
logger.info(" * Convert Huggingface RoBERTa format to Huggingface BERT format * ")
new_state_dict = {}
for key in state_dict:
value = state_dict[key]
if key == 'roberta.embeddings.position_embeddings.weight':
value = value[2:]
if key == 'roberta.embeddings.token_type_embeddings.weight':
continue
if key.startswith('roberta'):
key = 'bert.' + key[8:]
elif key.startswith('lm_head'):
if 'layer_norm' in key or 'dense' in key:
key = 'cls.predictions.transform.' + key[8:]
else:
key = 'cls.predictions.' + key[8:]
key = key.replace('layer_norm', 'LayerNorm')
new_state_dict[key] = value
return new_state_dict
def hf_distilbert_to_hf_bert(state_dict):
logger.info(" * Convert Huggingface DistilBERT format to Huggingface BERT format * ")
new_state_dict = {}
for key in state_dict:
value = state_dict[key]
if key == 'roberta.embeddings.position_embeddings.weight':
value = value[2:]
if key == 'roberta.embeddings.token_type_embeddings.weight':
continue
if key.startswith('roberta'):
key = 'bert.' + key[8:]
elif key.startswith('lm_head'):
if 'layer_norm' in key or 'dense' in key:
key = 'cls.predictions.transform.' + key[8:]
else:
key = 'cls.predictions.' + key[8:]
key = key.replace('layer_norm', 'LayerNorm')
new_state_dict[key] = value
return new_state_dict
def hf_bert_to_hf_bert(state_dict):
# NOTE: all cls states are used for prediction,
# we predict the index so omit all pretrained states for prediction.
new_state_dict = {}
for key in state_dict:
value = state_dict[key]
if key.startswith('cls'):
# NOTE: all cls states are used for prediction,
# we predict the index so omit all pretrained states for prediction.
continue
new_state_dict[key] = value
return new_state_dict
def hf_layoutlm_to_hf_bert(state_dict):
logger.info(" * Convert Huggingface LayoutLM format to Huggingface BERT format * ")
new_state_dict = {}
for key in state_dict:
value = state_dict[key]
if key.startswith('layoutlm'):
key = 'bert.' + key[9:]
elif key.startswith('cls'):
# NOTE: all cls states are used for prediction,
# we predict the index so omit all pretrained states for prediction.
continue
new_state_dict[key] = value
return new_state_dict
state_dict_convert = {
'bert': hf_bert_to_hf_bert,
'unilm': hf_bert_to_hf_bert,
'minilm': hf_bert_to_hf_bert,
'layoutlm': hf_layoutlm_to_hf_bert,
'roberta': hf_roberta_to_hf_bert,
'xlm-roberta': hf_roberta_to_hf_bert,
'distilbert': hf_distilbert_to_hf_bert,
}
+808
View File
@@ -0,0 +1,808 @@
from __future__ import absolute_import, division, print_function, unicode_literals
import logging
import math
import os
import torch
from torch import nn
from torch.nn.modules.loss import _Loss
import torch.nn.functional as F
from transformers import BertConfig
from transformers.modeling_bert import \
BertPreTrainedModel, BertSelfOutput, BertIntermediate, BertOutput, BertPredictionHeadTransform
from transformers.modeling_roberta import ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
from transformers.modeling_bert import BERT_PRETRAINED_MODEL_ARCHIVE_MAP
from transformers.modeling_xlm_roberta import XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
from s2s_ft.config import BertForSeq2SeqConfig
from s2s_ft.convert_state_dict import get_checkpoint_from_transformer_cache, state_dict_convert
logger = logging.getLogger(__name__)
BertLayerNorm = torch.nn.LayerNorm
UNILM_PRETRAINED_MODEL_ARCHIVE_MAP = {
'unilm-base-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm1-base-cased.bin?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm-large-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm1-large-cased.bin?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm1-base-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm1-base-cased.bin?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm1-large-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm1-large-cased.bin?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm1.2-base-uncased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm1.2-base-uncased.bin?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D"
}
MINILM_PRETRAINED_MODEL_ARCHIVE_MAP = {
'minilm-l12-h384-uncased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/minilm-l12-h384-uncased.bin?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
}
LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_MAP = {
'layoutlm-base-uncased': 'https://huggingface.co/microsoft/layoutlm-base-uncased/resolve/main/pytorch_model.bin',
'layoutlm-large-uncased': 'https://huggingface.co/microsoft/layoutlm-large-uncased/resolve/main/pytorch_model.bin'
}
LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP = {
'layoutlm-base-uncased': 'https://huggingface.co/microsoft/layoutlm-base-uncased/resolve/main/config.json',
'layoutlm-large-uncased': 'https://huggingface.co/microsoft/layoutlm-large-uncased/resolve/main/config.json'
}
class LayoutlmConfig(BertConfig):
pretrained_config_archive_map = LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP
model_type = "bert"
def __init__(self, max_2d_position_embeddings=1024, **kwargs):
super().__init__(**kwargs)
self.max_2d_position_embeddings = max_2d_position_embeddings
class BertPreTrainedForSeq2SeqModel(BertPreTrainedModel):
""" An abstract class to handle weights initialization and
a simple interface for dowloading and loading pretrained models.
"""
config_class = BertForSeq2SeqConfig
supported_convert_pretrained_model_archive_map = {
"bert": BERT_PRETRAINED_MODEL_ARCHIVE_MAP,
"roberta": ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP,
"xlm-roberta": XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP,
"unilm": UNILM_PRETRAINED_MODEL_ARCHIVE_MAP,
"minilm": MINILM_PRETRAINED_MODEL_ARCHIVE_MAP,
"layoutlm": LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_MAP,
}
base_model_prefix = "bert_for_seq2seq"
pretrained_model_archive_map = {
**ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP,
**XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP,
**BERT_PRETRAINED_MODEL_ARCHIVE_MAP,
**UNILM_PRETRAINED_MODEL_ARCHIVE_MAP,
**MINILM_PRETRAINED_MODEL_ARCHIVE_MAP,
**LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_MAP,
}
def _init_weights(self, module):
""" Initialize the weights """
if isinstance(module, (nn.Linear, nn.Embedding)):
# 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)
elif isinstance(module, BertLayerNorm):
module.bias.data.zero_()
module.weight.data.fill_(1.0)
if isinstance(module, nn.Linear) and module.bias is not None:
module.bias.data.zero_()
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, reuse_position_embedding=None,
*model_args, **kwargs):
model_type = kwargs.pop('model_type', None)
if model_type is not None and "state_dict" not in kwargs:
if model_type in cls.supported_convert_pretrained_model_archive_map:
pretrained_model_archive_map = cls.supported_convert_pretrained_model_archive_map[model_type]
if pretrained_model_name_or_path in pretrained_model_archive_map:
state_dict = get_checkpoint_from_transformer_cache(
archive_file=pretrained_model_archive_map[pretrained_model_name_or_path],
pretrained_model_name_or_path=pretrained_model_name_or_path,
pretrained_model_archive_map=pretrained_model_archive_map,
cache_dir=kwargs.get("cache_dir", None), force_download=kwargs.get("force_download", None),
proxies=kwargs.get("proxies", None), resume_download=kwargs.get("resume_download", None),
)
state_dict = state_dict_convert[model_type](state_dict)
kwargs["state_dict"] = state_dict
elif os.path.isfile(pretrained_model_name_or_path):
kwargs["state_dict"] = torch.load(pretrained_model_name_or_path, map_location='cpu')
if kwargs["state_dict"] is None:
logger.info("s2s-ft does't support the model !")
raise NotImplementedError()
config = kwargs["config"]
state_dict = kwargs["state_dict"]
# initialize new position embeddings (From Microsoft/UniLM)
_k = 'bert.embeddings.position_embeddings.weight'
if _k in state_dict:
if config.max_position_embeddings > state_dict[_k].shape[0]:
logger.info("Resize > position embeddings !")
old_vocab_size = state_dict[_k].shape[0]
new_position_embedding = state_dict[_k].data.new_tensor(torch.ones(
size=(config.max_position_embeddings, state_dict[_k].shape[1])), dtype=torch.float)
new_position_embedding = nn.Parameter(data=new_position_embedding, requires_grad=True)
new_position_embedding.data.normal_(mean=0.0, std=config.initializer_range)
max_range = config.max_position_embeddings if reuse_position_embedding else old_vocab_size
shift = 0
while shift < max_range:
delta = min(old_vocab_size, max_range - shift)
new_position_embedding.data[shift: shift + delta, :] = state_dict[_k][:delta, :]
logger.info(" CP [%d ~ %d] into [%d ~ %d] " % (0, delta, shift, shift + delta))
shift += delta
state_dict[_k] = new_position_embedding.data
del new_position_embedding
elif config.max_position_embeddings < state_dict[_k].shape[0]:
logger.info("Resize < position embeddings !")
old_vocab_size = state_dict[_k].shape[0]
new_position_embedding = state_dict[_k].data.new_tensor(torch.ones(
size=(config.max_position_embeddings, state_dict[_k].shape[1])), dtype=torch.float)
new_position_embedding = nn.Parameter(data=new_position_embedding, requires_grad=True)
new_position_embedding.data.normal_(mean=0.0, std=config.initializer_range)
new_position_embedding.data.copy_(state_dict[_k][:config.max_position_embeddings, :])
state_dict[_k] = new_position_embedding.data
del new_position_embedding
return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
class BertEmbeddings(nn.Module):
"""Construct the embeddings from word, position and token_type embeddings.
"""
def __init__(self, config):
super(BertEmbeddings, self).__init__()
self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)
self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
if config.type_vocab_size > 0:
self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
else:
self.token_type_embeddings = None
# self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
# any TensorFlow checkpoint file
self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
self.dropout = nn.Dropout(config.hidden_dropout_prob)
def forward(self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None):
if input_ids is not None:
input_shape = input_ids.size()
else:
input_shape = inputs_embeds.size()[:-1]
seq_length = input_shape[1]
device = input_ids.device if input_ids is not None else inputs_embeds.device
if position_ids is None:
position_ids = torch.arange(seq_length, dtype=torch.long, device=device)
position_ids = position_ids.unsqueeze(0).expand(input_shape)
if token_type_ids is None:
token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
if inputs_embeds is None:
inputs_embeds = self.word_embeddings(input_ids)
position_embeddings = self.position_embeddings(position_ids)
embeddings = inputs_embeds + position_embeddings
if self.token_type_embeddings:
embeddings = embeddings + self.token_type_embeddings(token_type_ids)
embeddings = self.LayerNorm(embeddings)
embeddings = self.dropout(embeddings)
return embeddings
class LayoutlmEmbeddings(nn.Module):
def __init__(self, config):
super(LayoutlmEmbeddings, self).__init__()
self.only_layout_flag = config.layoutlm_only_layout
if not config.layoutlm_only_layout:
self.word_embeddings = nn.Embedding(
config.vocab_size, config.hidden_size, padding_idx=0
)
else:
self.word_embeddings = None
self.position_embeddings = nn.Embedding(
config.max_position_embeddings, config.hidden_size
)
self.x_position_embeddings = nn.Embedding(
config.max_2d_position_embeddings, config.hidden_size
)
self.y_position_embeddings = nn.Embedding(
config.max_2d_position_embeddings, config.hidden_size
)
self.h_position_embeddings = nn.Embedding(
config.max_2d_position_embeddings, config.hidden_size
)
self.w_position_embeddings = nn.Embedding(
config.max_2d_position_embeddings, config.hidden_size
)
if config.type_vocab_size > 0:
self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
else:
self.token_type_embeddings = None
# self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
# any TensorFlow checkpoint file
self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
self.dropout = nn.Dropout(config.hidden_dropout_prob)
def forward(
self,
input_ids,
bbox,
token_type_ids=None,
position_ids=None,
inputs_embeds=None,
):
seq_length = input_ids.size(1)
if position_ids is None:
position_ids = torch.arange(
seq_length, dtype=torch.long, device=input_ids.device
)
position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
if token_type_ids is None:
token_type_ids = torch.zeros_like(input_ids)
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])
h_position_embeddings = self.h_position_embeddings(
bbox[:, :, 3] - bbox[:, :, 1]
)
w_position_embeddings = self.w_position_embeddings(
bbox[:, :, 2] - bbox[:, :, 0]
)
position_embeddings = self.position_embeddings(position_ids)
embeddings = (
left_position_embeddings
+ upper_position_embeddings
+ right_position_embeddings
+ lower_position_embeddings
+ h_position_embeddings
+ w_position_embeddings
+ position_embeddings
# + token_type_embeddings
)
if not self.only_layout_flag:
words_embeddings = self.word_embeddings(input_ids)
embeddings = embeddings + words_embeddings
if self.token_type_embeddings:
embeddings = embeddings + self.token_type_embeddings(token_type_ids)
embeddings = self.LayerNorm(embeddings)
embeddings = self.dropout(embeddings)
return embeddings
class BertSelfAttention(nn.Module):
def __init__(self, config):
super(BertSelfAttention, self).__init__()
if config.hidden_size % config.num_attention_heads != 0:
raise ValueError(
"The hidden size (%d) is not a multiple of the number of attention "
"heads (%d)" % (config.hidden_size, config.num_attention_heads))
self.output_attentions = config.output_attentions
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)
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 multi_head_attention(self, query, key, value, attention_mask):
query_layer = self.transpose_for_scores(query)
key_layer = self.transpose_for_scores(key)
value_layer = self.transpose_for_scores(value)
# 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))
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 BertModel forward() function)
attention_scores = attention_scores + attention_mask
# Normalize the attention scores to probabilities.
attention_probs = nn.Softmax(dim=-1)(attention_scores)
# 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)
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)
return (context_layer, attention_probs) if self.output_attentions else (context_layer,)
def forward(self, hidden_states, attention_mask=None, encoder_hidden_states=None, split_lengths=None):
mixed_query_layer = self.query(hidden_states)
if split_lengths:
assert not self.output_attentions
# 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.
if encoder_hidden_states is not None:
mixed_key_layer = self.key(encoder_hidden_states)
mixed_value_layer = self.value(encoder_hidden_states)
else:
mixed_key_layer = self.key(hidden_states)
mixed_value_layer = self.value(hidden_states)
if split_lengths:
query_parts = torch.split(mixed_query_layer, split_lengths, dim=1)
key_parts = torch.split(mixed_key_layer, split_lengths, dim=1)
value_parts = torch.split(mixed_value_layer, split_lengths, dim=1)
key = None
value = None
outputs = []
sum_length = 0
for (query, _key, _value, part_length) in zip(query_parts, key_parts, value_parts, split_lengths):
key = _key if key is None else torch.cat((key, _key), dim=1)
value = _value if value is None else torch.cat((value, _value), dim=1)
sum_length += part_length
outputs.append(self.multi_head_attention(
query, key, value, attention_mask[:, :, sum_length - part_length: sum_length, :sum_length]
)[0])
outputs = (torch.cat(outputs, dim=1), )
else:
outputs = self.multi_head_attention(
mixed_query_layer, mixed_key_layer, mixed_value_layer, attention_mask)
return outputs
class BertAttention(nn.Module):
def __init__(self, config):
super(BertAttention, self).__init__()
self.self = BertSelfAttention(config)
self.output = BertSelfOutput(config)
def forward(self, hidden_states, attention_mask=None, encoder_hidden_states=None, split_lengths=None):
self_outputs = self.self(
hidden_states, attention_mask=attention_mask,
encoder_hidden_states=encoder_hidden_states, split_lengths=split_lengths)
attention_output = self.output(self_outputs[0], hidden_states)
outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
return outputs
class BertLayer(nn.Module):
def __init__(self, config):
super(BertLayer, self).__init__()
self.attention = BertAttention(config)
self.intermediate = BertIntermediate(config)
self.output = BertOutput(config)
def forward(self, hidden_states, attention_mask=None, split_lengths=None):
self_attention_outputs = self.attention(
hidden_states, attention_mask, split_lengths=split_lengths)
attention_output = self_attention_outputs[0]
intermediate_output = self.intermediate(attention_output)
layer_output = self.output(intermediate_output, attention_output)
outputs = (layer_output,) + self_attention_outputs[1:]
return outputs
class BertEncoder(nn.Module):
def __init__(self, config):
super(BertEncoder, self).__init__()
self.output_attentions = config.output_attentions
self.output_hidden_states = config.output_hidden_states
self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
def forward(self, hidden_states, attention_mask=None, split_lengths=None):
all_hidden_states = ()
all_attentions = ()
for i, layer_module in enumerate(self.layer):
if self.output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
layer_outputs = layer_module(hidden_states, attention_mask, split_lengths=split_lengths)
hidden_states = layer_outputs[0]
if self.output_attentions:
all_attentions = all_attentions + (layer_outputs[1],)
# Add last layer
if self.output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
outputs = (hidden_states,)
if self.output_hidden_states:
outputs = outputs + (all_hidden_states,)
if self.output_attentions:
outputs = outputs + (all_attentions,)
return outputs # last-layer hidden state, (all hidden states), (all attentions)
class BertModel(BertPreTrainedForSeq2SeqModel):
r"""
Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
**last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``
Sequence of hidden-states at the output of the last layer of the model.
**pooler_output**: ``torch.FloatTensor`` of shape ``(batch_size, hidden_size)``
Last layer hidden-state of the first token of the sequence (classification token)
further processed by a Linear layer and a Tanh activation function. The Linear
layer weights are trained from the next sentence prediction (classification)
objective during Bert pretraining. This output is usually *not* a good summary
of the semantic content of the input, you're often better with averaging or pooling
the sequence of hidden-states for the whole input sequence.
**hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
of shape ``(batch_size, sequence_length, hidden_size)``:
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
**attentions**: (`optional`, returned when ``config.output_attentions=True``)
list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
Examples::
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0) # Batch size 1
outputs = model(input_ids)
last_hidden_states = outputs[0] # The last hidden-state is the first element of the output tuple
"""
def __init__(self, config):
super(BertModel, self).__init__(config)
self.config = config
self.embeddings = BertEmbeddings(config)
self.encoder = BertEncoder(config)
def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
position_ids=None, inputs_embeds=None, split_lengths=None, return_emb=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")
device = input_ids.device if input_ids is not None else inputs_embeds.device
if attention_mask is None:
attention_mask = torch.ones(input_shape, 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.
if attention_mask.dim() == 3:
extended_attention_mask = attention_mask[:, None, :, :]
# Provided a padding mask of dimensions [batch_size, seq_length]
# - if the model is a decoder, apply a causal mask in addition to the padding mask
# - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
if attention_mask.dim() == 2:
extended_attention_mask = attention_mask[:, None, None, :]
# Since attention_mask is 1.0 for positions we want to attend and 0.0 for
# masked positions, this operation will create a tensor which is 0.0 for
# positions we want to attend and -10000.0 for masked positions.
# Since we are adding it to the raw scores before the softmax, this is
# effectively the same as removing these entirely.
extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
embedding_output = self.embeddings(
input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds)
encoder_outputs = self.encoder(
embedding_output, attention_mask=extended_attention_mask, split_lengths=split_lengths)
sequence_output = encoder_outputs[0]
outputs = (sequence_output, ) + encoder_outputs[1:] # add hidden_states and attentions if they are here
if return_emb:
outputs += (embedding_output,)
return outputs # sequence_output, pooled_output, (hidden_states), (attentions)
class LayoutlmModel(BertPreTrainedForSeq2SeqModel):
r"""
Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
**last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``
Sequence of hidden-states at the output of the last layer of the model.
**pooler_output**: ``torch.FloatTensor`` of shape ``(batch_size, hidden_size)``
Last layer hidden-state of the first token of the sequence (classification token)
further processed by a Linear layer and a Tanh activation function. The Linear
layer weights are trained from the next sentence prediction (classification)
objective during Bert pretraining. This output is usually *not* a good summary
of the semantic content of the input, you're often better with averaging or pooling
the sequence of hidden-states for the whole input sequence.
**hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
of shape ``(batch_size, sequence_length, hidden_size)``:
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
**attentions**: (`optional`, returned when ``config.output_attentions=True``)
list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
Examples::
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0) # Batch size 1
outputs = model(input_ids)
last_hidden_states = outputs[0] # The last hidden-state is the first element of the output tuple
"""
def __init__(self, config):
super(LayoutlmModel, self).__init__(config)
self.config = config
self.embeddings = LayoutlmEmbeddings(config)
self.encoder = BertEncoder(config)
def forward(self,
input_ids=None,
bbox=None,
attention_mask=None,
token_type_ids=None,
position_ids=None,
inputs_embeds=None,
split_lengths=None,
return_emb=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")
device = input_ids.device if input_ids is not None else inputs_embeds.device
if attention_mask is None:
attention_mask = torch.ones(input_shape, 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.
if attention_mask.dim() == 3:
extended_attention_mask = attention_mask[:, None, :, :]
# Provided a padding mask of dimensions [batch_size, seq_length]
# - if the model is a decoder, apply a causal mask in addition to the padding mask
# - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
if attention_mask.dim() == 2:
extended_attention_mask = attention_mask[:, None, None, :]
# Since attention_mask is 1.0 for positions we want to attend and 0.0 for
# masked positions, this operation will create a tensor which is 0.0 for
# positions we want to attend and -10000.0 for masked positions.
# Since we are adding it to the raw scores before the softmax, this is
# effectively the same as removing these entirely.
extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
# embedding_output = self.embeddings(
# input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds)
embedding_output = self.embeddings(
input_ids, bbox, position_ids=position_ids, token_type_ids=token_type_ids
)
encoder_outputs = self.encoder(
embedding_output, attention_mask=extended_attention_mask, split_lengths=split_lengths)
sequence_output = encoder_outputs[0]
outputs = (sequence_output, ) + encoder_outputs[1:] # add hidden_states and attentions if they are here
if return_emb:
outputs += (embedding_output,)
return outputs # sequence_output, pooled_output, (hidden_states), (attentions)
class LabelSmoothingLoss(_Loss):
"""
With label smoothing,
KL-divergence between q_{smoothed ground truth prob.}(w)
and p_{prob. computed by model}(w) is minimized.
"""
def __init__(self, label_smoothing=0, tgt_size=0, ignore_index=0, size_average=None, reduce=None, reduction='mean'):
assert 0.0 < label_smoothing <= 1.0
self.ignore_index = ignore_index
super(LabelSmoothingLoss, self).__init__(
size_average=size_average, reduce=reduce, reduction=reduction)
assert label_smoothing > 0
assert tgt_size > 0
smoothing_value = label_smoothing / (tgt_size - 2)
one_hot = torch.full((tgt_size,), smoothing_value)
one_hot[self.ignore_index] = 0
self.register_buffer('one_hot', one_hot.unsqueeze(0))
self.confidence = 1.0 - label_smoothing
self.tgt_size = tgt_size
def forward(self, output, target):
"""
output (FloatTensor): batch_size * num_pos * n_classes
target (LongTensor): batch_size * num_pos
"""
assert self.tgt_size == output.size(2)
batch_size, num_pos = target.size(0), target.size(1)
output = output.view(-1, self.tgt_size)
target = target.view(-1)
model_prob = self.one_hot.float().repeat(target.size(0), 1)
model_prob.scatter_(1, target.unsqueeze(1), self.confidence)
model_prob.masked_fill_((target == self.ignore_index).unsqueeze(1), 0)
return F.kl_div(output, model_prob, reduction='none').view(batch_size, num_pos, -1).sum(2)
class LayoutlmSPLMPredictionHead(nn.Module):
def __init__(self, config, src_len):
super(LayoutlmSPLMPredictionHead, self).__init__()
self.transform = BertPredictionHeadTransform(config)
self.bias = nn.Parameter(torch.zeros(src_len))
def forward(self, hidden_states, src_emb):
hidden_states = self.transform(hidden_states)
hidden_states = torch.einsum('btf,bsf->bts', hidden_states, src_emb) + self.bias
# hidden_states = F.linear(hidden_states, weight=src_emb, bias=self.bias)
return hidden_states
class LayoutlmSPOnlyMLMHead(nn.Module):
def __init__(self, config, src_len):
super(LayoutlmSPOnlyMLMHead, self).__init__()
self.predictions = LayoutlmSPLMPredictionHead(config, src_len=src_len)
def forward(self, sequence_output, src_emb):
prediction_scores = self.predictions(sequence_output, src_emb=src_emb)
return prediction_scores
class LayoutlmForSequenceToSequence(BertPreTrainedForSeq2SeqModel):
def __init__(self, config):
super(LayoutlmForSequenceToSequence, self).__init__(config)
if config.base_model_type == 'layoutlm':
self.bert = LayoutlmModel(config)
else:
self.bert = BertModel(config)
self.cls = LayoutlmSPOnlyMLMHead(config, src_len=config.max_source_length)
self.init_weights()
self.log_softmax = nn.LogSoftmax()
# setattr(config, 'label_smoothing', 0.1)
self.source_type_id = config.source_type_id
self.target_type_id = config.target_type_id
if config.label_smoothing > 0:
self.crit_mask_lm_smoothed = LabelSmoothingLoss(
config.label_smoothing, config.max_source_length, ignore_index=0, reduction='none')
self.crit_mask_lm = None
else:
self.crit_mask_lm_smoothed = None
self.crit_mask_lm = nn.CrossEntropyLoss(reduction='none', ignore_index=0)
@staticmethod
def create_mask_and_position_ids(num_tokens, max_len, offset=None):
base_position_matrix = torch.arange(
0, max_len, dtype=num_tokens.dtype, device=num_tokens.device).view(1, -1)
mask = (base_position_matrix < num_tokens.view(-1, 1)).type_as(num_tokens)
if offset is not None:
base_position_matrix = base_position_matrix + offset.view(-1, 1)
position_ids = base_position_matrix * mask
return mask, position_ids
@staticmethod
def create_attention_mask(source_mask, target_mask, source_position_ids, target_span_ids):
weight = torch.cat((torch.zeros_like(source_position_ids), target_span_ids, -target_span_ids), dim=1)
from_weight = weight.unsqueeze(-1)
to_weight = weight.unsqueeze(1)
true_tokens = (0 <= to_weight) & (torch.cat((source_mask, target_mask, target_mask), dim=1) == 1).unsqueeze(1)
true_tokens_mask = (from_weight >= 0) & true_tokens & (to_weight <= from_weight)
pseudo_tokens_mask = (from_weight < 0) & true_tokens & (-to_weight > from_weight)
pseudo_tokens_mask = pseudo_tokens_mask | ((from_weight < 0) & (to_weight == from_weight))
return (true_tokens_mask | pseudo_tokens_mask).type_as(source_mask)
def forward(self, source_idxys, target_idxys, target_index, pseudo_idxys, num_source_tokens, num_target_tokens,
target_span_ids=None):
source_len = source_idxys.size(1)
target_len = target_idxys.size(1)
pseudo_len = pseudo_idxys.size(1)
assert target_len == pseudo_len
assert source_len > 0 and target_len > 0
split_lengths = (source_len, target_len, pseudo_len)
if self.config.base_model_type == 'layoutlm':
source_xys = source_idxys[:, :, 1:]
target_xys = target_idxys[:, :, 1:]
pseudo_xys = pseudo_idxys[:, :, 1:]
input_xys = torch.cat((source_xys, target_xys, pseudo_xys), dim=1)
source_ids = source_idxys[:, :, 0]
target_ids = target_idxys[:, :, 0]
pseudo_ids = pseudo_idxys[:, :, 0]
else:
source_ids = source_idxys
target_ids = target_idxys
pseudo_ids = pseudo_idxys
input_xys = None
input_ids = torch.cat((source_ids, target_ids, pseudo_ids), dim=1)
token_type_ids = torch.cat(
(torch.ones_like(source_ids) * self.source_type_id,
torch.ones_like(target_ids) * self.target_type_id,
torch.ones_like(pseudo_ids) * self.target_type_id), dim=1)
source_mask, source_position_ids = \
self.create_mask_and_position_ids(num_source_tokens, source_len)
target_mask, target_position_ids = \
self.create_mask_and_position_ids(num_target_tokens, target_len, offset=num_source_tokens)
position_ids = torch.cat((source_position_ids, target_position_ids, target_position_ids), dim=1)
if target_span_ids is None:
target_span_ids = target_position_ids
attention_mask = self.create_attention_mask(source_mask, target_mask, source_position_ids, target_span_ids)
if self.config.base_model_type == 'layoutlm':
outputs = self.bert(
input_ids, input_xys, attention_mask=attention_mask, token_type_ids=token_type_ids,
position_ids=position_ids, split_lengths=split_lengths, return_emb=True)
else:
outputs = self.bert(
input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids,
position_ids=position_ids, split_lengths=split_lengths, return_emb=True)
sequence_output = outputs[0]
pseudo_sequence_output = sequence_output[:, source_len + target_len:, ]
sequence_embedding = outputs[-1]
source_embedding = sequence_embedding[:, :source_len, :]
def loss_mask_and_normalize(loss, mask):
mask = mask.type_as(loss)
loss = loss * mask
denominator = torch.sum(mask) + 1e-5
return (loss / denominator).sum()
# TODO: do we need to mask the impossible pos with the real input length
prediction_scores_masked = self.cls(pseudo_sequence_output, source_embedding)
if self.crit_mask_lm_smoothed:
masked_lm_loss = self.crit_mask_lm_smoothed(
F.log_softmax(prediction_scores_masked.float(), dim=-1), target_index)
else:
masked_lm_loss = self.crit_mask_lm(
prediction_scores_masked.transpose(1, 2).float(), target_index)
pseudo_lm_loss = loss_mask_and_normalize(
masked_lm_loss.float(), target_mask)
return pseudo_lm_loss
File diff suppressed because it is too large Load Diff
+165
View File
@@ -0,0 +1,165 @@
import numpy as np
from random import randint, shuffle, choice
from random import random as rand
import math
import logging
import torch
import torch.utils.data
logger = logging.getLogger(__name__)
def get_random_word(vocab_words):
i = randint(0, len(vocab_words)-1)
return vocab_words[i]
def batch_list_to_batch_tensors(batch):
batch_tensors = []
for x in zip(*batch):
if x[0] is None:
batch_tensors.append(None)
elif isinstance(x[0], torch.Tensor):
batch_tensors.append(torch.stack(x))
else:
batch_tensors.append(torch.tensor(x, dtype=torch.long))
return batch_tensors
def _get_word_split_index(tokens, st, end):
split_idx = []
i = st
while i < end:
if (not tokens[i].startswith('##')) or (i == st):
split_idx.append(i)
i += 1
split_idx.append(end)
return split_idx
def _expand_whole_word(tokens, st, end):
new_st, new_end = st, end
while (new_st >= 0) and tokens[new_st].startswith('##'):
new_st -= 1
while (new_end < len(tokens)) and tokens[new_end].startswith('##'):
new_end += 1
return new_st, new_end
class Pipeline():
""" Pre-process Pipeline Class : callable """
def __init__(self):
super().__init__()
self.skipgram_prb = None
self.skipgram_size = None
self.pre_whole_word = None
self.mask_whole_word = None
self.word_subsample_prb = None
self.sp_prob = None
self.pieces_dir = None
self.vocab_words = None
self.pieces_threshold = 10
self.call_count = 0
self.offline_mode = False
self.skipgram_size_geo_list = None
self.span_same_mask = False
def __call__(self, instance):
raise NotImplementedError
class Preprocess4Seq2seqDecoder(Pipeline):
""" Pre-processing steps for pretraining transformer """
def __init__(self, vocab_words, indexer, max_len=512, max_tgt_length=128,
mode="s2s", pos_shift=False, source_type_id=0, target_type_id=1,
cls_token='[CLS]', sep_token='[SEP]', pad_token='[PAD]', layout_flag=False):
super().__init__()
self.max_len = max_len
self.vocab_words = vocab_words # vocabulary (sub)words
self.indexer = indexer # function from token to token index
self.max_len = max_len
self._tril_matrix = torch.tril(torch.ones((max_len, max_len), dtype=torch.long))
self.task_idx = 3 # relax projection layer for different tasks
assert mode in ("s2s", "l2r")
self.mode = mode
self.max_tgt_length = max_tgt_length
self.pos_shift = pos_shift
self.layout_flag = layout_flag
if layout_flag:
self.cls_token = [cls_token, 0, 0, 0, 0]
self.sep_token = [sep_token, 1000, 1000, 1000, 1000]
self.pad_token = [pad_token, 0, 0, 0, 0]
else:
self.cls_token = cls_token
self.sep_token = sep_token
self.pad_token = pad_token
self.source_type_id = source_type_id
self.target_type_id = target_type_id
self.cc = 0
def __call__(self, instance):
tokens_a, max_a_len = instance
# NOTE: must pad to the max src length
max_a_len = 511
padded_tokens_a = [self.cls_token] + tokens_a + [self.sep_token]
assert len(padded_tokens_a) <= max_a_len + 2
if max_a_len + 2 > len(padded_tokens_a):
padded_tokens_a += [self.pad_token] * \
(max_a_len + 2 - len(padded_tokens_a))
assert len(padded_tokens_a) == max_a_len + 2
max_len_in_batch = min(self.max_tgt_length + max_a_len + 2, self.max_len)
tokens = padded_tokens_a
segment_ids = [self.source_type_id] * (len(padded_tokens_a)) \
+ [self.target_type_id] * (max_len_in_batch - len(padded_tokens_a))
mask_qkv = None
position_ids = []
for i in range(len(tokens_a) + 2):
position_ids.append(i)
for i in range(len(tokens_a) + 2, max_a_len + 2):
position_ids.append(0)
for i in range(max_a_len + 2, max_len_in_batch):
position_ids.append(i - (max_a_len + 2) + len(tokens_a) + 2)
# Token Indexing
if not self.layout_flag:
input_ids = self.indexer(tokens)
else:
raw_text = [x[0] for x in tokens]
raw_text_ids = self.indexer(raw_text)
input_ids = [[i] + x[1:] for i, x in zip(raw_text_ids, tokens)]
self.cc += 1
if self.cc < 5:
if not self.layout_flag:
logger.info("Input src = %s" % " ".join(self.vocab_words[tk_id] for tk_id in input_ids))
else:
logger.info("Input src = %s" % " ".join(self.vocab_words[tk_id[0]] for tk_id in input_ids))
# Zero Padding
input_mask = torch.zeros(
max_len_in_batch, max_len_in_batch, dtype=torch.long)
if self.mode == "s2s":
input_mask[:, :len(tokens_a)+2].fill_(1)
else:
st, end = 0, len(tokens_a) + 2
input_mask[st:end, st:end].copy_(
self._tril_matrix[:end, :end])
input_mask[end:, :len(tokens_a)+2].fill_(1)
second_st, second_end = len(padded_tokens_a), max_len_in_batch
input_mask[second_st:second_end, second_st:second_end].copy_(
self._tril_matrix[:second_end-second_st, :second_end-second_st])
return input_ids, segment_ids, position_ids, input_mask, mask_qkv, self.task_idx
@@ -0,0 +1,72 @@
# coding=utf-8
# The MIT License (MIT)
# Copyright (c) Microsoft Corporation
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
"""Tokenization classes for MiniLM."""
from __future__ import absolute_import, division, print_function, unicode_literals
import collections
import logging
import os
import unicodedata
from io import open
from transformers.tokenization_bert import BertTokenizer, whitespace_tokenize
logger = logging.getLogger(__name__)
VOCAB_FILES_NAMES = {'vocab_file': 'vocab.txt'}
PRETRAINED_VOCAB_FILES_MAP = {
'vocab_file':
{
'minilm-l12-h384-uncased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/minilm-l12-h384-uncased-vocab.txt?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
}
}
PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
'minilm-l12-h384-uncased': 512,
}
class MinilmTokenizer(BertTokenizer):
r"""
Constructs a MinilmTokenizer.
:class:`~transformers.MinilmTokenizer` is identical to BertTokenizer and runs end-to-end tokenization: punctuation splitting + wordpiece
Args:
vocab_file: Path to a one-wordpiece-per-line vocabulary file
do_lower_case: Whether to lower case the input. Only has an effect when do_wordpiece_only=False
do_basic_tokenize: Whether to do basic tokenization before wordpiece.
max_len: An artificial maximum length to truncate tokenized sequences to; Effective maximum length is always the
minimum of this value (if specified) and the underlying BERT model's sequence length.
never_split: List of tokens which will never be split during tokenization. Only has an effect when
do_wordpiece_only=False
"""
vocab_files_names = VOCAB_FILES_NAMES
pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
class WhitespaceTokenizer(object):
def tokenize(self, text):
return whitespace_tokenize(text)
+80
View File
@@ -0,0 +1,80 @@
# coding=utf-8
# The MIT License (MIT)
# Copyright (c) Microsoft Corporation
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
"""Tokenization classes for UniLM."""
from __future__ import absolute_import, division, print_function, unicode_literals
import collections
import logging
import os
import unicodedata
from io import open
from transformers.tokenization_bert import BertTokenizer, whitespace_tokenize
logger = logging.getLogger(__name__)
VOCAB_FILES_NAMES = {'vocab_file': 'vocab.txt'}
PRETRAINED_VOCAB_FILES_MAP = {
'vocab_file':
{
'unilm-large-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm-large-cased-vocab.txt?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm-base-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm-base-cased-vocab.txt?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm1-large-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm1-large-cased-vocab.txt?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm1-base-cased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm1-base-cased-vocab.txt?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D",
'unilm1.2-base-uncased': "https://conversationhub.blob.core.windows.net/beit-share-public/ckpt/unilm1.2-base-uncased-vocab.txt?sv=2021-10-04&st=2023-06-08T11%3A16%3A02Z&se=2033-06-09T11%3A16%3A00Z&sr=c&sp=r&sig=N4pfCVmSeq4L4tS8QbrFVsX6f6q844eft8xSuXdxU48%3D"
}
}
PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
'unilm-large-cased': 512,
'unilm-base-cased': 512,
'unilm1-large-cased': 512,
'unilm1-base-cased': 512,
'unilm1.2-base-uncased': 512,
}
class UnilmTokenizer(BertTokenizer):
r"""
Constructs a UnilmTokenizer.
:class:`~transformers.UnilmTokenizer` is identical to BertTokenizer and runs end-to-end tokenization: punctuation splitting + wordpiece
Args:
vocab_file: Path to a one-wordpiece-per-line vocabulary file
do_lower_case: Whether to lower case the input. Only has an effect when do_wordpiece_only=False
do_basic_tokenize: Whether to do basic tokenization before wordpiece.
max_len: An artificial maximum length to truncate tokenized sequences to; Effective maximum length is always the
minimum of this value (if specified) and the underlying BERT model's sequence length.
never_split: List of tokens which will never be split during tokenization. Only has an effect when
do_wordpiece_only=False
"""
vocab_files_names = VOCAB_FILES_NAMES
pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
class WhitespaceTokenizer(object):
def tokenize(self, text):
return whitespace_tokenize(text)
+608
View File
@@ -0,0 +1,608 @@
from __future__ import absolute_import, division, print_function
import logging
import os
import json
import random
import glob
import re
import torch
import tqdm
import torch.utils.data
logger = logging.getLogger(__name__)
class Seq2seqDatasetForBert(torch.utils.data.Dataset):
def __init__(
self, features, max_source_len, max_target_len,
vocab_size, cls_id, sep_id, pad_id, mask_id,
random_prob, keep_prob, offset, num_training_instances,
span_len=1, span_prob=1.0):
self.features = features
self.max_source_len = max_source_len
self.max_target_len = max_target_len
self.offset = offset
if offset > 0:
logger.info(" **** Set offset %d in Seq2seqDatasetForBert **** ", offset)
self.cls_id = cls_id
self.sep_id = sep_id
self.pad_id = pad_id
self.random_prob = random_prob
self.keep_prob = keep_prob
self.mask_id = mask_id
self.vocab_size = vocab_size
self.num_training_instances = num_training_instances
self.span_len = span_len
self.span_prob = span_prob
def __len__(self):
return int(self.num_training_instances)
def __trunk(self, ids, max_len):
if len(ids) > max_len - 1:
ids = ids[:max_len - 1]
ids = ids + [self.sep_id]
return ids
def __pad(self, ids, max_len):
if len(ids) < max_len:
return ids + [self.pad_id] * (max_len - len(ids))
else:
assert len(ids) == max_len
return ids
def __getitem__(self, idx):
idx = (self.offset + idx) % len(self.features)
feature = self.features[idx]
source_ids = self.__trunk([self.cls_id] + feature["source_ids"], self.max_source_len)
target_ids = self.__trunk(feature["target_ids"], self.max_target_len)
pseudo_ids = []
for tk_id in target_ids:
p = random.random()
if p < self.keep_prob:
pseudo_ids.append(tk_id)
elif p < self.keep_prob + self.random_prob:
pseudo_ids.append(random.randint(0, self.vocab_size - 1))
else:
pseudo_ids.append(self.mask_id)
num_source_tokens = len(source_ids)
num_target_tokens = len(target_ids)
source_ids = self.__pad(source_ids, self.max_source_len)
target_ids = self.__pad(target_ids, self.max_target_len)
pseudo_ids = self.__pad(pseudo_ids, self.max_target_len)
if self.span_len > 1:
span_ids = []
span_id = 1
while len(span_ids) < num_target_tokens:
p = random.random()
if p < self.span_prob:
span_len = random.randint(2, self.span_len)
span_len = min(span_len, num_target_tokens - len(span_ids))
else:
span_len = 1
span_ids.extend([span_id] * span_len)
span_id += 1
span_ids = self.__pad(span_ids, self.max_target_len)
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, span_ids
else:
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens
# DONE: finish this!!! the 2D input id settings.
class Seq2seqDatasetForLayoutlm(torch.utils.data.Dataset):
def __init__(
self, features, max_source_len, max_target_len,
vocab_size, cls_id, sep_id, pad_id, mask_id,
random_prob, keep_prob, offset, num_training_instances, layout_flag=True,
span_len=1, span_prob=1.0):
self.layout_flag = layout_flag
self.features = features
self.max_source_len = max_source_len
self.max_target_len = max_target_len
self.offset = offset
if offset > 0:
logger.info(" **** Set offset %d in Seq2seqDatasetForBert **** ", offset)
self.cls_id = cls_id
self.sep_id = sep_id
self.pad_id = pad_id
self.random_prob = random_prob
self.keep_prob = keep_prob
self.mask_id = mask_id
self.vocab_size = vocab_size
self.num_training_instances = num_training_instances
self.span_len = span_len
self.span_prob = span_prob
self.index_sp_id = 0
def __len__(self):
return int(self.num_training_instances)
def __clip_index(self, ids):
replace_value = 0
for i in range(len(ids)):
if ids[i] > self.max_source_len - 1:
ids[i] = replace_value
return ids
def __trunk(self, ids, max_len, simple=False, value=None):
trunk_value = value if value is not None else self.sep_id
if len(ids) > max_len - 1:
ids = ids[:max_len - 1]
if simple:
ids = ids + [trunk_value]
else:
ids = ids + [[trunk_value, 1000, 1000, 1000, 1000]]
return ids
def __pad(self, ids, max_len, simple=False, value=None):
pad_value = value if value is not None else self.pad_id
if len(ids) < max_len:
if simple:
return ids + [pad_value] * (max_len - len(ids))
else:
return ids + [[pad_value, 0, 0, 0, 0]] * (max_len - len(ids))
else:
assert len(ids) == max_len
return ids
def __getitem__(self, idx):
if self.layout_flag:
return self.__getitem_layout__(idx)
else:
return self.__getitem_bert__(idx)
def __getitem_bert__(self, idx):
idx = (self.offset + idx) % len(self.features)
feature = self.features[idx]
source_ids = self.__trunk([self.cls_id] + feature["source_ids"], self.max_source_len, simple=True)
target_ids = self.__trunk(feature["target_ids"], self.max_target_len, simple=True)
target_index = self.__trunk(feature['target_index'], self.max_target_len, simple=True, value=self.index_sp_id)
pseudo_ids = []
for tk_id in target_ids:
p = random.random()
if p < self.keep_prob:
pseudo_ids.append(tk_id)
elif p < self.keep_prob + self.random_prob:
pseudo_ids.append(random.randint(0, self.vocab_size - 1))
else:
pseudo_ids.append(self.mask_id)
num_source_tokens = len(source_ids)
num_target_tokens = len(target_ids)
source_ids = self.__pad(source_ids, self.max_source_len, simple=True)
target_ids = self.__pad(target_ids, self.max_target_len, simple=True)
pseudo_ids = self.__pad(pseudo_ids, self.max_target_len, simple=True)
target_index = self.__pad(target_index, self.max_target_len, simple=True, value=self.index_sp_id)
target_index = self.__clip_index(target_index)
if self.span_len > 1:
span_ids = []
span_id = 1
while len(span_ids) < num_target_tokens:
p = random.random()
if p < self.span_prob:
span_len = random.randint(2, self.span_len)
span_len = min(span_len, num_target_tokens - len(span_ids))
else:
span_len = 1
span_ids.extend([span_id] * span_len)
span_id += 1
span_ids = self.__pad(span_ids, self.max_target_len)
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, span_ids, target_index
else:
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, target_index
def __getitem_layout__(self, idx):
# TODO: how to initialize the random and masked tokens' pos emb
# Simple Solution: only mask the text
idx = (self.offset + idx) % len(self.features)
feature = self.features[idx]
source_ids = self.__trunk([[self.cls_id, 0, 0, 0, 0]] + feature["source_ids"], self.max_source_len)
target_ids = self.__trunk(feature["target_ids"], self.max_target_len)
target_index = self.__trunk(feature['target_index'], self.max_target_len, simple=True, value=self.index_sp_id)
pseudo_ids = []
for tk_id in target_ids:
p = random.random()
if p < self.keep_prob:
pseudo_ids.append(tk_id)
elif p < self.keep_prob + self.random_prob:
pseudo_ids.append([random.randint(0, self.vocab_size - 1)] + [0, 0, 0, 0]) # tk_id[1:])
else:
pseudo_ids.append([self.mask_id] + [0, 0, 0, 0]) # tk_id[1:])
num_source_tokens = len(source_ids)
num_target_tokens = len(target_ids)
source_ids = self.__pad(source_ids, self.max_source_len)
target_ids = self.__pad(target_ids, self.max_target_len)
pseudo_ids = self.__pad(pseudo_ids, self.max_target_len)
target_index = self.__pad(target_index, self.max_target_len, simple=True, value=self.index_sp_id)
target_index = self.__clip_index(target_index)
if self.span_len > 1:
span_ids = []
span_id = 1
while len(span_ids) < num_target_tokens:
p = random.random()
if p < self.span_prob:
span_len = random.randint(2, self.span_len)
span_len = min(span_len, num_target_tokens - len(span_ids))
else:
span_len = 1
span_ids.extend([span_id] * span_len)
span_id += 1
span_ids = self.__pad(span_ids, self.max_target_len)
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, span_ids, target_index
else:
return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, target_index
def batch_list_to_batch_tensors(batch):
batch_tensors = []
for x in zip(*batch):
if isinstance(x[0], torch.Tensor):
batch_tensors.append(torch.stack(x))
else:
batch_tensors.append(torch.tensor(x, dtype=torch.long))
return batch_tensors
def get_max_epoch_model(output_dir):
fn_model_list = glob.glob(os.path.join(output_dir, "model.*.bin"))
fn_optim_list = glob.glob(os.path.join(output_dir, "optim.*.bin"))
if (not fn_model_list) or (not fn_optim_list):
return None
os.path.basename(output_dir)
both_set = set([int(os.path.basename(fn).split('.')[1]) for fn in fn_model_list]
) & set([int(os.path.basename(fn).split('.')[1]) for fn in fn_optim_list])
if both_set:
return max(both_set)
else:
return None
def load_and_cache_examples(
example_file, tokenizer, local_rank, cached_features_file, shuffle=True):
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
if local_rank not in [-1, 0]:
torch.distributed.barrier()
if cached_features_file is not None and os.path.exists(cached_features_file):
logger.info("Loading features from cached file %s", cached_features_file)
features = torch.load(cached_features_file)
else:
logger.info("Creating features from dataset file at %s", example_file)
examples = []
with open(example_file, mode="r", encoding="utf-8") as reader:
for i, line in enumerate(reader):
if i == 100:
break
examples.append(json.loads(line))
features = []
for example in tqdm.tqdm(examples):
if isinstance(example["src"], list):
source_tokens = example["src"]
target_tokens = example["tgt"]
else:
source_tokens = tokenizer.tokenize(example["src"])
target_tokens = tokenizer.tokenize(example["tgt"])
features.append({
"source_ids": tokenizer.convert_tokens_to_ids(source_tokens),
"target_ids": tokenizer.convert_tokens_to_ids(target_tokens),
})
if shuffle:
random.shuffle(features)
if local_rank in [-1, 0] and cached_features_file is not None:
logger.info("Saving features into cached file %s", cached_features_file)
torch.save(features, cached_features_file)
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
if local_rank == 0:
torch.distributed.barrier()
return features
def load_and_cache_line_order_examples(
example_path, tokenizer, local_rank, cached_features_file, max_src_length=1024,
layout_flag=True, shuffle=True,
src_shuffle_rate=0,
file_info_flag=False,
):
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
if local_rank not in [-1, 0]:
torch.distributed.barrier()
if cached_features_file is not None and os.path.exists(cached_features_file) and False:
logger.info("Loading features from cached file %s", cached_features_file)
features = torch.load(cached_features_file)
else:
logger.info("Creating features from dataset at %s", example_path)
examples = []
with open(example_path, 'r') as layout_reader:
logger.info(f'Start loading {example_path}')
for i, line in enumerate(layout_reader):
examples.append(json.loads(line))
features = []
for layout in tqdm.tqdm(examples):
bleu = layout['bleu']
if random.random() < src_shuffle_rate:
# print('Random!!!')
# DONE: the random src! here has bug! index also need shuffle
src_layout = layout['src']
tgt_index = layout['tgt_index']
source_length = len(src_layout)
shuffle_index = list(range(source_length))
random.shuffle(shuffle_index)
shuffle_layout = ['' for _ in range(source_length)]
for i, j in enumerate(shuffle_index):
# NOTE: map i-th token to j-th token
shuffle_layout[j] = src_layout[i]
shuffle_target_index = [shuffle_index[i] for i in tgt_index]
layout['tgt_index'] = shuffle_target_index
layout['src'] = shuffle_layout
mask = tokenizer.mask_token_id
src_ids = [tokenizer.convert_tokens_to_ids([str(tmp_i)])[:1] + src_layout for tmp_i, src_layout in enumerate(layout['src'])]
tgt_ids = [tokenizer.convert_tokens_to_ids([str(tmp_i)])[:1] + tgt_layout for tmp_i, tgt_layout in enumerate(layout['tgt'])]
tgt_index = layout['tgt_index']
feature = {
"source_ids": src_ids,
"target_ids": tgt_ids,
"target_index": tgt_index,
'bleu': bleu
}
if file_info_flag:
file_info = {'original_filename': layout['filename'], 'filename': layout['filename'],
'page_idx': 0}
feature['file_info'] = file_info
features.append(feature)
if shuffle:
random.shuffle(features)
if local_rank in [-1, 0] and cached_features_file is not None:
logger.info("Saving features into cached file %s", cached_features_file)
torch.save(features, cached_features_file)
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
if local_rank == 0:
torch.distributed.barrier()
return features
def load_and_cache_layoutlm_examples(
example_path, tokenizer, local_rank, cached_features_file, max_src_length=1024,
layout_flag=True, shuffle=True,
src_shuffle_rate=0,
file_info_flag=False
):
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
if local_rank not in [-1, 0]:
torch.distributed.barrier()
if cached_features_file is not None and os.path.exists(cached_features_file):
logger.info("Loading features from cached file %s", cached_features_file)
features = torch.load(cached_features_file)
else:
logger.info("Creating features from dataset at %s", example_path)
examples = []
if os.path.isdir(example_path):
text_files = glob.glob(f'{example_path}/*text*.json')
layout_files = [re.sub('text|txt', 'layout', x, 1) for x in text_files]
else:
text_files = [example_path]
layout_files = [re.sub('text|txt', 'layout', example_path, 1)]
for text_file, layout_file in zip(text_files, layout_files):
with open(text_file, mode='r', encoding='utf-8') as text_reader, \
open(layout_file, mode='r', encoding='utf-8') as layout_reader:
logger.info(f'Start loading {text_file}')
for i, (text_line, layout_line) in enumerate(zip(text_reader, layout_reader)):
if (i + 1) % 10000 == 0:
logger.info(f'{i + 1} lines ...')
examples.append((json.loads(text_line), json.loads(layout_line)))
features = []
def tokenize_text_and_layout_src(_text, _layout, _layout_flag):
ret = []
index_split = {}
words = _text.split()
# note: (OLD) the index should start from 1: 0-the cls token in src
# note: (NEW) we need to remove the src embedding's CLS SEP token so we can still start from 0
# note: (NEWER) we need to at least one blank pos for ignore index in loss function (we use sep's index)
# NOTE: (NEWER-ER) 1 for all padding tgt index
new_token_index = 1 # first ordinary index
for i, (word, box) in enumerate(zip(words, _layout)):
if (not box[2] >= box[0]) or (not box[3] >= box[1]):
continue
tokens = tokenizer.tokenize(word)
tokens = tokenizer.convert_tokens_to_ids(tokens)
new_token_ids = []
for token in tokens:
if _layout_flag:
ret.append([token] + box)
else:
ret.append(token)
new_token_ids.append(new_token_index)
new_token_index += 1
index_split[i] = new_token_ids
return ret, index_split
def tokenize_text_and_layout_tgt(_text, _layout, _index, _index_split, _layout_flag):
ret = []
ret_index = []
words = _text.split()
for word, box, i in zip(words, _layout, _index):
if (not box[2] >= box[0]) or (not box[3] >= box[1]):
continue
tokens = tokenizer.tokenize(word)
tokens = tokenizer.convert_tokens_to_ids(tokens)
for token, ii in zip(tokens, _index_split[i]):
if _layout_flag:
ret.append([token] + box)
else:
ret.append(token)
ii = min(ii, max_src_length - 1)
ret_index.append(ii)
return ret, ret_index
for text, layout in tqdm.tqdm(examples):
if 'bleu' in text:
bleu = text['bleu']
else:
bleu = 0
if random.random() < src_shuffle_rate:
# print('Random!!!')
# DONE: the random src! here has bug! index also need shuffle
src_text = text['src']
src_layout = layout['src']
tgt_index = text['tgt_index']
src_text = src_text.split()
source_length = len(src_text)
shuffle_index = list(range(source_length))
random.shuffle(shuffle_index)
shuffle_text = ['' for _ in range(source_length)]
shuffle_layout = ['' for _ in range(source_length)]
for i, j in enumerate(shuffle_index):
# NOTE: map i-th token to j-th token
shuffle_text[j] = src_text[i]
shuffle_layout[j] = src_layout[i]
shuffle_target_index = [shuffle_index[i] for i in tgt_index]
text['src'] = ' '.join(shuffle_text)
text['tgt_index'] = shuffle_target_index
layout['src'] = shuffle_layout
src_ids, src_index_split = tokenize_text_and_layout_src(text['src'], layout['src'],
_layout_flag=layout_flag)
tgt_ids, tgt_index = tokenize_text_and_layout_tgt(text['tgt'], layout['tgt'], text['tgt_index'],
src_index_split, _layout_flag=layout_flag)
feature = {
"source_ids": src_ids,
"target_ids": tgt_ids,
"target_index": tgt_index,
'bleu': bleu
}
if file_info_flag:
file_info = {'original_filename': text['original_filename'], 'filename': text['filename'], 'page_idx': text['page_idx']}
feature['file_info'] = file_info
features.append(feature)
if shuffle:
random.shuffle(features)
if local_rank in [-1, 0] and cached_features_file is not None:
if not os.path.exists(os.path.dirname(cached_features_file)):
os.makedirs(os.path.dirname(cached_features_file))
logger.info("Saving features into cached file %s", cached_features_file)
torch.save(features, cached_features_file)
# Make sure only the first process in distributed training process the dataset, and the others will use the cache
if local_rank == 0:
torch.distributed.barrier()
return features
def convert_src_layout_inputs_to_tokens(inputs, converter, max_src_length, layout_flag=True):
ret = []
if not layout_flag:
for line in inputs:
ret.append(converter(line["source_ids"])[: max_src_length])
else:
for line in inputs:
raw_text_ids = [x[0] for x in line['source_ids']]
raw_text = converter(raw_text_ids)
new_line = [[t] + x[1:] for t, x in zip(raw_text, line['source_ids'])][: max_src_length]
ret.append(new_line)
return ret
def convert_tgt_layout_inputs_to_tokens(inputs, converter, max_tgt_length, layout_flag=True):
ret = []
if not layout_flag:
for line in inputs:
ret.append(converter(line["target_ids"])[: max_tgt_length])
else:
for line in inputs:
raw_text_ids = [x[0] for x in line['target_ids']]
ret.append(converter(raw_text_ids)[: max_tgt_length])
return ret
def get_tokens_from_src_and_index(src, index, modifier=None):
result = []
for i in index:
i = modifier(i)
i = min(i, len(src) - 1)
if isinstance(src[i], list):
result.append(src[i][0])
else:
result.append(src[i])
return result
def get_layout_from_src_and_index(src, index, modifier=None):
result = []
s = set()
for i in index:
i = modifier(i)
i = min(i, len(src) - 1)
layout = src[i][1:]
if repr(layout) not in s:
result.append(layout)
s.add(repr(layout))
return result
def get_everything_from_src_and_index(src, index, modifier=None):
result = []
for i in index:
i = modifier(i)
i = min(i, len(src) - 1)
result.append(src[i])
return result