chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,138 @@
|
||||
# LayoutReader
|
||||
|
||||
LayoutReader captures the text and layout information for reading order prediction using the seq2seq model. It significantly improves both open-source and commercial OCR engines in ordering text lines in their results in our experiments.
|
||||
|
||||
|
||||
Our paper "[LayoutReader: Pre-training of Text and Layout for Reading Order Detection](https://arxiv.org/pdf/2108.11591.pdf)" has been accepted by EMNLP 2021.
|
||||
|
||||
**ReadingBank** is a benchmark dataset for reading order detection built with weak supervision from WORD documents, which contains 500K document images with a wide range of document types as well as the corresponding reading order information. For more details, please refer to [ReadingBank](https://aka.ms/readingbank).
|
||||
|
||||
## Installation
|
||||
~~~
|
||||
conda create -n LayoutReader python=3.7
|
||||
conda activate LayoutReader
|
||||
conda install pytorch==1.7.1 -c pytorch
|
||||
pip install nltk
|
||||
python -c "import nltk; nltk.download('punkt')"
|
||||
git clone https://github.com/NVIDIA/apex.git && cd apex && python setup.py install --cuda_ext --cpp_ext
|
||||
pip install transformers==2.10.0
|
||||
git clone https://github.com/microsoft/unilm.git
|
||||
cd unilm/layoutreader
|
||||
pip install -e .
|
||||
~~~
|
||||
|
||||
## Run
|
||||
1. Download the pre-processed data ([`ReadingBank.zip`](https://mail2sysueducn-my.sharepoint.com/:u:/g/personal/huangyp28_mail2_sysu_edu_cn/Efh3ZWjsA-xFrH2FSjyhSVoBMak6ypmbABWmJEmPwtKhhw?e=tbthMD)). For more details of the dataset, please refer to [ReadingBank](https://aka.ms/readingbank).
|
||||
2. (Optional) Download our pre-trained model ([`layoutreader-base-readingbank.zip`](https://mail2sysueducn-my.sharepoint.com/:u:/g/personal/huangyp28_mail2_sysu_edu_cn/ET9XynvgSZFLhPy7p30zbtoBs-T_Yxj6gl_k-b2-N53ChQ?e=gKafBy)) and evaluate it refer to step 4.
|
||||
3. Training
|
||||
~~~
|
||||
export CUDA_VISIBLE_DEVICE=0,1,2,3
|
||||
export OMP_NUM_THREADS=4
|
||||
export MKL_NUM_THREADS=4
|
||||
|
||||
python -m torch.distributed.launch --nproc_per_node=4 run_seq2seq.py \
|
||||
--model_type layoutlm \
|
||||
--model_name_or_path layoutlm-base-uncased \
|
||||
--train_folder /path/to/ReadingBank/train \
|
||||
--output_dir /path/to/output/LayoutReader/layoutlm \
|
||||
--do_lower_case \
|
||||
--fp16 \
|
||||
--fp16_opt_level O2 \
|
||||
--max_source_seq_length 513 \
|
||||
--max_target_seq_length 511 \
|
||||
--per_gpu_train_batch_size 2 \
|
||||
--gradient_accumulation_steps 1 \
|
||||
--learning_rate 7e-5 \
|
||||
--num_warmup_steps 500 \
|
||||
--num_training_steps 75000 \
|
||||
--cache_dir /path/to/output/LayoutReader/cache \
|
||||
--label_smoothing 0.1 \
|
||||
--save_steps 5000 \
|
||||
--cached_train_features_file /path/to/ReadingBank/features_train.pt
|
||||
~~~
|
||||
4. Decoding
|
||||
~~~
|
||||
export CUDA_VISIBLE_DEVICES=0
|
||||
export OMP_NUM_THREADS=4
|
||||
export MKL_NUM_THREADS=4
|
||||
|
||||
python decode_seq2seq.py --fp16 \
|
||||
--model_type layoutlm \
|
||||
--tokenizer_name bert-base-uncased \
|
||||
--input_folder /path/to/ReadingBank/test \
|
||||
--cached_feature_file /path/to/ReadingBank/features_test.pt \
|
||||
--output_file /path/to/output/LayoutReader/layoutlm/output.txt \
|
||||
--split test \
|
||||
--do_lower_case \
|
||||
--model_path /path/to/output/LayoutReader/layoutlm/ckpt-75000 \
|
||||
--cache_dir /path/to/output/LayoutReader/cache \
|
||||
--max_seq_length 1024 \
|
||||
--max_tgt_length 511 \
|
||||
--batch_size 32 \
|
||||
--beam_size 1 \
|
||||
--length_penalty 0 \
|
||||
--forbid_duplicate_ngrams \
|
||||
--mode s2s \
|
||||
--forbid_ignore_word "."
|
||||
~~~
|
||||
|
||||
## Results
|
||||
Our released [pre-trained model](https://mail2sysueducn-my.sharepoint.com/:u:/g/personal/huangyp28_mail2_sysu_edu_cn/ET9XynvgSZFLhPy7p30zbtoBs-T_Yxj6gl_k-b2-N53ChQ?e=gKafBy) achieves 98.2% Average Page-level BLEU score. Detailed results are reported as follow:
|
||||
|
||||
* Evaluation results of the LayoutReader on the reading order detection task, where the source-side of training/testing data is in the left-to-right and top-to-bottom order
|
||||
|
||||
| Method | Encoder | Avg. Page-level BLEU ↑ | ARD ↓ |
|
||||
| -------------------------- | ---------------------- | ---------------------- | ----- |
|
||||
| Heuristic Method | - | 0.6972 | 8.46 |
|
||||
| LayoutReader (text only) | BERT | 0.8510 | 12.08 |
|
||||
| LayoutReader (text only) | UniLM | 0.8765 | 10.65 |
|
||||
| LayoutReader (layout only) | LayoutLM (layout only) | 0.9732 | 2.31 |
|
||||
| LayoutReader | LayoutLM | 0.9819 | 1.75 |
|
||||
|
||||
* Input order study with left-to-right and top-to-bottom inputs in evaluation, where r is the proportion of
|
||||
shuffled samples in training.
|
||||
|
||||
| Method | Avg. Page-level BLEU ↑ | Avg. Page-level BLEU ↑ | Avg. Page-level BLEU ↑ | ARD ↓ | ARD ↓ | ARD ↓ |
|
||||
|---------------------------------|------------------------|------------------------|------------------------|--------|-------|-------|
|
||||
| | r=100% | r=50% | r=0% | r=100% | r=50% | r=0% |
|
||||
| LayoutReader (text only, BERT) | 0.3355 | 0.8397 | 0.8510 | 77.97 | 15.62 | 12.08 |
|
||||
| LayoutReader (text only, UniLM) | 0.3440 | 0.8588 | 0.8765 | 78.67 | 13.65 | 10.65 |
|
||||
| LayoutReader (layout only) | 0.9701 | 0.9729 | 0.9732 | 2.85 | 2.61 | 2.31 |
|
||||
| LayoutReader | 0.9765 | 0.9788 | 0.9819 | 2.50 | 2.24 | 1.75 |
|
||||
|
||||
* Input order study with token-shuffled inputs in evaluation, where r is the proportion of shuffled samples in training.
|
||||
|
||||
| Method | Avg. Page-level BLEU ↑ | Avg. Page-level BLEU ↑ | Avg. Page-level BLEU ↑ | ARD ↓ | ARD ↓ | ARD ↓ |
|
||||
|---------------------------------|------------------------|------------------------|------------------------|--------|-------|--------|
|
||||
| | r=100% | r=50% | r=0% | r=100% | r=50% | r=0% |
|
||||
| LayoutReader (text only, BERT) | 0.3085 | 0.2730 | 0.1711 | 78.69 | 85.44 | 67.96 |
|
||||
| LayoutReader (text only, UniLM) | 0.3119 | 0.2855 | 0.1728 | 80.00 | 85.60 | 71.13 |
|
||||
| LayoutReader (layout only) | 0.9718 | 0.9714 | 0.1331 | 2.72 | 2.82 | 105.40 |
|
||||
| LayoutReader | 0.9772 | 0.9770 | 0.1783 | 2.48 | 2.46 | 72.94 |
|
||||
|
||||
## Citation
|
||||
|
||||
If you find LayoutReader helpful, please cite us:
|
||||
```
|
||||
@misc{wang2021layoutreader,
|
||||
title={LayoutReader: Pre-training of Text and Layout for Reading Order Detection},
|
||||
author={Zilong Wang and Yiheng Xu and Lei Cui and Jingbo Shang and Furu Wei},
|
||||
year={2021},
|
||||
eprint={2108.11591},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## License
|
||||
|
||||
This project is licensed under the license found in the LICENSE file in the root directory of this source tree.
|
||||
Portions of the source code are based on the [transformers](https://github.com/huggingface/transformers) and [s2s-ft](../s2s-ft) projects.
|
||||
[Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct)
|
||||
|
||||
## Contact
|
||||
|
||||
For help or issues using LayoutReader, please submit a GitHub issue.
|
||||
|
||||
For other communications related to LayoutLM, please contact Lei Cui (`lecu@microsoft.com`), Furu Wei (`fuwei@microsoft.com`).
|
||||
@@ -0,0 +1,345 @@
|
||||
"""BERT finetuning runner."""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import pickle
|
||||
import random
|
||||
from time import sleep
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from nltk.translate.bleu_score import sentence_bleu
|
||||
from tqdm import tqdm
|
||||
from transformers import \
|
||||
BertTokenizer, RobertaTokenizer
|
||||
from transformers.tokenization_bert import whitespace_tokenize
|
||||
|
||||
import s2s_ft.s2s_loader as seq2seq_loader
|
||||
from s2s_ft.modeling_decoding import LayoutlmForSeq2SeqDecoder, BertConfig
|
||||
from s2s_ft.tokenization_minilm import MinilmTokenizer
|
||||
from s2s_ft.tokenization_unilm import UnilmTokenizer
|
||||
from s2s_ft.utils import load_and_cache_layoutlm_examples, convert_src_layout_inputs_to_tokens, \
|
||||
get_tokens_from_src_and_index, convert_tgt_layout_inputs_to_tokens
|
||||
|
||||
TOKENIZER_CLASSES = {
|
||||
'bert': BertTokenizer,
|
||||
'minilm': MinilmTokenizer,
|
||||
'roberta': RobertaTokenizer,
|
||||
'unilm': UnilmTokenizer,
|
||||
'layoutlm': BertTokenizer,
|
||||
}
|
||||
|
||||
|
||||
class WhitespaceTokenizer(object):
|
||||
def tokenize(self, text):
|
||||
return whitespace_tokenize(text)
|
||||
|
||||
|
||||
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
|
||||
datefmt='%m/%d/%Y %H:%M:%S',
|
||||
level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def detokenize(tk_list):
|
||||
r_list = []
|
||||
for tk in tk_list:
|
||||
if tk.startswith('##') and len(r_list) > 0:
|
||||
r_list[-1] = r_list[-1] + tk[2:]
|
||||
else:
|
||||
r_list.append(tk)
|
||||
return r_list
|
||||
|
||||
|
||||
def ascii_print(text):
|
||||
text = text.encode("ascii", "ignore")
|
||||
print(text)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
# Required parameters
|
||||
parser.add_argument("--model_type", default=None, type=str, required=True,
|
||||
help="Model type selected in the list: " + ", ".join(TOKENIZER_CLASSES.keys()))
|
||||
parser.add_argument("--model_path", default=None, type=str, required=True,
|
||||
help="Path to the model checkpoint.")
|
||||
parser.add_argument("--config_path", default=None, type=str,
|
||||
help="Path to config.json for the model.")
|
||||
|
||||
parser.add_argument("--sentence_shuffle_rate", default=0, type=float)
|
||||
parser.add_argument("--layoutlm_only_layout", action='store_true')
|
||||
|
||||
# tokenizer_name
|
||||
parser.add_argument("--tokenizer_name", default=None, type=str, required=True,
|
||||
help="tokenizer name")
|
||||
parser.add_argument("--max_seq_length", default=512, type=int,
|
||||
help="The maximum total input sequence length after WordPiece tokenization. \n"
|
||||
"Sequences longer than this will be truncated, and sequences shorter \n"
|
||||
"than this will be padded.")
|
||||
|
||||
# decoding parameters
|
||||
parser.add_argument('--fp16', action='store_true',
|
||||
help="Whether to use 16-bit float precision instead of 32-bit")
|
||||
parser.add_argument('--amp', action='store_true',
|
||||
help="Whether to use amp for fp16")
|
||||
parser.add_argument("--input_file", type=str, help="Input file")
|
||||
parser.add_argument("--input_folder", type=str, help="Input folder")
|
||||
parser.add_argument("--cached_feature_file", type=str)
|
||||
parser.add_argument('--subset', type=int, default=0,
|
||||
help="Decode a subset of the input dataset.")
|
||||
parser.add_argument("--output_file", type=str, help="output file")
|
||||
parser.add_argument("--split", type=str, default="",
|
||||
help="Data split (train/val/test).")
|
||||
parser.add_argument('--tokenized_input', action='store_true',
|
||||
help="Whether the input is tokenized.")
|
||||
parser.add_argument('--seed', type=int, default=123,
|
||||
help="random seed for initialization")
|
||||
parser.add_argument("--do_lower_case", action='store_true',
|
||||
help="Set this flag if you are using an uncased model.")
|
||||
parser.add_argument('--batch_size', type=int, default=4,
|
||||
help="Batch size for decoding.")
|
||||
parser.add_argument('--beam_size', type=int, default=1,
|
||||
help="Beam size for searching")
|
||||
parser.add_argument('--length_penalty', type=float, default=0,
|
||||
help="Length penalty for beam search")
|
||||
|
||||
parser.add_argument('--forbid_duplicate_ngrams', action='store_true')
|
||||
parser.add_argument('--forbid_ignore_word', type=str, default=None,
|
||||
help="Forbid the word during forbid_duplicate_ngrams")
|
||||
parser.add_argument("--min_len", default=1, type=int)
|
||||
parser.add_argument('--need_score_traces', action='store_true')
|
||||
parser.add_argument('--ngram_size', type=int, default=3)
|
||||
parser.add_argument('--mode', default="s2s",
|
||||
choices=["s2s", "l2r", "both"])
|
||||
parser.add_argument('--max_tgt_length', type=int, default=128,
|
||||
help="maximum length of target sequence")
|
||||
parser.add_argument('--s2s_special_token', action='store_true',
|
||||
help="New special tokens ([S2S_SEP]/[S2S_CLS]) of S2S.")
|
||||
parser.add_argument('--s2s_add_segment', action='store_true',
|
||||
help="Additional segmental for the encoder of S2S.")
|
||||
parser.add_argument('--s2s_share_segment', action='store_true',
|
||||
help="Sharing segment embeddings for the encoder of S2S (used with --s2s_add_segment).")
|
||||
parser.add_argument('--pos_shift', action='store_true',
|
||||
help="Using position shift for fine-tuning.")
|
||||
parser.add_argument("--cache_dir", default=None, type=str,
|
||||
help="Where do you want to store the pre-trained models downloaded from s3")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
model_path = args.model_path
|
||||
assert os.path.exists(model_path), 'model_path ' + model_path + ' not exists!'
|
||||
|
||||
if args.need_score_traces and args.beam_size <= 1:
|
||||
raise ValueError(
|
||||
"Score trace is only available for beam search with beam size > 1.")
|
||||
if args.max_tgt_length >= args.max_seq_length - 2:
|
||||
raise ValueError("Maximum tgt length exceeds max seq length - 2.")
|
||||
|
||||
device = torch.device(
|
||||
"cuda" if torch.cuda.is_available() else "cpu")
|
||||
n_gpu = torch.cuda.device_count()
|
||||
|
||||
if args.seed > 0:
|
||||
random.seed(args.seed)
|
||||
np.random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
if n_gpu > 0:
|
||||
torch.cuda.manual_seed_all(args.seed)
|
||||
else:
|
||||
random_seed = random.randint(0, 10000)
|
||||
logger.info("Set random seed as: {}".format(random_seed))
|
||||
random.seed(random_seed)
|
||||
np.random.seed(random_seed)
|
||||
torch.manual_seed(random_seed)
|
||||
if n_gpu > 0:
|
||||
torch.cuda.manual_seed_all(args.seed)
|
||||
|
||||
tokenizer = TOKENIZER_CLASSES[args.model_type].from_pretrained(
|
||||
args.tokenizer_name, do_lower_case=args.do_lower_case,
|
||||
cache_dir=args.cache_dir if args.cache_dir else None,
|
||||
max_len=args.max_seq_length
|
||||
)
|
||||
|
||||
if args.model_type == "roberta":
|
||||
vocab = tokenizer.encoder
|
||||
else:
|
||||
vocab = tokenizer.vocab
|
||||
|
||||
# NOTE: tokenizer cannot setattr, so move this to the initialization step
|
||||
# tokenizer.max_len = args.max_seq_length
|
||||
|
||||
config_file = args.config_path if args.config_path else os.path.join(args.model_path, "config.json")
|
||||
logger.info("Read decoding config from: %s" % config_file)
|
||||
config = BertConfig.from_json_file(config_file,
|
||||
# base_model_type=args.model_type
|
||||
layoutlm_only_layout_flag=args.layoutlm_only_layout
|
||||
)
|
||||
|
||||
bi_uni_pipeline = []
|
||||
bi_uni_pipeline.append(seq2seq_loader.Preprocess4Seq2seqDecoder(
|
||||
list(vocab.keys()), tokenizer.convert_tokens_to_ids, args.max_seq_length,
|
||||
max_tgt_length=args.max_tgt_length, pos_shift=args.pos_shift,
|
||||
source_type_id=config.source_type_id, target_type_id=config.target_type_id,
|
||||
cls_token=tokenizer.cls_token, sep_token=tokenizer.sep_token, pad_token=tokenizer.pad_token,
|
||||
layout_flag=args.model_type == 'layoutlm'
|
||||
))
|
||||
|
||||
mask_word_id, eos_word_ids, sos_word_id = tokenizer.convert_tokens_to_ids(
|
||||
[tokenizer.mask_token, tokenizer.sep_token, tokenizer.sep_token])
|
||||
forbid_ignore_set = None
|
||||
if args.forbid_ignore_word:
|
||||
w_list = []
|
||||
for w in args.forbid_ignore_word.split('|'):
|
||||
if w.startswith('[') and w.endswith(']'):
|
||||
w_list.append(w.upper())
|
||||
else:
|
||||
w_list.append(w)
|
||||
forbid_ignore_set = set(tokenizer.convert_tokens_to_ids(w_list))
|
||||
print(args.model_path)
|
||||
found_checkpoint_flag = False
|
||||
for model_recover_path in [args.model_path.strip()]:
|
||||
logger.info("***** Recover model: %s *****", model_recover_path)
|
||||
found_checkpoint_flag = True
|
||||
model = LayoutlmForSeq2SeqDecoder.from_pretrained(
|
||||
model_recover_path, config=config, mask_word_id=mask_word_id, search_beam_size=args.beam_size,
|
||||
length_penalty=args.length_penalty, eos_id=eos_word_ids, sos_id=sos_word_id,
|
||||
forbid_duplicate_ngrams=args.forbid_duplicate_ngrams, forbid_ignore_set=forbid_ignore_set,
|
||||
ngram_size=args.ngram_size, min_len=args.min_len, mode=args.mode,
|
||||
max_position_embeddings=args.max_seq_length, pos_shift=args.pos_shift,
|
||||
)
|
||||
|
||||
if args.fp16:
|
||||
model.half()
|
||||
model.to(device)
|
||||
if n_gpu > 1:
|
||||
model = torch.nn.DataParallel(model)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
model.eval()
|
||||
next_i = 0
|
||||
max_src_length = args.max_seq_length - 2 - args.max_tgt_length
|
||||
max_tgt_length = args.max_tgt_length
|
||||
|
||||
example_path = args.input_file if args.input_file else args.input_folder
|
||||
|
||||
to_pred = load_and_cache_layoutlm_examples(
|
||||
example_path, tokenizer, local_rank=-1,
|
||||
cached_features_file=args.cached_feature_file, shuffle=False, layout_flag=args.model_type == 'layoutlm',
|
||||
src_shuffle_rate=args.sentence_shuffle_rate
|
||||
)
|
||||
|
||||
input_lines = convert_src_layout_inputs_to_tokens(to_pred, tokenizer.convert_ids_to_tokens, max_src_length,
|
||||
layout_flag=args.model_type == 'layoutlm')
|
||||
target_lines = convert_tgt_layout_inputs_to_tokens(to_pred, tokenizer.convert_ids_to_tokens, max_tgt_length,
|
||||
layout_flag=args.model_type == 'layoutlm')
|
||||
target_geo_scores = [x['bleu'] for x in to_pred]
|
||||
|
||||
if args.subset > 0:
|
||||
logger.info("Decoding subset: %d", args.subset)
|
||||
input_lines = input_lines[:args.subset]
|
||||
|
||||
# NOTE: add the sequence index through enumerate
|
||||
input_lines = sorted(list(enumerate(input_lines)), key=lambda x: -len(x[1]))
|
||||
|
||||
score_trace_list = [None] * len(input_lines)
|
||||
total_batch = math.ceil(len(input_lines) / args.batch_size)
|
||||
|
||||
fn_out = args.output_file
|
||||
fout = open(fn_out, "w", encoding="utf-8")
|
||||
|
||||
with tqdm(total=total_batch) as pbar:
|
||||
batch_count = 0
|
||||
first_batch = True
|
||||
while first_batch or (next_i + args.batch_size <= len(input_lines)):
|
||||
# while next_i < len(input_lines):
|
||||
_chunk = input_lines[next_i:next_i + args.batch_size]
|
||||
buf_id = [x[0] for x in _chunk]
|
||||
buf = [x[1] for x in _chunk]
|
||||
next_i += args.batch_size
|
||||
batch_count += 1
|
||||
max_a_len = max([len(x) for x in buf])
|
||||
instances = []
|
||||
for instance in [(x, max_a_len) for x in buf]:
|
||||
for proc in bi_uni_pipeline:
|
||||
instances.append(proc(instance))
|
||||
with torch.no_grad():
|
||||
batch = seq2seq_loader.batch_list_to_batch_tensors(
|
||||
instances)
|
||||
batch = [
|
||||
t.to(device) if t is not None else None for t in batch]
|
||||
input_ids, token_type_ids, position_ids, input_mask, mask_qkv, task_idx = batch
|
||||
traces = model(input_ids, token_type_ids,
|
||||
position_ids, input_mask, task_idx=task_idx, mask_qkv=mask_qkv)
|
||||
if args.beam_size > 1:
|
||||
traces = {k: v.tolist() for k, v in traces.items()}
|
||||
output_ids = traces['pred_seq']
|
||||
else:
|
||||
output_ids = traces.tolist()
|
||||
for i in range(len(buf)):
|
||||
w_ids = output_ids[i]
|
||||
output_buf = get_tokens_from_src_and_index(src=buf[i], index=w_ids, modifier=lambda x: x-1)
|
||||
output_tokens = []
|
||||
for t in output_buf:
|
||||
if t in (tokenizer.sep_token, tokenizer.pad_token):
|
||||
break
|
||||
output_tokens.append(t)
|
||||
output_tokens = output_tokens[:len(target_lines[buf_id[i]])]
|
||||
if args.model_type == "roberta":
|
||||
output_sequence = tokenizer.convert_tokens_to_string(output_tokens)
|
||||
else:
|
||||
output_sequence = ' '.join(detokenize(output_tokens))
|
||||
if '\n' in output_sequence:
|
||||
output_sequence = " [X_SEP] ".join(output_sequence.split('\n'))
|
||||
|
||||
target = target_lines[buf_id[i]]
|
||||
target = detokenize(target)
|
||||
result = output_sequence.split()
|
||||
score = sentence_bleu([target], result)
|
||||
|
||||
geo_score = target_geo_scores[buf_id[i]]
|
||||
target_sequence = ' '.join(target)
|
||||
|
||||
fout.write('{}\t{:.8f}\t{:.8f}\t{}\t{}\n'.format(buf_id[i], score, geo_score, output_sequence, target_sequence))
|
||||
|
||||
if first_batch or batch_count % 50 == 0:
|
||||
logger.info("{}: BLEU={:.4f} GEO={:.4f} | {}"
|
||||
.format(buf_id[i], score, target_geo_scores[buf_id[i]], output_sequence))
|
||||
if args.need_score_traces:
|
||||
score_trace_list[buf_id[i]] = {
|
||||
'scores': traces['scores'][i], 'wids': traces['wids'][i], 'ptrs': traces['ptrs'][i]}
|
||||
pbar.update(1)
|
||||
first_batch = False
|
||||
|
||||
outscore = open(fn_out, encoding='utf-8')
|
||||
bleu_score = geo_score = {}
|
||||
total_bleu = total_geo = 0.0
|
||||
for line in outscore.readlines():
|
||||
id, bleu, geo, out_seq, tgt_seq = line.split('\t')
|
||||
bleu_score[int(id)] = float(bleu)
|
||||
total_bleu += float(bleu)
|
||||
geo_score[int(id)] = float(geo)
|
||||
total_geo += float(geo)
|
||||
print("avg_bleu", round(100 * total_bleu / len(bleu_score), 1))
|
||||
print("avg_geo", round(100 * total_geo / len(geo_score), 1))
|
||||
# released model (layoutreader-base-readingbank): avg_bleu 98.2, avg_geo 69.7
|
||||
|
||||
if args.need_score_traces:
|
||||
with open(fn_out + ".trace.pickle", "wb") as fout_trace:
|
||||
pickle.dump(
|
||||
{"version": 0.0, "num_samples": len(input_lines)}, fout_trace)
|
||||
for x in score_trace_list:
|
||||
pickle.dump(x, fout_trace)
|
||||
|
||||
if not found_checkpoint_flag:
|
||||
logger.info("Not found the model checkpoint file!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,447 @@
|
||||
from __future__ import absolute_import, division, print_function
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import (DataLoader, SequentialSampler)
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
|
||||
try:
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
except:
|
||||
from tensorboardX import SummaryWriter
|
||||
|
||||
import tqdm
|
||||
|
||||
from s2s_ft.modeling import LayoutlmForSequenceToSequence, LayoutlmConfig
|
||||
from transformers import AdamW, get_linear_schedule_with_warmup
|
||||
from transformers import \
|
||||
RobertaConfig, BertConfig, \
|
||||
BertTokenizer, RobertaTokenizer, \
|
||||
XLMRobertaConfig, XLMRobertaTokenizer
|
||||
|
||||
from s2s_ft.configuration_unilm import UnilmConfig
|
||||
from s2s_ft.tokenization_unilm import UnilmTokenizer
|
||||
from s2s_ft.configuration_minilm import MinilmConfig
|
||||
from s2s_ft.tokenization_minilm import MinilmTokenizer
|
||||
|
||||
from s2s_ft import utils
|
||||
from s2s_ft.config import BertForSeq2SeqConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MODEL_CLASSES = {
|
||||
'bert': (BertConfig, BertTokenizer),
|
||||
'minilm': (MinilmConfig, MinilmTokenizer),
|
||||
'roberta': (RobertaConfig, RobertaTokenizer),
|
||||
'xlm-roberta': (XLMRobertaConfig, XLMRobertaTokenizer),
|
||||
'unilm': (UnilmConfig, UnilmTokenizer),
|
||||
'layoutlm': (LayoutlmConfig, BertTokenizer),
|
||||
}
|
||||
|
||||
|
||||
def prepare_for_training(args, model, checkpoint_state_dict, amp):
|
||||
no_decay = ['bias', 'LayerNorm.weight']
|
||||
optimizer_grouped_parameters = [
|
||||
{'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
|
||||
'weight_decay': args.weight_decay},
|
||||
{'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
|
||||
]
|
||||
optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
|
||||
|
||||
if amp:
|
||||
model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)
|
||||
if checkpoint_state_dict:
|
||||
amp.load_state_dict(checkpoint_state_dict['amp'])
|
||||
|
||||
if checkpoint_state_dict:
|
||||
optimizer.load_state_dict(checkpoint_state_dict['optimizer'])
|
||||
model.load_state_dict(checkpoint_state_dict['model'])
|
||||
|
||||
# multi-gpu training (should be after apex fp16 initialization)
|
||||
if args.n_gpu > 1:
|
||||
model = torch.nn.DataParallel(model)
|
||||
|
||||
# Distributed training (should be after apex fp16 initialization)
|
||||
if args.local_rank != -1:
|
||||
model = torch.nn.parallel.DistributedDataParallel(
|
||||
model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True)
|
||||
|
||||
return model, optimizer
|
||||
|
||||
|
||||
def train(args, training_features, model, tokenizer):
|
||||
""" Train the model """
|
||||
if args.local_rank in [-1, 0] and args.log_dir:
|
||||
tb_writer = SummaryWriter(log_dir=args.log_dir)
|
||||
else:
|
||||
tb_writer = None
|
||||
|
||||
if args.fp16:
|
||||
try:
|
||||
from apex import amp
|
||||
except ImportError:
|
||||
raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
|
||||
else:
|
||||
amp = None
|
||||
|
||||
# model recover
|
||||
recover_step = utils.get_max_epoch_model(args.output_dir)
|
||||
checkpoint_state_dict = None
|
||||
|
||||
model.to(args.device)
|
||||
model, optimizer = prepare_for_training(args, model, checkpoint_state_dict, amp=amp)
|
||||
|
||||
if args.n_gpu == 0 or args.no_cuda:
|
||||
per_node_train_batch_size = args.per_gpu_train_batch_size * args.gradient_accumulation_steps
|
||||
else:
|
||||
per_node_train_batch_size = args.per_gpu_train_batch_size * args.n_gpu * args.gradient_accumulation_steps
|
||||
|
||||
train_batch_size = per_node_train_batch_size * (torch.distributed.get_world_size() if args.local_rank != -1 else 1)
|
||||
global_step = recover_step if recover_step else 0
|
||||
|
||||
if args.num_training_steps == -1:
|
||||
args.num_training_steps = int(args.num_training_epochs * len(training_features) / train_batch_size)
|
||||
|
||||
scheduler = get_linear_schedule_with_warmup(
|
||||
optimizer, num_warmup_steps=args.num_warmup_steps,
|
||||
num_training_steps=args.num_training_steps, last_epoch=-1)
|
||||
|
||||
if checkpoint_state_dict:
|
||||
scheduler.load_state_dict(checkpoint_state_dict["lr_scheduler"])
|
||||
|
||||
train_dataset = utils.Seq2seqDatasetForLayoutlm(
|
||||
features=training_features, max_source_len=args.max_source_seq_length,
|
||||
max_target_len=args.max_target_seq_length, vocab_size=tokenizer.vocab_size,
|
||||
cls_id=tokenizer.cls_token_id, sep_id=tokenizer.sep_token_id, pad_id=tokenizer.pad_token_id,
|
||||
mask_id=tokenizer.mask_token_id, random_prob=args.random_prob, keep_prob=args.keep_prob,
|
||||
offset=train_batch_size * global_step, num_training_instances=train_batch_size * args.num_training_steps,
|
||||
layout_flag=args.model_type == 'layoutlm'
|
||||
)
|
||||
|
||||
logger.info("Check dataset:")
|
||||
for i in range(5):
|
||||
source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, target_index = train_dataset.__getitem__(
|
||||
i)
|
||||
logger.info("Instance-%d" % i)
|
||||
try:
|
||||
src = [sid[0] for sid in source_ids]
|
||||
tgt = [tid[0] for tid in target_ids]
|
||||
except TypeError:
|
||||
src = source_ids
|
||||
tgt = target_ids
|
||||
logger.info("Source tokens = %s" % " ".join(tokenizer.convert_ids_to_tokens(src)))
|
||||
logger.info("Target tokens = %s" % " ".join(tokenizer.convert_ids_to_tokens(tgt)))
|
||||
|
||||
logger.info("Mode = %s" % str(model))
|
||||
|
||||
# Train!
|
||||
logger.info(" ***** Running training ***** *")
|
||||
logger.info(" Num examples = %d", len(training_features))
|
||||
logger.info(" Num Epochs = %.2f", len(train_dataset) / len(training_features))
|
||||
logger.info(" Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
|
||||
logger.info(" Batch size per node = %d", per_node_train_batch_size)
|
||||
logger.info(" Total train batch size (w. parallel, distributed & accumulation) = %d", train_batch_size)
|
||||
logger.info(" Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
|
||||
logger.info(" Total optimization steps = %d", args.num_training_steps)
|
||||
|
||||
if args.num_training_steps <= global_step:
|
||||
logger.info("Training is done. Please use a new dir or clean this dir!")
|
||||
else:
|
||||
# The training features are shuffled
|
||||
train_sampler = SequentialSampler(train_dataset) \
|
||||
if args.local_rank == -1 else DistributedSampler(train_dataset, shuffle=False)
|
||||
train_dataloader = DataLoader(
|
||||
train_dataset, sampler=train_sampler,
|
||||
batch_size=per_node_train_batch_size // args.gradient_accumulation_steps,
|
||||
collate_fn=utils.batch_list_to_batch_tensors)
|
||||
|
||||
train_iterator = tqdm.tqdm(
|
||||
train_dataloader, initial=global_step,
|
||||
desc="Iter (loss=X.XXX, lr=X.XXXXXXX)", disable=args.local_rank not in [-1, 0])
|
||||
|
||||
model.train()
|
||||
model.zero_grad()
|
||||
|
||||
tr_loss, logging_loss = 0.0, 0.0
|
||||
|
||||
for step, batch in enumerate(train_iterator):
|
||||
batch = tuple(t.to(args.device) for t in batch)
|
||||
inputs = {'source_idxys': batch[0],
|
||||
'target_idxys': batch[1],
|
||||
'pseudo_idxys': batch[2],
|
||||
'num_source_tokens': batch[3],
|
||||
'num_target_tokens': batch[4],
|
||||
'target_index': batch[-1]}
|
||||
loss = model(**inputs)
|
||||
if args.n_gpu > 1:
|
||||
loss = loss.mean() # mean() to average on multi-gpu parallel (not distributed) training
|
||||
|
||||
train_iterator.set_description('Iter (loss=%5.3f) lr=%9.7f' % (loss.item(), scheduler.get_lr()[0]))
|
||||
|
||||
if args.gradient_accumulation_steps > 1:
|
||||
loss = loss / args.gradient_accumulation_steps
|
||||
|
||||
if args.fp16:
|
||||
with amp.scale_loss(loss, optimizer) as scaled_loss:
|
||||
scaled_loss.backward()
|
||||
else:
|
||||
loss.backward()
|
||||
|
||||
logging_loss += loss.item()
|
||||
if (step + 1) % args.gradient_accumulation_steps == 0:
|
||||
if args.fp16:
|
||||
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
|
||||
else:
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
|
||||
|
||||
optimizer.step()
|
||||
scheduler.step() # Update learning rate schedule
|
||||
model.zero_grad()
|
||||
global_step += 1
|
||||
|
||||
if args.local_rank in [-1, 0] and global_step % args.logging_steps == 0 and tb_writer is not None:
|
||||
logging_loss = 0.0
|
||||
tb_writer.add_scalar('train/lr', scheduler.get_lr()[0], global_step=global_step)
|
||||
tb_writer.add_scalar('train/loss', loss.item(), global_step=global_step)
|
||||
|
||||
if args.local_rank in [-1, 0] and args.save_steps > 0 and \
|
||||
(global_step % args.save_steps == 0 or global_step == args.num_training_steps):
|
||||
save_path = os.path.join(args.output_dir, "ckpt-%d" % global_step)
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
model_to_save = model.module if hasattr(model, "module") else model
|
||||
model_to_save.save_pretrained(save_path)
|
||||
|
||||
optim_to_save = {
|
||||
"optimizer": optimizer.state_dict(),
|
||||
"lr_scheduler": scheduler.state_dict(),
|
||||
}
|
||||
if args.fp16:
|
||||
optim_to_save["amp"] = amp.state_dict()
|
||||
torch.save(
|
||||
optim_to_save, os.path.join(args.output_dir, 'optim.{}.bin'.format(global_step)))
|
||||
|
||||
logger.info("Saving model checkpoint %d into %s", global_step, save_path)
|
||||
|
||||
if args.local_rank in [-1, 0] and tb_writer:
|
||||
tb_writer.close()
|
||||
|
||||
|
||||
def get_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument("--train_file", default=None, type=str,
|
||||
help="Training data (json format) for training. Keys: source and target")
|
||||
parser.add_argument("--train_folder", default=None, type=str,
|
||||
help="Training data folder for training. Keys: source and target")
|
||||
parser.add_argument("--sentence_shuffle_rate", default=0, type=float)
|
||||
parser.add_argument("--model_type", default=None, type=str, required=True,
|
||||
help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()))
|
||||
|
||||
parser.add_argument("--layoutlm_only_layout", action='store_true')
|
||||
parser.add_argument("--layout_only_dataset", action='store_true')
|
||||
|
||||
parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
|
||||
help="Path to pre-trained model or shortcut name selected in the list:")
|
||||
parser.add_argument("--output_dir", default=None, type=str, required=True,
|
||||
help="The output directory where the model checkpoints and predictions will be written.")
|
||||
parser.add_argument("--log_dir", default=None, type=str,
|
||||
help="The output directory where the log will be written.")
|
||||
|
||||
## Other parameters
|
||||
parser.add_argument("--config_name", default=None, type=str,
|
||||
help="Pretrained config name or path if not the same as model_name")
|
||||
parser.add_argument("--tokenizer_name", default=None, type=str,
|
||||
help="Pretrained tokenizer name or path if not the same as model_name")
|
||||
parser.add_argument("--cache_dir", default=None, type=str,
|
||||
help="Where do you want to store the pre-trained models downloaded from s3")
|
||||
|
||||
parser.add_argument("--max_source_seq_length", default=464, type=int,
|
||||
help="The maximum total source sequence length after WordPiece tokenization. Sequences "
|
||||
"longer than this will be truncated, and sequences shorter than this will be padded.")
|
||||
parser.add_argument("--max_target_seq_length", default=48, type=int,
|
||||
help="The maximum total target sequence length after WordPiece tokenization. Sequences "
|
||||
"longer than this will be truncated, and sequences shorter than this will be padded.")
|
||||
|
||||
parser.add_argument("--cached_train_features_file", default=None, type=str,
|
||||
help="Cached training features file")
|
||||
parser.add_argument("--do_lower_case", action='store_true',
|
||||
help="Set this flag if you are using an uncased model.")
|
||||
|
||||
parser.add_argument("--per_gpu_train_batch_size", default=8, type=int,
|
||||
help="Batch size per GPU/CPU for training.")
|
||||
parser.add_argument("--learning_rate", default=5e-5, type=float,
|
||||
help="The initial learning rate for Adam.")
|
||||
parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
|
||||
help="Number of updates steps to accumulate before performing a backward/update pass.")
|
||||
parser.add_argument("--weight_decay", default=0.01, type=float,
|
||||
help="Weight decay if we apply some.")
|
||||
parser.add_argument("--adam_epsilon", default=1e-8, type=float,
|
||||
help="Epsilon for Adam optimizer.")
|
||||
parser.add_argument("--max_grad_norm", default=1.0, type=float,
|
||||
help="Max gradient norm.")
|
||||
parser.add_argument("--label_smoothing", default=0.1, type=float,
|
||||
help="Label smoothing.")
|
||||
parser.add_argument("--num_training_steps", default=-1, type=int,
|
||||
help="set total number of training steps to perform")
|
||||
parser.add_argument("--num_training_epochs", default=10, type=int,
|
||||
help="set total number of training epochs to perform (--num_training_steps has higher priority)")
|
||||
parser.add_argument("--num_warmup_steps", default=0, type=int,
|
||||
help="Linear warmup over warmup_steps.")
|
||||
|
||||
parser.add_argument("--random_prob", default=0.1, type=float,
|
||||
help="prob to random replace a masked token")
|
||||
parser.add_argument("--keep_prob", default=0.1, type=float,
|
||||
help="prob to keep no change for a masked token")
|
||||
|
||||
parser.add_argument('--logging_steps', type=int, default=500,
|
||||
help="Log every X updates steps.")
|
||||
parser.add_argument('--save_steps', type=int, default=1500,
|
||||
help="Save checkpoint every X updates steps.")
|
||||
parser.add_argument("--no_cuda", action='store_true',
|
||||
help="Whether not to use CUDA when available")
|
||||
parser.add_argument('--seed', type=int, default=42,
|
||||
help="random seed for initialization")
|
||||
|
||||
parser.add_argument("--local_rank", type=int, default=-1,
|
||||
help="local_rank for distributed training on gpus")
|
||||
parser.add_argument('--fp16', action='store_true',
|
||||
help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
|
||||
parser.add_argument('--fp16_opt_level', type=str, default='O1',
|
||||
help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
|
||||
"See details at https://nvidia.github.io/apex/amp.html")
|
||||
parser.add_argument('--server_ip', type=str, default='', help="Can be used for distant debugging.")
|
||||
parser.add_argument('--server_port', type=str, default='', help="Can be used for distant debugging.")
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def prepare(args):
|
||||
# Setup distant debugging if needed
|
||||
if args.server_ip and args.server_port:
|
||||
# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
|
||||
import ptvsd
|
||||
print("Waiting for debugger attach")
|
||||
ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
|
||||
ptvsd.wait_for_attach()
|
||||
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
json.dump(args.__dict__, open(os.path.join(
|
||||
args.output_dir, 'train_opt.json'), 'w'), sort_keys=True, indent=2)
|
||||
|
||||
# Setup CUDA, GPU & distributed training
|
||||
if args.local_rank == -1 or args.no_cuda:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
|
||||
args.n_gpu = torch.cuda.device_count()
|
||||
else: # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
|
||||
torch.cuda.set_device(args.local_rank)
|
||||
device = torch.device("cuda", args.local_rank)
|
||||
torch.distributed.init_process_group(backend='nccl')
|
||||
args.n_gpu = 1
|
||||
args.device = device
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
|
||||
datefmt='%m/%d/%Y %H:%M:%S',
|
||||
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
|
||||
logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
|
||||
args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16)
|
||||
|
||||
# Set seed
|
||||
random.seed(args.seed)
|
||||
np.random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
if args.n_gpu > 0:
|
||||
torch.cuda.manual_seed_all(args.seed)
|
||||
|
||||
logger.info("Training/evaluation parameters %s", args)
|
||||
|
||||
# Before we do anything with models, we want to ensure that we get fp16 execution of torch.einsum if args.fp16 is set.
|
||||
# Otherwise it'll default to "promote" mode, and we'll get fp32 operations. Note that running `--fp16_opt_level="O2"` will
|
||||
# remove the need for this code, but it is still valid.
|
||||
if args.fp16:
|
||||
try:
|
||||
import apex
|
||||
apex.amp.register_half_function(torch, 'einsum')
|
||||
except ImportError:
|
||||
raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
|
||||
|
||||
|
||||
def get_model_and_tokenizer(args):
|
||||
config_class, tokenizer_class = MODEL_CLASSES[args.model_type]
|
||||
model_config = config_class.from_pretrained(
|
||||
args.config_name if args.config_name else args.model_name_or_path,
|
||||
cache_dir=args.cache_dir if args.cache_dir else None)
|
||||
config = BertForSeq2SeqConfig.from_exist_config(
|
||||
config=model_config, label_smoothing=args.label_smoothing,
|
||||
max_position_embeddings=args.max_source_seq_length + args.max_target_seq_length,
|
||||
max_source_length=args.max_source_seq_length,
|
||||
base_model_type=args.model_type,
|
||||
layoutlm_only_layout_flag=args.layoutlm_only_layout,
|
||||
)
|
||||
|
||||
logger.info("Model config for seq2seq: %s", str(config))
|
||||
|
||||
if args.model_type == 'layoutlm':
|
||||
if args.tokenizer_name is not None:
|
||||
tokenizer_name = args.tokenizer_name
|
||||
else:
|
||||
tokenizer_name = 'bert' + args.model_name_or_path[8:]
|
||||
tokenizer = tokenizer_class.from_pretrained(
|
||||
tokenizer_name, do_lower_case=args.do_lower_case, cache_dir=args.cache_dir if args.cache_dir else None)
|
||||
else:
|
||||
tokenizer = tokenizer_class.from_pretrained(
|
||||
args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
|
||||
do_lower_case=args.do_lower_case, cache_dir=args.cache_dir if args.cache_dir else None)
|
||||
|
||||
model = LayoutlmForSequenceToSequence.from_pretrained(
|
||||
args.model_name_or_path, config=config, model_type=args.model_type,
|
||||
reuse_position_embedding=True,
|
||||
cache_dir=args.cache_dir if args.cache_dir else None,
|
||||
)
|
||||
|
||||
return model, tokenizer
|
||||
|
||||
|
||||
def main():
|
||||
args = get_args()
|
||||
prepare(args)
|
||||
|
||||
if args.local_rank not in [-1, 0]:
|
||||
torch.distributed.barrier()
|
||||
# Make sure only the first process in distributed training will download model & vocab
|
||||
# Load pretrained model and tokenizer
|
||||
model, tokenizer = get_model_and_tokenizer(args)
|
||||
|
||||
if args.local_rank == 0:
|
||||
torch.distributed.barrier()
|
||||
# Make sure only the first process in distributed training will download model & vocab
|
||||
|
||||
if args.cached_train_features_file is None:
|
||||
args.cached_train_features_file = os.path.join(args.output_dir, "cached_features_for_training.pt")
|
||||
|
||||
example_path = args.train_file if args.train_file else args.train_folder
|
||||
if args.layout_only_dataset:
|
||||
training_features = utils.load_and_cache_line_order_examples(
|
||||
example_path=example_path, tokenizer=tokenizer, local_rank=args.local_rank,
|
||||
cached_features_file=args.cached_train_features_file, max_src_length=args.max_source_seq_length,
|
||||
layout_flag=args.model_type == 'layoutlm', shuffle=True,
|
||||
src_shuffle_rate=args.sentence_shuffle_rate)
|
||||
else:
|
||||
training_features = utils.load_and_cache_layoutlm_examples(
|
||||
example_path=example_path, tokenizer=tokenizer, local_rank=args.local_rank,
|
||||
cached_features_file=args.cached_train_features_file, max_src_length=args.max_source_seq_length,
|
||||
layout_flag=args.model_type == 'layoutlm', shuffle=True,
|
||||
src_shuffle_rate=args.sentence_shuffle_rate
|
||||
)
|
||||
|
||||
train(args, training_features, model, tokenizer)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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)
|
||||
@@ -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)")
|
||||
@@ -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)")
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
from io import open
|
||||
from setuptools import find_packages, setup
|
||||
|
||||
|
||||
extras = {
|
||||
'serving': ['pydantic', 'uvicorn', 'fastapi'],
|
||||
'serving-tf': ['pydantic', 'uvicorn', 'fastapi'],
|
||||
'serving-torch': ['pydantic', 'uvicorn', 'fastapi', 'torch']
|
||||
}
|
||||
extras['all'] = [package for package in extras.values()]
|
||||
|
||||
setup(
|
||||
name="s2s-ft",
|
||||
version="0.0.1",
|
||||
author="UniLM Team",
|
||||
author_email="unilm@microsoft.com",
|
||||
description="Fine-Tuning Bidirectional Transformers for Sequence-to-Sequence Learning",
|
||||
long_description=open("README.md", "r", encoding='utf-8').read(),
|
||||
long_description_content_type="text/markdown",
|
||||
keywords='Fine-Tuning Bidirectional Transformers for Sequence-to-Sequence Learning',
|
||||
license='Apache',
|
||||
url="https://github.com/microsoft/unilm/tree/master/s2s-ft",
|
||||
packages=find_packages(exclude=["*.tests", "*.tests.*",
|
||||
"tests.*", "tests"]),
|
||||
install_requires=['numpy',
|
||||
'boto3',
|
||||
'requests',
|
||||
'tqdm',
|
||||
'regex != 2019.12.17',
|
||||
'sentencepiece',
|
||||
'sacremoses',
|
||||
'tensorboardX',
|
||||
'transformers <= 2.10.0'],
|
||||
extras_require=extras,
|
||||
python_requires='>=3.5.0',
|
||||
classifiers=[
|
||||
'Programming Language :: Python :: 3',
|
||||
],
|
||||
)
|
||||
Reference in New Issue
Block a user