chore: import upstream snapshot with attribution

xdoc/README.md

@@ -0,0 +1,232 @@
+# XDoc
+
+## Introduction
+
+XDoc is a unified pre-trained model that deals with different document formats in a single model. With only 36.7% parameters, XDoc achieves comparable or better performance on downstream tasks, which is cost-effective for real-world deployment.
+
+[XDoc: Unified Pre-training for Cross-Format Document Understanding](https://arxiv.org/abs/2210.02849)
+Jingye Chen, Tengchao Lv, Lei Cui, Cha Zhang, Furu Wei, [EMNLP 2022](#)
+
+The overview of our framework is as follows:
+
+<div align="center">
+<img src="./architecture.png" width="100%" height="100%" />
+</div>
+
+## Download
+
+
+### Pre-trained Model
+|   Model  | Download |
+| -------- | -------- |
+| xdoc-pretrain-roberta-1M    | [xdoc-base](https://huggingface.co/microsoft/xdoc-base) |
+
+### Fine-tuning Models
+|   Model  | Download |
+| -------- | -------- |
+| xdoc-squad1.1    | [xdoc-squad1.1](https://huggingface.co/microsoft/xdoc-base-squad1.1) |
+| xdoc-squad2.0    | [xdoc-squad2.0](https://huggingface.co/microsoft/xdoc-base-squad2.0) |
+| xdoc-funsd    | [xdoc-funsd](https://huggingface.co/microsoft/xdoc-base-funsd) |
+| xdoc-websrc | [xdoc-websrc](https://huggingface.co/microsoft/xdoc-base-websrc) |
+
+
+
+## Fine-tune
+
+
+### SQuAD
+The dataset will be **automatically downloaded**. Please refer to ```./fine_tuning/squad/```.
+
+#### Installation
+```
+pip install -r requirements.txt
+```
+
+#### Train
+To train XDoc on SQuADv1.1
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python run_squad.py \
+  --model_name_or_path microsoft/xdoc-base \
+  --dataset_name squad \
+  --do_train \
+  --do_eval \
+  --per_device_train_batch_size 16 \
+  --learning_rate 3e-5 \
+  --num_train_epochs 2 \
+  --max_seq_length 384 \
+  --doc_stride 128 \
+  --output_dir ./v1_result \
+  --overwrite_output_dir
+```
+
+To train XDoc on SQuADv2.0
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python run_squad.py \
+  --model_name_or_path microsoft/xdoc-base \
+  --dataset_name squad_v2 \
+  --do_train \
+  --do_eval \
+  --version_2_with_negative \
+  --per_device_train_batch_size 16 \
+  --learning_rate 3e-5 \
+  --num_train_epochs 4 \
+  --max_seq_length 384 \
+  --doc_stride 128 \
+  --output_dir ./v2_result \
+  --overwrite_output_dir
+```
+
+#### Test
+To test XDoc on SQuADv1.1
+
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python run_squad.py \
+  --model_name_or_path microsoft/xdoc-base-squad1.1 \
+  --dataset_name squad \
+  --do_eval \
+  --per_device_train_batch_size 16 \
+  --learning_rate 3e-5 \
+  --num_train_epochs 2 \
+  --max_seq_length 384 \
+  --doc_stride 128 \
+  --output_dir ./squadv1.1_result \
+  --overwrite_output_dir
+```
+
+To test XDoc on SQuADv2.0
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python run_squad.py \
+  --model_name_or_path microsoft/xdoc-base-squad2.0 \
+  --dataset_name squad_v2 \
+  --do_eval \
+  --version_2_with_negative \
+  --per_device_train_batch_size 16 \
+  --learning_rate 3e-5 \
+  --num_train_epochs 4 \
+  --max_seq_length 384 \
+  --doc_stride 128 \
+  --output_dir ./squadv2.0_result \
+  --overwrite_output_dir
+```
+
+
+
+### FUNSD
+The dataset will be **automatically downloaded**. Please refer to ```./fine_tuning/funsd/```.
+
+#### Installation
+
+```bash
+pip install -r requirements.txt
+```
+
+Also, you need to install ```detectron2```. For example, if you use torch1.8 with cuda version 10.1, you can use the following command
+
+```bash
+pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.8/index.html
+```
+
+#### Train
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 --master_port 5678 run_funsd.py \
+        --model_name_or_path microsoft/xdoc-base \
+        --output_dir camera_ready_funsd_1M \
+        --do_train \
+        --do_eval \
+        --max_steps 1000 \
+        --warmup_ratio 0.1 \
+        --fp16 \
+        --overwrite_output_dir \
+        --seed 42
+```
+
+#### Test
+
+```
+CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 --master_port 5678 run_funsd.py \
+        --model_name_or_path microsoft/xdoc-base-funsd \
+        --output_dir camera_ready_funsd_1M \
+        --do_eval \
+        --max_steps 1000 \
+        --warmup_ratio 0.1 \
+        --fp16 \
+        --overwrite_output_dir \
+        --seed 42
+```
+
+### WebSRC
+The dataset will be **manually downloaded**. After downloading, please modify the argument ```--web_train_file```, ```--web_eval_file```, ```web_root_dir```, and ```root_dir``` in args.py.
+
+#### Installation
+
+```bash
+pip install -r requirements.txt
+```
+
+#### Train
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python run_docvqa.py --do_train True --do_eval True  --model_name_or_path microsoft/xdoc-base
+```
+
+#### Test
+```bash
+CUDA_VISIBLE_DEVICES=0 python run_docvqa.py --do_train False --do_eval True --model_name_or_path microsoft/xdoc-base-websrc
+```
+
+
+## Result
+
+* To verify the model accuracy, we select the GLUE benchmark and SQuAD to evaluate plain text understanding, FUNSD and DocVQA to evaluate doc-
+ument understanding, and WebSRC for web text understanding. Experimental
+results have demonstrated that XDoc achieves comparable or even better performance on these tasks.
+
+| Model      | MNLI-m | QNLI | SST2 | MRPC | SQUAD1.1/2.0 | FUNSD | DocVQA | WebSRC |
+| :----------: | :------: | :----: | :----: | :----: | :------------: | :-----: | :------: | :------: |
+| RoBERTa    | **87.6**   | **92.8** | 94.8 | 90.2 | **92.2**/83.4    | -     | -      | -      |
+| LayoutLM   | -      | -    | -    | -    | -            | 79.3  | 69.2   | -      |
+| MarkupLM   | -      | -    | -    | -    | -            | -     | -      | 74.5   |
+| **XDoc(Ours)** | 86.8   | 92.3 | **95.3** | **91.1** | 92.0/**83.5**    | **89.4**  | **72.7**   | **74.8**   |
+
+* With only 36.7% parameters, XDoc achieves comparable or even better performance on a variety of downstream tasks compared with the individual pre-trained models, which is cost effective for real-world deployment.
+
+| Model      | Word | 1D Position | Transformer | 2D Position | XPath | Adaptive | Total |
+| :----------: | :----: | :-----------: | :-----------: | :-----------: | :-----: | :--------: | :-----: |
+| RoBERTa    | √  | √         | √         | -           | -     | -        | 128M  |
+| LayoutLM   | √  | √         | √         | √         | -     | -        | 131M  |
+| MarkupLM   | √  | √         | √         | -           | √   | -        | 139M  |
+| **XDoc(Ours)** | √  | √         | √         | √         | √   | √      | 146M  |
+
+
+
+## Citation
+
+If you find XDoc helpful, please cite us:
+```
+@article{chen2022xdoc,
+  title={XDoc: Unified Pre-training for Cross-Format Document Understanding},
+  author={Chen, Jingye and Lv, Tengchao and Cui, Lei and Zhang, Cha and Wei, Furu},
+  journal={arXiv preprint arXiv:2210.02849},
+  year={2022}
+}
+```
+
+
+## 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).
+[Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct)
+
+## Contact
+
+For help or issues using XDoc, please submit a GitHub issue.
+
+For other communications, please contact Lei Cui (`lecu@microsoft.com`), Furu Wei (`fuwei@microsoft.com`).
+
+

xdoc/fine_tuning/README.md

@@ -0,0 +1,3 @@
+# Fine-tuning code of XDoc
+
+Please follow the ```requirements.txt``` in each folder to install the corresponding packages. In addition, please note that different downstream tasks may use different versions of ```transformers```.

xdoc/fine_tuning/funsd/layoutlmft/__init__.py

@@ -0,0 +1,46 @@
+from collections import OrderedDict
+
+from transformers import CONFIG_MAPPING, MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING, MODEL_NAMES_MAPPING, TOKENIZER_MAPPING
+from transformers.convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS, BertConverter, XLMRobertaConverter
+from transformers.models.auto.modeling_auto import auto_class_factory
+
+from .models.layoutlmv2 import (
+    LayoutLMv2Config,
+    LayoutLMv2ForRelationExtraction,
+    LayoutLMv2ForTokenClassification,
+    LayoutLMv2Tokenizer,
+    LayoutLMv2TokenizerFast,
+)
+from .models.layoutxlm import (
+    LayoutXLMConfig,
+    LayoutXLMForRelationExtraction,
+    LayoutXLMForTokenClassification,
+    LayoutXLMTokenizer,
+    LayoutXLMTokenizerFast,
+)
+
+
+CONFIG_MAPPING.update([("layoutlmv2", LayoutLMv2Config), ("layoutxlm", LayoutXLMConfig)])
+MODEL_NAMES_MAPPING.update([("layoutlmv2", "LayoutLMv2"), ("layoutxlm", "LayoutXLM")])
+TOKENIZER_MAPPING.update(
+    [
+        (LayoutLMv2Config, (LayoutLMv2Tokenizer, LayoutLMv2TokenizerFast)),
+        (LayoutXLMConfig, (LayoutXLMTokenizer, LayoutXLMTokenizerFast)),
+    ]
+)
+SLOW_TO_FAST_CONVERTERS.update({"LayoutLMv2Tokenizer": BertConverter, "LayoutXLMTokenizer": XLMRobertaConverter})
+MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.update(
+    [(LayoutLMv2Config, LayoutLMv2ForTokenClassification), (LayoutXLMConfig, LayoutXLMForTokenClassification)]
+)
+
+MODEL_FOR_RELATION_EXTRACTION_MAPPING = OrderedDict(
+    [(LayoutLMv2Config, LayoutLMv2ForRelationExtraction), (LayoutXLMConfig, LayoutXLMForRelationExtraction)]
+)
+
+AutoModelForTokenClassification = auto_class_factory(
+    "AutoModelForTokenClassification", MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING, head_doc="token classification"
+)
+
+AutoModelForRelationExtraction = auto_class_factory(
+    "AutoModelForRelationExtraction", MODEL_FOR_RELATION_EXTRACTION_MAPPING, head_doc="relation extraction"
+)

xdoc/fine_tuning/funsd/layoutlmft/data/__init__.py

@@ -0,0 +1,3 @@
+# flake8: noqa
+from .data_collator import DataCollatorForKeyValueExtraction
+from .datasets import *

xdoc/fine_tuning/funsd/layoutlmft/data/data_args.py

@@ -0,0 +1,81 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    task_name: Optional[str] = field(default="ner", metadata={"help": "The name of the task (ner, pos...)."})
+    dataset_name: Optional[str] = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    train_file: Optional[str] = field(
+        default=None, metadata={"help": "The input training data file (a csv or JSON file)."}
+    )
+    validation_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input evaluation data file to evaluate on (a csv or JSON file)."},
+    )
+    test_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input test data file to predict on (a csv or JSON file)."},
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    pad_to_max_length: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to pad all samples to model maximum sentence length. "
+            "If False, will pad the samples dynamically when batching to the maximum length in the batch. More "
+            "efficient on GPU but very bad for TPU."
+        },
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
+            "value if set."
+        },
+    )
+    max_val_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of validation examples to this "
+            "value if set."
+        },
+    )
+    max_test_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of test examples to this "
+            "value if set."
+        },
+    )
+    label_all_tokens: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to put the label for one word on all tokens of generated by that word or just on the "
+            "one (in which case the other tokens will have a padding index)."
+        },
+    )
+    return_entity_level_metrics: bool = field(
+        default=False,
+        metadata={"help": "Whether to return all the entity levels during evaluation or just the overall ones."},
+    )
+
+
+@dataclass
+class XFUNDataTrainingArguments(DataTrainingArguments):
+    lang: Optional[str] = field(default="en")
+    additional_langs: Optional[str] = field(default=None)

xdoc/fine_tuning/funsd/layoutlmft/data/data_collator.py

@@ -0,0 +1,82 @@
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+
+from detectron2.structures import ImageList
+from transformers import PreTrainedTokenizerBase
+from transformers.file_utils import PaddingStrategy
+
+
+@dataclass
+class DataCollatorForKeyValueExtraction:
+    """
+    Data collator that will dynamically pad the inputs received, as well as the labels.
+
+    Args:
+        tokenizer (:class:`~transformers.PreTrainedTokenizer` or :class:`~transformers.PreTrainedTokenizerFast`):
+            The tokenizer used for encoding the data.
+        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.file_utils.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
+            among:
+
+            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+              sequence if provided).
+            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
+              maximum acceptable input length for the model if that argument is not provided.
+            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
+              different lengths).
+        max_length (:obj:`int`, `optional`):
+            Maximum length of the returned list and optionally padding length (see above).
+        pad_to_multiple_of (:obj:`int`, `optional`):
+            If set will pad the sequence to a multiple of the provided value.
+
+            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
+            7.5 (Volta).
+        label_pad_token_id (:obj:`int`, `optional`, defaults to -100):
+            The id to use when padding the labels (-100 will be automatically ignore by PyTorch loss functions).
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    label_pad_token_id: int = -100
+
+    def __call__(self, features):
+        label_name = "label" if "label" in features[0].keys() else "labels"
+        labels = [feature[label_name] for feature in features] if label_name in features[0].keys() else None
+
+        has_image_input = "image" in features[0]
+        has_bbox_input = "bbox" in features[0]
+        if has_image_input:
+            image = ImageList.from_tensors([torch.tensor(feature["image"]) for feature in features], 32)
+            for feature in features:
+                del feature["image"]
+        batch = self.tokenizer.pad(
+            features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            # Conversion to tensors will fail if we have labels as they are not of the same length yet.
+            return_tensors="pt" if labels is None else None,
+        )
+
+        if labels is None:
+            return batch
+
+        sequence_length = torch.tensor(batch["input_ids"]).shape[1]
+        padding_side = self.tokenizer.padding_side
+        if padding_side == "right":
+            batch["labels"] = [label + [self.label_pad_token_id] * (sequence_length - len(label)) for label in labels]
+            if has_bbox_input:
+                batch["bbox"] = [bbox + [[0, 0, 0, 0]] * (sequence_length - len(bbox)) for bbox in batch["bbox"]]
+        else:
+            batch["labels"] = [[self.label_pad_token_id] * (sequence_length - len(label)) + label for label in labels]
+            if has_bbox_input:
+                batch["bbox"] = [[[0, 0, 0, 0]] * (sequence_length - len(bbox)) + bbox for bbox in batch["bbox"]]
+
+        batch = {k: torch.tensor(v, dtype=torch.int64) if isinstance(v[0], list) else v for k, v in batch.items()}
+        if has_image_input:
+            batch["image"] = image
+        return batch

xdoc/fine_tuning/funsd/layoutlmft/data/datasets/funsd.py

@@ -0,0 +1,122 @@
+# coding=utf-8
+
+import json
+import os
+
+import datasets
+
+from layoutlmft.data.utils import load_image, normalize_bbox
+
+
+logger = datasets.logging.get_logger(__name__)
+
+
+_CITATION = """\
+@article{Jaume2019FUNSDAD,
+  title={FUNSD: A Dataset for Form Understanding in Noisy Scanned Documents},
+  author={Guillaume Jaume and H. K. Ekenel and J. Thiran},
+  journal={2019 International Conference on Document Analysis and Recognition Workshops (ICDARW)},
+  year={2019},
+  volume={2},
+  pages={1-6}
+}
+"""
+
+_DESCRIPTION = """\
+https://guillaumejaume.github.io/FUNSD/
+"""
+
+
+class FunsdConfig(datasets.BuilderConfig):
+    """BuilderConfig for FUNSD"""
+
+    def __init__(self, **kwargs):
+        """BuilderConfig for FUNSD.
+
+        Args:
+          **kwargs: keyword arguments forwarded to super.
+        """
+        super(FunsdConfig, self).__init__(**kwargs)
+
+
+class Funsd(datasets.GeneratorBasedBuilder):
+    """Conll2003 dataset."""
+
+    BUILDER_CONFIGS = [
+        FunsdConfig(name="funsd", version=datasets.Version("1.0.0"), description="FUNSD dataset"),
+    ]
+
+    def _info(self):
+        return datasets.DatasetInfo(
+            description=_DESCRIPTION,
+            features=datasets.Features(
+                {
+                    "id": datasets.Value("string"),
+                    "tokens": datasets.Sequence(datasets.Value("string")),
+                    "bboxes": datasets.Sequence(datasets.Sequence(datasets.Value("int64"))),
+                    "ner_tags": datasets.Sequence(
+                        datasets.features.ClassLabel(
+                            names=["O", "B-HEADER", "I-HEADER", "B-QUESTION", "I-QUESTION", "B-ANSWER", "I-ANSWER"]
+                        )
+                    ),
+                    "image": datasets.Array3D(shape=(3, 224, 224), dtype="uint8"),
+                }
+            ),
+            supervised_keys=None,
+            homepage="https://guillaumejaume.github.io/FUNSD/",
+            citation=_CITATION,
+        )
+
+    def _split_generators(self, dl_manager):
+        """Returns SplitGenerators."""
+        downloaded_file = dl_manager.download_and_extract("https://guillaumejaume.github.io/FUNSD/dataset.zip")
+        return [
+            datasets.SplitGenerator(
+                name=datasets.Split.TRAIN, gen_kwargs={"filepath": f"{downloaded_file}/dataset/training_data/"}
+            ),
+            datasets.SplitGenerator(
+                name=datasets.Split.TEST, gen_kwargs={"filepath": f"{downloaded_file}/dataset/testing_data/"}
+            ),
+        ]
+
+    def _generate_examples(self, filepath):
+
+        logger.info("⏳ Generating examples from = %s", filepath)
+        ann_dir = os.path.join(filepath, "annotations")
+        img_dir = os.path.join(filepath, "images")
+        for guid, file in enumerate(sorted(os.listdir(ann_dir))):
+            tokens = []
+            bboxes = []
+            ner_tags = []
+
+            file_path = os.path.join(ann_dir, file)
+            with open(file_path, "r", encoding="utf8") as f:
+                data = json.load(f)
+            image_path = os.path.join(img_dir, file)
+            image_path = image_path.replace("json", "png")
+            image, size = load_image(image_path)
+            for item in data["form"]:
+                words, label = item["words"], item["label"]
+                whole_box = item["box"]
+
+                words = [w for w in words if w["text"].strip() != ""]
+                if len(words) == 0:
+                    continue
+                if label == "other":
+                    for w in words:
+                        tokens.append(w["text"])
+                        ner_tags.append("O")
+                        # bboxes.append(normalize_bbox(w["box"], size))
+                        bboxes.append(normalize_bbox(whole_box, size))
+                else:
+                    tokens.append(words[0]["text"])
+                    ner_tags.append("B-" + label.upper())
+                    # bboxes.append(normalize_bbox(words[0]["box"], size))
+                    bboxes.append(normalize_bbox(whole_box, size))
+                    for w in words[1:]:
+                        tokens.append(w["text"])
+                        ner_tags.append("I-" + label.upper())
+                        # bboxes.append(normalize_bbox(w["box"], size))
+                        bboxes.append(normalize_bbox(whole_box, size))
+
+            yield guid, {"id": str(guid), "tokens": tokens, "bboxes": bboxes, "ner_tags": ner_tags, "image": image}

xdoc/fine_tuning/funsd/layoutlmft/data/datasets/xfun.py

@@ -0,0 +1,245 @@
+# Lint as: python3
+import json
+import logging
+import os
+
+import datasets
+
+from layoutlmft.data.utils import load_image, merge_bbox, normalize_bbox, simplify_bbox
+from transformers import AutoTokenizer
+
+
+_URL = "https://github.com/doc-analysis/XFUN/releases/download/v1.0/"
+
+_LANG = ["zh", "de", "es", "fr", "en", "it", "ja", "pt"]
+logger = logging.getLogger(__name__)
+
+
+class XFUNConfig(datasets.BuilderConfig):
+    """BuilderConfig for XFUN."""
+
+    def __init__(self, lang, additional_langs=None, **kwargs):
+        """
+        Args:
+            lang: string, language for the input text
+            **kwargs: keyword arguments forwarded to super.
+        """
+        super(XFUNConfig, self).__init__(**kwargs)
+        self.lang = lang
+        self.additional_langs = additional_langs
+
+
+class XFUN(datasets.GeneratorBasedBuilder):
+    """XFUN dataset."""
+
+    BUILDER_CONFIGS = [XFUNConfig(name=f"xfun.{lang}", lang=lang) for lang in _LANG]
+
+    tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
+
+    def _info(self):
+        return datasets.DatasetInfo(
+            features=datasets.Features(
+                {
+                    "id": datasets.Value("string"),
+                    "input_ids": datasets.Sequence(datasets.Value("int64")),
+                    "bbox": datasets.Sequence(datasets.Sequence(datasets.Value("int64"))),
+                    "labels": datasets.Sequence(
+                        datasets.ClassLabel(
+                            names=["O", "B-QUESTION", "B-ANSWER", "B-HEADER", "I-ANSWER", "I-QUESTION", "I-HEADER"]
+                        )
+                    ),
+                    "image": datasets.Array3D(shape=(3, 224, 224), dtype="uint8"),
+                    "entities": datasets.Sequence(
+                        {
+                            "start": datasets.Value("int64"),
+                            "end": datasets.Value("int64"),
+                            "label": datasets.ClassLabel(names=["HEADER", "QUESTION", "ANSWER"]),
+                        }
+                    ),
+                    "relations": datasets.Sequence(
+                        {
+                            "head": datasets.Value("int64"),
+                            "tail": datasets.Value("int64"),
+                            "start_index": datasets.Value("int64"),
+                            "end_index": datasets.Value("int64"),
+                        }
+                    ),
+                }
+            ),
+            supervised_keys=None,
+        )
+
+    def _split_generators(self, dl_manager):
+        """Returns SplitGenerators."""
+        urls_to_download = {
+            "train": [f"{_URL}{self.config.lang}.train.json", f"{_URL}{self.config.lang}.train.zip"],
+            "val": [f"{_URL}{self.config.lang}.val.json", f"{_URL}{self.config.lang}.val.zip"],
+            # "test": [f"{_URL}{self.config.lang}.test.json", f"{_URL}{self.config.lang}.test.zip"],
+        }
+        downloaded_files = dl_manager.download_and_extract(urls_to_download)
+        train_files_for_many_langs = [downloaded_files["train"]]
+        val_files_for_many_langs = [downloaded_files["val"]]
+        # test_files_for_many_langs = [downloaded_files["test"]]
+        if self.config.additional_langs:
+            additional_langs = self.config.additional_langs.split("+")
+            if "all" in additional_langs:
+                additional_langs = [lang for lang in _LANG if lang != self.config.lang]
+            for lang in additional_langs:
+                urls_to_download = {"train": [f"{_URL}{lang}.train.json", f"{_URL}{lang}.train.zip"]}
+                additional_downloaded_files = dl_manager.download_and_extract(urls_to_download)
+                train_files_for_many_langs.append(additional_downloaded_files["train"])
+
+        logger.info(f"Training on {self.config.lang} with additional langs({self.config.additional_langs})")
+        logger.info(f"Evaluating on {self.config.lang}")
+        logger.info(f"Testing on {self.config.lang}")
+        return [
+            datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={"filepaths": train_files_for_many_langs}),
+            datasets.SplitGenerator(
+                name=datasets.Split.VALIDATION, gen_kwargs={"filepaths": val_files_for_many_langs}
+            ),
+            # datasets.SplitGenerator(name=datasets.Split.TEST, gen_kwargs={"filepaths": test_files_for_many_langs}),
+        ]
+
+    def _generate_examples(self, filepaths):
+        for filepath in filepaths:
+            logger.info("Generating examples from = %s", filepath)
+            with open(filepath[0], "r", encoding="utf-8") as f:
+                data = json.load(f)
+
+            for doc in data["documents"]:
+                doc["img"]["fpath"] = os.path.join(filepath[1], doc["img"]["fname"])
+                image, size = load_image(doc["img"]["fpath"])
+                document = doc["document"]
+                tokenized_doc = {"input_ids": [], "bbox": [], "labels": []}
+                entities = []
+                relations = []
+                id2label = {}
+                entity_id_to_index_map = {}
+                empty_entity = set()
+                for line in document:
+                    if len(line["text"]) == 0:
+                        empty_entity.add(line["id"])
+                        continue
+                    id2label[line["id"]] = line["label"]
+                    relations.extend([tuple(sorted(l)) for l in line["linking"]])
+                    tokenized_inputs = self.tokenizer(
+                        line["text"],
+                        add_special_tokens=False,
+                        return_offsets_mapping=True,
+                        return_attention_mask=False,
+                    )
+                    text_length = 0
+                    ocr_length = 0
+                    bbox = []
+                    last_box = None
+                    for token_id, offset in zip(tokenized_inputs["input_ids"], tokenized_inputs["offset_mapping"]):
+                        if token_id == 6:
+                            bbox.append(None)
+                            continue
+                        text_length += offset[1] - offset[0]
+                        tmp_box = []
+                        while ocr_length < text_length:
+                            ocr_word = line["words"].pop(0)
+                            ocr_length += len(
+                                self.tokenizer._tokenizer.normalizer.normalize_str(ocr_word["text"].strip())
+                            )
+                            tmp_box.append(simplify_bbox(ocr_word["box"]))
+                        if len(tmp_box) == 0:
+                            tmp_box = last_box
+                        bbox.append(normalize_bbox(merge_bbox(tmp_box), size))
+                        last_box = tmp_box
+                    bbox = [
+                        [bbox[i + 1][0], bbox[i + 1][1], bbox[i + 1][0], bbox[i + 1][1]] if b is None else b
+                        for i, b in enumerate(bbox)
+                    ]
+                    if line["label"] == "other":
+                        label = ["O"] * len(bbox)
+                    else:
+                        label = [f"I-{line['label'].upper()}"] * len(bbox)
+                        label[0] = f"B-{line['label'].upper()}"
+                    tokenized_inputs.update({"bbox": bbox, "labels": label})
+                    if label[0] != "O":
+                        entity_id_to_index_map[line["id"]] = len(entities)
+                        entities.append(
+                            {
+                                "start": len(tokenized_doc["input_ids"]),
+                                "end": len(tokenized_doc["input_ids"]) + len(tokenized_inputs["input_ids"]),
+                                "label": line["label"].upper(),
+                            }
+                        )
+                    for i in tokenized_doc:
+                        tokenized_doc[i] = tokenized_doc[i] + tokenized_inputs[i]
+                relations = list(set(relations))
+                relations = [rel for rel in relations if rel[0] not in empty_entity and rel[1] not in empty_entity]
+                kvrelations = []
+                for rel in relations:
+                    pair = [id2label[rel[0]], id2label[rel[1]]]
+                    if pair == ["question", "answer"]:
+                        kvrelations.append(
+                            {"head": entity_id_to_index_map[rel[0]], "tail": entity_id_to_index_map[rel[1]]}
+                        )
+                    elif pair == ["answer", "question"]:
+                        kvrelations.append(
+                            {"head": entity_id_to_index_map[rel[1]], "tail": entity_id_to_index_map[rel[0]]}
+                        )
+                    else:
+                        continue
+
+                def get_relation_span(rel):
+                    bound = []
+                    for entity_index in [rel["head"], rel["tail"]]:
+                        bound.append(entities[entity_index]["start"])
+                        bound.append(entities[entity_index]["end"])
+                    return min(bound), max(bound)
+
+                relations = sorted(
+                    [
+                        {
+                            "head": rel["head"],
+                            "tail": rel["tail"],
+                            "start_index": get_relation_span(rel)[0],
+                            "end_index": get_relation_span(rel)[1],
+                        }
+                        for rel in kvrelations
+                    ],
+                    key=lambda x: x["head"],
+                )
+                chunk_size = 512
+                for chunk_id, index in enumerate(range(0, len(tokenized_doc["input_ids"]), chunk_size)):
+                    item = {}
+                    for k in tokenized_doc:
+                        item[k] = tokenized_doc[k][index : index + chunk_size]
+                    entities_in_this_span = []
+                    global_to_local_map = {}
+                    for entity_id, entity in enumerate(entities):
+                        if (
+                            index <= entity["start"] < index + chunk_size
+                            and index <= entity["end"] < index + chunk_size
+                        ):
+                            entity["start"] = entity["start"] - index
+                            entity["end"] = entity["end"] - index
+                            global_to_local_map[entity_id] = len(entities_in_this_span)
+                            entities_in_this_span.append(entity)
+                    relations_in_this_span = []
+                    for relation in relations:
+                        if (
+                            index <= relation["start_index"] < index + chunk_size
+                            and index <= relation["end_index"] < index + chunk_size
+                        ):
+                            relations_in_this_span.append(
+                                {
+                                    "head": global_to_local_map[relation["head"]],
+                                    "tail": global_to_local_map[relation["tail"]],
+                                    "start_index": relation["start_index"] - index,
+                                    "end_index": relation["end_index"] - index,
+                                }
+                            )
+                    item.update(
+                        {
+                            "id": f"{doc['id']}_{chunk_id}",
+                            "image": image,
+                            "entities": entities_in_this_span,
+                            "relations": relations_in_this_span,
+                        }
+                    )
+                    yield f"{doc['id']}_{chunk_id}", item

xdoc/fine_tuning/funsd/layoutlmft/data/utils.py

@@ -0,0 +1,36 @@
+import torch
+
+from detectron2.data.detection_utils import read_image
+from detectron2.data.transforms import ResizeTransform, TransformList
+
+
+def normalize_bbox(bbox, size):
+    return [
+        int(1000 * bbox[0] / size[0]),
+        int(1000 * bbox[1] / size[1]),
+        int(1000 * bbox[2] / size[0]),
+        int(1000 * bbox[3] / size[1]),
+    ]
+
+
+def simplify_bbox(bbox):
+    return [
+        min(bbox[0::2]),
+        min(bbox[1::2]),
+        max(bbox[2::2]),
+        max(bbox[3::2]),
+    ]
+
+
+def merge_bbox(bbox_list):
+    x0, y0, x1, y1 = list(zip(*bbox_list))
+    return [min(x0), min(y0), max(x1), max(y1)]
+
+
+def load_image(image_path):
+    image = read_image(image_path, format="BGR")
+    h = image.shape[0]
+    w = image.shape[1]
+    img_trans = TransformList([ResizeTransform(h=h, w=w, new_h=224, new_w=224)])
+    image = torch.tensor(img_trans.apply_image(image).copy()).permute(2, 0, 1)  # copy to make it writeable
+    return image, (w, h)

xdoc/fine_tuning/funsd/layoutlmft/evaluation.py

@@ -0,0 +1,150 @@
+import os
+import re
+
+import numpy as np
+
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+PREFIX_CHECKPOINT_DIR = "checkpoint"
+_re_checkpoint = re.compile(r"^" + PREFIX_CHECKPOINT_DIR + r"\-(\d+)$")
+
+
+def get_last_checkpoint(folder):
+    content = os.listdir(folder)
+    checkpoints = [
+        path
+        for path in content
+        if _re_checkpoint.search(path) is not None and os.path.isdir(os.path.join(folder, path))
+    ]
+    if len(checkpoints) == 0:
+        return
+    return os.path.join(folder, max(checkpoints, key=lambda x: int(_re_checkpoint.search(x).groups()[0])))
+
+
+def re_score(pred_relations, gt_relations, mode="strict"):
+    """Evaluate RE predictions
+
+    Args:
+        pred_relations (list) :  list of list of predicted relations (several relations in each sentence)
+        gt_relations (list) :    list of list of ground truth relations
+
+            rel = { "head": (start_idx (inclusive), end_idx (exclusive)),
+                    "tail": (start_idx (inclusive), end_idx (exclusive)),
+                    "head_type": ent_type,
+                    "tail_type": ent_type,
+                    "type": rel_type}
+
+        vocab (Vocab) :         dataset vocabulary
+        mode (str) :            in 'strict' or 'boundaries'"""
+
+    assert mode in ["strict", "boundaries"]
+
+    relation_types = [v for v in [0, 1] if not v == 0]
+    scores = {rel: {"tp": 0, "fp": 0, "fn": 0} for rel in relation_types + ["ALL"]}
+
+    # Count GT relations and Predicted relations
+    n_sents = len(gt_relations)
+    n_rels = sum([len([rel for rel in sent]) for sent in gt_relations])
+    n_found = sum([len([rel for rel in sent]) for sent in pred_relations])
+
+    # Count TP, FP and FN per type
+    for pred_sent, gt_sent in zip(pred_relations, gt_relations):
+        for rel_type in relation_types:
+            # strict mode takes argument types into account
+            if mode == "strict":
+                pred_rels = {
+                    (rel["head"], rel["head_type"], rel["tail"], rel["tail_type"])
+                    for rel in pred_sent
+                    if rel["type"] == rel_type
+                }
+                gt_rels = {
+                    (rel["head"], rel["head_type"], rel["tail"], rel["tail_type"])
+                    for rel in gt_sent
+                    if rel["type"] == rel_type
+                }
+
+            # boundaries mode only takes argument spans into account
+            elif mode == "boundaries":
+                pred_rels = {(rel["head"], rel["tail"]) for rel in pred_sent if rel["type"] == rel_type}
+                gt_rels = {(rel["head"], rel["tail"]) for rel in gt_sent if rel["type"] == rel_type}
+
+            scores[rel_type]["tp"] += len(pred_rels & gt_rels)
+            scores[rel_type]["fp"] += len(pred_rels - gt_rels)
+            scores[rel_type]["fn"] += len(gt_rels - pred_rels)
+
+    # Compute per entity Precision / Recall / F1
+    for rel_type in scores.keys():
+        if scores[rel_type]["tp"]:
+            scores[rel_type]["p"] = scores[rel_type]["tp"] / (scores[rel_type]["fp"] + scores[rel_type]["tp"])
+            scores[rel_type]["r"] = scores[rel_type]["tp"] / (scores[rel_type]["fn"] + scores[rel_type]["tp"])
+        else:
+            scores[rel_type]["p"], scores[rel_type]["r"] = 0, 0
+
+        if not scores[rel_type]["p"] + scores[rel_type]["r"] == 0:
+            scores[rel_type]["f1"] = (
+                2 * scores[rel_type]["p"] * scores[rel_type]["r"] / (scores[rel_type]["p"] + scores[rel_type]["r"])
+            )
+        else:
+            scores[rel_type]["f1"] = 0
+
+    # Compute micro F1 Scores
+    tp = sum([scores[rel_type]["tp"] for rel_type in relation_types])
+    fp = sum([scores[rel_type]["fp"] for rel_type in relation_types])
+    fn = sum([scores[rel_type]["fn"] for rel_type in relation_types])
+
+    if tp:
+        precision = tp / (tp + fp)
+        recall = tp / (tp + fn)
+        f1 = 2 * precision * recall / (precision + recall)
+
+    else:
+        precision, recall, f1 = 0, 0, 0
+
+    scores["ALL"]["p"] = precision
+    scores["ALL"]["r"] = recall
+    scores["ALL"]["f1"] = f1
+    scores["ALL"]["tp"] = tp
+    scores["ALL"]["fp"] = fp
+    scores["ALL"]["fn"] = fn
+
+    # Compute Macro F1 Scores
+    scores["ALL"]["Macro_f1"] = np.mean([scores[ent_type]["f1"] for ent_type in relation_types])
+    scores["ALL"]["Macro_p"] = np.mean([scores[ent_type]["p"] for ent_type in relation_types])
+    scores["ALL"]["Macro_r"] = np.mean([scores[ent_type]["r"] for ent_type in relation_types])
+
+    logger.info(f"RE Evaluation in *** {mode.upper()} *** mode")
+
+    logger.info(
+        "processed {} sentences with {} relations; found: {} relations; correct: {}.".format(
+            n_sents, n_rels, n_found, tp
+        )
+    )
+    logger.info(
+        "\tALL\t TP: {};\tFP: {};\tFN: {}".format(scores["ALL"]["tp"], scores["ALL"]["fp"], scores["ALL"]["fn"])
+    )
+    logger.info("\t\t(m avg): precision: {:.2f};\trecall: {:.2f};\tf1: {:.2f} (micro)".format(precision, recall, f1))
+    logger.info(
+        "\t\t(M avg): precision: {:.2f};\trecall: {:.2f};\tf1: {:.2f} (Macro)\n".format(
+            scores["ALL"]["Macro_p"], scores["ALL"]["Macro_r"], scores["ALL"]["Macro_f1"]
+        )
+    )
+
+    for rel_type in relation_types:
+        logger.info(
+            "\t{}: \tTP: {};\tFP: {};\tFN: {};\tprecision: {:.2f};\trecall: {:.2f};\tf1: {:.2f};\t{}".format(
+                rel_type,
+                scores[rel_type]["tp"],
+                scores[rel_type]["fp"],
+                scores[rel_type]["fn"],
+                scores[rel_type]["p"],
+                scores[rel_type]["r"],
+                scores[rel_type]["f1"],
+                scores[rel_type]["tp"] + scores[rel_type]["fp"],
+            )
+        )
+
+    return scores

xdoc/fine_tuning/funsd/layoutlmft/models/layoutlm/__init__.py

@@ -0,0 +1 @@
+from transformers.models.layoutlm import *

xdoc/fine_tuning/funsd/layoutlmft/models/layoutlmv2/__init__.py

@@ -0,0 +1,4 @@
+from .configuration_layoutlmv2 import LayoutLMv2Config
+from .modeling_layoutlmv2 import LayoutLMv2ForRelationExtraction, LayoutLMv2ForTokenClassification, LayoutLMv2Model
+from .tokenization_layoutlmv2 import LayoutLMv2Tokenizer
+from .tokenization_layoutlmv2_fast import LayoutLMv2TokenizerFast

xdoc/fine_tuning/funsd/layoutlmft/models/layoutlmv2/configuration_layoutlmv2.py

@@ -0,0 +1,77 @@
+# coding=utf-8
+from transformers.models.layoutlm.configuration_layoutlm import LayoutLMConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+LAYOUTLMV2_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "layoutlmv2-base-uncased": "https://huggingface.co/microsoft/layoutlmv2-base-uncased/resolve/main/config.json",
+    "layoutlmv2-large-uncased": "https://huggingface.co/microsoft/layoutlmv2-large-uncased/resolve/main/config.json",
+}
+
+
+class LayoutLMv2Config(LayoutLMConfig):
+    model_type = "layoutlmv2"
+
+    def __init__(
+        self,
+        vocab_size=30522,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        max_position_embeddings=512,
+        type_vocab_size=2,
+        initializer_range=0.02,
+        layer_norm_eps=1e-12,
+        pad_token_id=0,
+        gradient_checkpointing=False,
+        max_2d_position_embeddings=1024,
+        max_rel_pos=128,
+        rel_pos_bins=32,
+        fast_qkv=True,
+        max_rel_2d_pos=256,
+        rel_2d_pos_bins=64,
+        convert_sync_batchnorm=True,
+        image_feature_pool_shape=[7, 7, 256],
+        coordinate_size=128,
+        shape_size=128,
+        has_relative_attention_bias=True,
+        has_spatial_attention_bias=True,
+        has_visual_segment_embedding=False,
+        **kwargs
+    ):
+        super().__init__(
+            vocab_size=vocab_size,
+            hidden_size=hidden_size,
+            num_hidden_layers=num_hidden_layers,
+            num_attention_heads=num_attention_heads,
+            intermediate_size=intermediate_size,
+            hidden_act=hidden_act,
+            hidden_dropout_prob=hidden_dropout_prob,
+            attention_probs_dropout_prob=attention_probs_dropout_prob,
+            max_position_embeddings=max_position_embeddings,
+            type_vocab_size=type_vocab_size,
+            initializer_range=initializer_range,
+            layer_norm_eps=layer_norm_eps,
+            pad_token_id=pad_token_id,
+            gradient_checkpointing=gradient_checkpointing,
+            **kwargs,
+        )
+        self.max_2d_position_embeddings = max_2d_position_embeddings
+        self.max_rel_pos = max_rel_pos
+        self.rel_pos_bins = rel_pos_bins
+        self.fast_qkv = fast_qkv
+        self.max_rel_2d_pos = max_rel_2d_pos
+        self.rel_2d_pos_bins = rel_2d_pos_bins
+        self.convert_sync_batchnorm = convert_sync_batchnorm
+        self.image_feature_pool_shape = image_feature_pool_shape
+        self.coordinate_size = coordinate_size
+        self.shape_size = shape_size
+        self.has_relative_attention_bias = has_relative_attention_bias
+        self.has_spatial_attention_bias = has_spatial_attention_bias
+        self.has_visual_segment_embedding = has_visual_segment_embedding

xdoc/fine_tuning/funsd/layoutlmft/models/layoutlmv2/detectron2_config.py

@@ -0,0 +1,102 @@
+# -*- coding: utf-8 -*-
+def add_layoutlmv2_config(cfg):
+    _C = cfg
+    # -----------------------------------------------------------------------------
+    # Config definition
+    # -----------------------------------------------------------------------------
+    _C.MODEL.MASK_ON = True
+
+    # When using pre-trained models in Detectron1 or any MSRA models,
+    # std has been absorbed into its conv1 weights, so the std needs to be set 1.
+    # Otherwise, you can use [57.375, 57.120, 58.395] (ImageNet std)
+    _C.MODEL.PIXEL_STD = [57.375, 57.120, 58.395]
+
+    # ---------------------------------------------------------------------------- #
+    # Backbone options
+    # ---------------------------------------------------------------------------- #
+    _C.MODEL.BACKBONE.NAME = "build_resnet_fpn_backbone"
+
+    # ---------------------------------------------------------------------------- #
+    # FPN options
+    # ---------------------------------------------------------------------------- #
+    # Names of the input feature maps to be used by FPN
+    # They must have contiguous power of 2 strides
+    # e.g., ["res2", "res3", "res4", "res5"]
+    _C.MODEL.FPN.IN_FEATURES = ["res2", "res3", "res4", "res5"]
+
+    # ---------------------------------------------------------------------------- #
+    # Anchor generator options
+    # ---------------------------------------------------------------------------- #
+    # Anchor sizes (i.e. sqrt of area) in absolute pixels w.r.t. the network input.
+    # Format: list[list[float]]. SIZES[i] specifies the list of sizes
+    # to use for IN_FEATURES[i]; len(SIZES) == len(IN_FEATURES) must be true,
+    # or len(SIZES) == 1 is true and size list SIZES[0] is used for all
+    # IN_FEATURES.
+    _C.MODEL.ANCHOR_GENERATOR.SIZES = [[32], [64], [128], [256], [512]]
+
+    # ---------------------------------------------------------------------------- #
+    # RPN options
+    # ---------------------------------------------------------------------------- #
+    # Names of the input feature maps to be used by RPN
+    # e.g., ["p2", "p3", "p4", "p5", "p6"] for FPN
+    _C.MODEL.RPN.IN_FEATURES = ["p2", "p3", "p4", "p5", "p6"]
+    # Number of top scoring RPN proposals to keep before applying NMS
+    # When FPN is used, this is *per FPN level* (not total)
+    _C.MODEL.RPN.PRE_NMS_TOPK_TRAIN = 2000
+    _C.MODEL.RPN.PRE_NMS_TOPK_TEST = 1000
+    # Number of top scoring RPN proposals to keep after applying NMS
+    # When FPN is used, this limit is applied per level and then again to the union
+    # of proposals from all levels
+    # NOTE: When FPN is used, the meaning of this config is different from Detectron1.
+    # It means per-batch topk in Detectron1, but per-image topk here.
+    # See the "find_top_rpn_proposals" function for details.
+    _C.MODEL.RPN.POST_NMS_TOPK_TRAIN = 1000
+    _C.MODEL.RPN.POST_NMS_TOPK_TEST = 1000
+
+    # ---------------------------------------------------------------------------- #
+    # ROI HEADS options
+    # ---------------------------------------------------------------------------- #
+    _C.MODEL.ROI_HEADS.NAME = "StandardROIHeads"
+    # Number of foreground classes
+    _C.MODEL.ROI_HEADS.NUM_CLASSES = 5
+    # Names of the input feature maps to be used by ROI heads
+    # Currently all heads (box, mask, ...) use the same input feature map list
+    # e.g., ["p2", "p3", "p4", "p5"] is commonly used for FPN
+    _C.MODEL.ROI_HEADS.IN_FEATURES = ["p2", "p3", "p4", "p5"]
+
+    # ---------------------------------------------------------------------------- #
+    # Box Head
+    # ---------------------------------------------------------------------------- #
+    # C4 don't use head name option
+    # Options for non-C4 models: FastRCNNConvFCHead,
+    _C.MODEL.ROI_BOX_HEAD.NAME = "FastRCNNConvFCHead"
+    _C.MODEL.ROI_BOX_HEAD.NUM_FC = 2
+    _C.MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION = 14
+
+    # ---------------------------------------------------------------------------- #
+    # Mask Head
+    # ---------------------------------------------------------------------------- #
+    _C.MODEL.ROI_MASK_HEAD.NAME = "MaskRCNNConvUpsampleHead"
+    _C.MODEL.ROI_MASK_HEAD.NUM_CONV = 4  # The number of convs in the mask head
+    _C.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION = 7
+
+    # ---------------------------------------------------------------------------- #
+    # ResNe[X]t options (ResNets = {ResNet, ResNeXt}
+    # Note that parts of a resnet may be used for both the backbone and the head
+    # These options apply to both
+    # ---------------------------------------------------------------------------- #
+    _C.MODEL.RESNETS.DEPTH = 101
+    _C.MODEL.RESNETS.SIZES = [[32], [64], [128], [256], [512]]
+    _C.MODEL.RESNETS.ASPECT_RATIOS = [[0.5, 1.0, 2.0]]
+    _C.MODEL.RESNETS.OUT_FEATURES = ["res2", "res3", "res4", "res5"]  # res4 for C4 backbone, res2..5 for FPN backbone
+
+    # Number of groups to use; 1 ==> ResNet; > 1 ==> ResNeXt
+    _C.MODEL.RESNETS.NUM_GROUPS = 32
+
+    # Baseline width of each group.
+    # Scaling this parameters will scale the width of all bottleneck layers.
+    _C.MODEL.RESNETS.WIDTH_PER_GROUP = 8
+
+    # Place the stride 2 conv on the 1x1 filter
+    # Use True only for the original MSRA ResNet; use False for C2 and Torch models
+    _C.MODEL.RESNETS.STRIDE_IN_1X1 = False

xdoc/fine_tuning/funsd/layoutlmft/models/layoutlmv2/modeling_layoutlmv2.py

@@ -0,0 +1,937 @@
+# coding=utf-8
+import math
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+import detectron2
+from detectron2.modeling import META_ARCH_REGISTRY
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
+from transformers.models.layoutlm.modeling_layoutlm import LayoutLMIntermediate as LayoutLMv2Intermediate
+from transformers.models.layoutlm.modeling_layoutlm import LayoutLMOutput as LayoutLMv2Output
+from transformers.models.layoutlm.modeling_layoutlm import LayoutLMPooler as LayoutLMv2Pooler
+from transformers.models.layoutlm.modeling_layoutlm import LayoutLMSelfOutput as LayoutLMv2SelfOutput
+from transformers.utils import logging
+
+from ...modules.decoders.re import REDecoder
+from ...utils import ReOutput
+from .configuration_layoutlmv2 import LayoutLMv2Config
+from .detectron2_config import add_layoutlmv2_config
+
+
+logger = logging.get_logger(__name__)
+
+LAYOUTLMV2_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "layoutlmv2-base-uncased",
+    "layoutlmv2-large-uncased",
+]
+
+
+LayoutLMv2LayerNorm = torch.nn.LayerNorm
+
+
+class LayoutLMv2Embeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings."""
+
+    def __init__(self, config):
+        super(LayoutLMv2Embeddings, self).__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+
+        self.x_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.coordinate_size)
+        self.y_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.coordinate_size)
+        self.h_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.shape_size)
+        self.w_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.shape_size)
+        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+        self.LayerNorm = LayoutLMv2LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
+
+    def _cal_spatial_position_embeddings(self, bbox):
+        try:
+            left_position_embeddings = self.x_position_embeddings(bbox[:, :, 0])
+            upper_position_embeddings = self.y_position_embeddings(bbox[:, :, 1])
+            right_position_embeddings = self.x_position_embeddings(bbox[:, :, 2])
+            lower_position_embeddings = self.y_position_embeddings(bbox[:, :, 3])
+        except IndexError as e:
+            raise IndexError("The :obj:`bbox`coordinate values should be within 0-1000 range.") from e
+
+        h_position_embeddings = self.h_position_embeddings(bbox[:, :, 3] - bbox[:, :, 1])
+        w_position_embeddings = self.w_position_embeddings(bbox[:, :, 2] - bbox[:, :, 0])
+
+        spatial_position_embeddings = torch.cat(
+            [
+                left_position_embeddings,
+                upper_position_embeddings,
+                right_position_embeddings,
+                lower_position_embeddings,
+                h_position_embeddings,
+                w_position_embeddings,
+            ],
+            dim=-1,
+        )
+        return spatial_position_embeddings
+
+
+class LayoutLMv2SelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+            raise ValueError(
+                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+                f"heads ({config.num_attention_heads})"
+            )
+        self.fast_qkv = config.fast_qkv
+        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.has_relative_attention_bias = config.has_relative_attention_bias
+        self.has_spatial_attention_bias = config.has_spatial_attention_bias
+
+        if config.fast_qkv:
+            self.qkv_linear = nn.Linear(config.hidden_size, 3 * self.all_head_size, bias=False)
+            self.q_bias = nn.Parameter(torch.zeros(1, 1, self.all_head_size))
+            self.v_bias = nn.Parameter(torch.zeros(1, 1, self.all_head_size))
+        else:
+            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 compute_qkv(self, hidden_states):
+        if self.fast_qkv:
+            qkv = self.qkv_linear(hidden_states)
+            q, k, v = torch.chunk(qkv, 3, dim=-1)
+            if q.ndimension() == self.q_bias.ndimension():
+                q = q + self.q_bias
+                v = v + self.v_bias
+            else:
+                _sz = (1,) * (q.ndimension() - 1) + (-1,)
+                q = q + self.q_bias.view(*_sz)
+                v = v + self.v_bias.view(*_sz)
+        else:
+            q = self.query(hidden_states)
+            k = self.key(hidden_states)
+            v = self.value(hidden_states)
+        return q, k, v
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+        rel_pos=None,
+        rel_2d_pos=None,
+    ):
+        q, k, v = self.compute_qkv(hidden_states)
+
+        # (B, L, H*D) -> (B, H, L, D)
+        query_layer = self.transpose_for_scores(q)
+        key_layer = self.transpose_for_scores(k)
+        value_layer = self.transpose_for_scores(v)
+
+        query_layer = query_layer / math.sqrt(self.attention_head_size)
+        # [BSZ, NAT, L, L]
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        if self.has_relative_attention_bias:
+            attention_scores += rel_pos
+        if self.has_spatial_attention_bias:
+            attention_scores += rel_2d_pos
+        attention_scores = attention_scores.float().masked_fill_(attention_mask.to(torch.bool), float("-inf"))
+        attention_probs = F.softmax(attention_scores, dim=-1, dtype=torch.float32).type_as(value_layer)
+        # 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)
+
+        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+        return outputs
+
+
+class LayoutLMv2Attention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self = LayoutLMv2SelfAttention(config)
+        self.output = LayoutLMv2SelfOutput(config)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
+        )
+
+        # Prune linear layers
+        self.self.query = prune_linear_layer(self.self.query, index)
+        self.self.key = prune_linear_layer(self.self.key, index)
+        self.self.value = prune_linear_layer(self.self.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
+        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+        rel_pos=None,
+        rel_2d_pos=None,
+    ):
+        self_outputs = self.self(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            past_key_value,
+            output_attentions,
+            rel_pos=rel_pos,
+            rel_2d_pos=rel_2d_pos,
+        )
+        attention_output = self.output(self_outputs[0], hidden_states)
+        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+class LayoutLMv2Layer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = LayoutLMv2Attention(config)
+        self.is_decoder = config.is_decoder
+        self.add_cross_attention = config.add_cross_attention
+        if self.add_cross_attention:
+            assert self.is_decoder, f"{self} should be used as a decoder model if cross attention is added"
+            self.crossattention = LayoutLMv2Attention(config)
+        self.intermediate = LayoutLMv2Intermediate(config)
+        self.output = LayoutLMv2Output(config)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+        rel_pos=None,
+        rel_2d_pos=None,
+    ):
+        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
+        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
+        self_attention_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            output_attentions=output_attentions,
+            past_key_value=self_attn_past_key_value,
+            rel_pos=rel_pos,
+            rel_2d_pos=rel_2d_pos,
+        )
+        attention_output = self_attention_outputs[0]
+
+        # if decoder, the last output is tuple of self-attn cache
+        if self.is_decoder:
+            outputs = self_attention_outputs[1:-1]
+            present_key_value = self_attention_outputs[-1]
+        else:
+            outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights
+
+        cross_attn_present_key_value = None
+        if self.is_decoder and encoder_hidden_states is not None:
+            assert hasattr(
+                self, "crossattention"
+            ), f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"
+
+            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
+            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
+            cross_attention_outputs = self.crossattention(
+                attention_output,
+                attention_mask,
+                head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                cross_attn_past_key_value,
+                output_attentions,
+            )
+            attention_output = cross_attention_outputs[0]
+            outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights
+
+            # add cross-attn cache to positions 3,4 of present_key_value tuple
+            cross_attn_present_key_value = cross_attention_outputs[-1]
+            present_key_value = present_key_value + cross_attn_present_key_value
+
+        layer_output = apply_chunking_to_forward(
+            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+        )
+        outputs = (layer_output,) + outputs
+
+        # if decoder, return the attn key/values as the last output
+        if self.is_decoder:
+            outputs = outputs + (present_key_value,)
+
+        return outputs
+
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+
+def relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
+    ret = 0
+    if bidirectional:
+        num_buckets //= 2
+        ret += (relative_position > 0).long() * num_buckets
+        n = torch.abs(relative_position)
+    else:
+        n = torch.max(-relative_position, torch.zeros_like(relative_position))
+    # now n is in the range [0, inf)
+
+    # half of the buckets are for exact increments in positions
+    max_exact = num_buckets // 2
+    is_small = n < max_exact
+
+    # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance
+    val_if_large = max_exact + (
+        torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact)
+    ).to(torch.long)
+    val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1))
+
+    ret += torch.where(is_small, n, val_if_large)
+    return ret
+
+
+class LayoutLMv2Encoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList([LayoutLMv2Layer(config) for _ in range(config.num_hidden_layers)])
+
+        self.has_relative_attention_bias = config.has_relative_attention_bias
+        self.has_spatial_attention_bias = config.has_spatial_attention_bias
+
+        if self.has_relative_attention_bias:
+            self.rel_pos_bins = config.rel_pos_bins
+            self.max_rel_pos = config.max_rel_pos
+            self.rel_pos_onehot_size = config.rel_pos_bins
+            self.rel_pos_bias = nn.Linear(self.rel_pos_onehot_size, config.num_attention_heads, bias=False)
+
+        if self.has_spatial_attention_bias:
+            self.max_rel_2d_pos = config.max_rel_2d_pos
+            self.rel_2d_pos_bins = config.rel_2d_pos_bins
+            self.rel_2d_pos_onehot_size = config.rel_2d_pos_bins
+            self.rel_pos_x_bias = nn.Linear(self.rel_2d_pos_onehot_size, config.num_attention_heads, bias=False)
+            self.rel_pos_y_bias = nn.Linear(self.rel_2d_pos_onehot_size, config.num_attention_heads, bias=False)
+
+    def _cal_1d_pos_emb(self, hidden_states, position_ids):
+        rel_pos_mat = position_ids.unsqueeze(-2) - position_ids.unsqueeze(-1)
+        rel_pos = relative_position_bucket(
+            rel_pos_mat,
+            num_buckets=self.rel_pos_bins,
+            max_distance=self.max_rel_pos,
+        )
+        rel_pos = F.one_hot(rel_pos, num_classes=self.rel_pos_onehot_size).type_as(hidden_states)
+        rel_pos = self.rel_pos_bias(rel_pos).permute(0, 3, 1, 2)
+        rel_pos = rel_pos.contiguous()
+        return rel_pos
+
+    def _cal_2d_pos_emb(self, hidden_states, bbox):
+        position_coord_x = bbox[:, :, 0]
+        position_coord_y = bbox[:, :, 3]
+        rel_pos_x_2d_mat = position_coord_x.unsqueeze(-2) - position_coord_x.unsqueeze(-1)
+        rel_pos_y_2d_mat = position_coord_y.unsqueeze(-2) - position_coord_y.unsqueeze(-1)
+        rel_pos_x = relative_position_bucket(
+            rel_pos_x_2d_mat,
+            num_buckets=self.rel_2d_pos_bins,
+            max_distance=self.max_rel_2d_pos,
+        )
+        rel_pos_y = relative_position_bucket(
+            rel_pos_y_2d_mat,
+            num_buckets=self.rel_2d_pos_bins,
+            max_distance=self.max_rel_2d_pos,
+        )
+        rel_pos_x = F.one_hot(rel_pos_x, num_classes=self.rel_2d_pos_onehot_size).type_as(hidden_states)
+        rel_pos_y = F.one_hot(rel_pos_y, num_classes=self.rel_2d_pos_onehot_size).type_as(hidden_states)
+        rel_pos_x = self.rel_pos_x_bias(rel_pos_x).permute(0, 3, 1, 2)
+        rel_pos_y = self.rel_pos_y_bias(rel_pos_y).permute(0, 3, 1, 2)
+        rel_pos_x = rel_pos_x.contiguous()
+        rel_pos_y = rel_pos_y.contiguous()
+        rel_2d_pos = rel_pos_x + rel_pos_y
+        return rel_2d_pos
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=True,
+        bbox=None,
+        position_ids=None,
+    ):
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
+
+        next_decoder_cache = () if use_cache else None
+
+        rel_pos = self._cal_1d_pos_emb(hidden_states, position_ids) if self.has_relative_attention_bias else None
+        rel_2d_pos = self._cal_2d_pos_emb(hidden_states, bbox) if self.has_spatial_attention_bias else None
+
+        for i, layer_module in enumerate(self.layer):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            past_key_value = past_key_values[i] if past_key_values is not None else None
+
+            if getattr(self.config, "gradient_checkpointing", False) and self.training:
+
+                if use_cache:
+                    logger.warn(
+                        "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "
+                        "`use_cache=False`..."
+                    )
+                    use_cache = False
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, past_key_value, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    rel_pos=rel_pos,
+                    rel_2d_pos=rel_2d_pos,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    past_key_value,
+                    output_attentions,
+                    rel_pos=rel_pos,
+                    rel_2d_pos=rel_2d_pos,
+                )
+
+            hidden_states = layer_outputs[0]
+            if use_cache:
+                next_decoder_cache += (layer_outputs[-1],)
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+                if self.config.add_cross_attention:
+                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    next_decoder_cache,
+                    all_hidden_states,
+                    all_self_attentions,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+class LayoutLMv2PreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = LayoutLMv2Config
+    pretrained_model_archive_map = LAYOUTLMV2_PRETRAINED_MODEL_ARCHIVE_LIST
+    base_model_prefix = "layoutlmv2"
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, LayoutLMv2LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+def my_convert_sync_batchnorm(module, process_group=None):
+    # same as `nn.modules.SyncBatchNorm.convert_sync_batchnorm` but allowing converting from `detectron2.layers.FrozenBatchNorm2d`
+    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
+        return nn.modules.SyncBatchNorm.convert_sync_batchnorm(module, process_group)
+    module_output = module
+    if isinstance(module, detectron2.layers.FrozenBatchNorm2d):
+        module_output = torch.nn.SyncBatchNorm(
+            num_features=module.num_features,
+            eps=module.eps,
+            affine=True,
+            track_running_stats=True,
+            process_group=process_group,
+        )
+        module_output.weight = torch.nn.Parameter(module.weight)
+        module_output.bias = torch.nn.Parameter(module.bias)
+        module_output.running_mean = module.running_mean
+        module_output.running_var = module.running_var
+        module_output.num_batches_tracked = torch.tensor(0, dtype=torch.long, device=module.running_mean.device)
+    for name, child in module.named_children():
+        module_output.add_module(name, my_convert_sync_batchnorm(child, process_group))
+    del module
+    return module_output
+
+
+class VisualBackbone(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.cfg = detectron2.config.get_cfg()
+        add_layoutlmv2_config(self.cfg)
+        meta_arch = self.cfg.MODEL.META_ARCHITECTURE
+        model = META_ARCH_REGISTRY.get(meta_arch)(self.cfg)
+        assert isinstance(model.backbone, detectron2.modeling.backbone.FPN)
+        self.backbone = model.backbone
+        if (
+            config.convert_sync_batchnorm
+            and torch.distributed.is_available()
+            and torch.distributed.is_initialized()
+            and torch.distributed.get_rank() > -1
+        ):
+            self_rank = torch.distributed.get_rank()
+            node_size = torch.cuda.device_count()
+            world_size = torch.distributed.get_world_size()
+            assert world_size % node_size == 0
+
+            node_global_ranks = [
+                list(range(i * node_size, (i + 1) * node_size)) for i in range(world_size // node_size)
+            ]
+            sync_bn_groups = [
+                torch.distributed.new_group(ranks=node_global_ranks[i]) for i in range(world_size // node_size)
+            ]
+            node_rank = self_rank // node_size
+            assert self_rank in node_global_ranks[node_rank]
+
+            self.backbone = my_convert_sync_batchnorm(self.backbone, process_group=sync_bn_groups[node_rank])
+
+        assert len(self.cfg.MODEL.PIXEL_MEAN) == len(self.cfg.MODEL.PIXEL_STD)
+        num_channels = len(self.cfg.MODEL.PIXEL_MEAN)
+        self.register_buffer(
+            "pixel_mean",
+            torch.Tensor(self.cfg.MODEL.PIXEL_MEAN).view(num_channels, 1, 1),
+        )
+        self.register_buffer("pixel_std", torch.Tensor(self.cfg.MODEL.PIXEL_STD).view(num_channels, 1, 1))
+        self.out_feature_key = "p2"
+        if torch.is_deterministic():
+            logger.warning("using `AvgPool2d` instead of `AdaptiveAvgPool2d`")
+            input_shape = (224, 224)
+            backbone_stride = self.backbone.output_shape()[self.out_feature_key].stride
+            self.pool = nn.AvgPool2d(
+                (
+                    math.ceil(math.ceil(input_shape[0] / backbone_stride) / config.image_feature_pool_shape[0]),
+                    math.ceil(math.ceil(input_shape[1] / backbone_stride) / config.image_feature_pool_shape[1]),
+                )
+            )
+        else:
+            self.pool = nn.AdaptiveAvgPool2d(config.image_feature_pool_shape[:2])
+        if len(config.image_feature_pool_shape) == 2:
+            config.image_feature_pool_shape.append(self.backbone.output_shape()[self.out_feature_key].channels)
+        assert self.backbone.output_shape()[self.out_feature_key].channels == config.image_feature_pool_shape[2]
+
+    def forward(self, images):      
+        images_input = ((images if torch.is_tensor(images) else images.tensor) - self.pixel_mean) / self.pixel_std
+        features = self.backbone(images_input)
+        features = features[self.out_feature_key]
+        features = self.pool(features).flatten(start_dim=2).transpose(1, 2).contiguous()
+        return features
+
+
+class LayoutLMv2Model(LayoutLMv2PreTrainedModel):
+    def __init__(self, config):
+        super(LayoutLMv2Model, self).__init__(config)
+        self.config = config
+        self.has_visual_segment_embedding = config.has_visual_segment_embedding
+        self.embeddings = LayoutLMv2Embeddings(config)
+
+        self.visual = VisualBackbone(config)
+        self.visual_proj = nn.Linear(config.image_feature_pool_shape[-1], config.hidden_size)
+        if self.has_visual_segment_embedding:
+            self.visual_segment_embedding = nn.Parameter(nn.Embedding(1, config.hidden_size).weight[0])
+        self.visual_LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.visual_dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        self.encoder = LayoutLMv2Encoder(config)
+        self.pooler = LayoutLMv2Pooler(config)
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    def _calc_text_embeddings(self, input_ids, bbox, position_ids, token_type_ids):
+        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)
+
+        words_embeddings = self.embeddings.word_embeddings(input_ids)
+        position_embeddings = self.embeddings.position_embeddings(position_ids)
+        spatial_position_embeddings = self.embeddings._cal_spatial_position_embeddings(bbox)
+        token_type_embeddings = self.embeddings.token_type_embeddings(token_type_ids)
+        embeddings = words_embeddings + position_embeddings + spatial_position_embeddings + token_type_embeddings
+        embeddings = self.embeddings.LayerNorm(embeddings)
+        embeddings = self.embeddings.dropout(embeddings)
+        return embeddings
+
+    def _calc_img_embeddings(self, image, bbox, position_ids):
+        visual_embeddings = self.visual_proj(self.visual(image))
+        position_embeddings = self.embeddings.position_embeddings(position_ids)
+        spatial_position_embeddings = self.embeddings._cal_spatial_position_embeddings(bbox)
+        embeddings = visual_embeddings + position_embeddings + spatial_position_embeddings
+        if self.has_visual_segment_embedding:
+            embeddings += self.visual_segment_embedding
+        embeddings = self.visual_LayerNorm(embeddings)
+        embeddings = self.visual_dropout(embeddings)
+        return embeddings
+
+    def forward(
+        self,
+        input_ids=None,
+        bbox=None,
+        image=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if 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
+
+        visual_shape = list(input_shape)
+        visual_shape[1] = self.config.image_feature_pool_shape[0] * self.config.image_feature_pool_shape[1]
+        visual_shape = torch.Size(visual_shape)
+        final_shape = list(input_shape)
+        final_shape[1] += visual_shape[1]
+        final_shape = torch.Size(final_shape)
+
+        visual_bbox_x = (
+            torch.arange(
+                0,
+                1000 * (self.config.image_feature_pool_shape[1] + 1),
+                1000,
+                device=device,
+                dtype=bbox.dtype,
+            )
+            // self.config.image_feature_pool_shape[1]
+        )
+        visual_bbox_y = (
+            torch.arange(
+                0,
+                1000 * (self.config.image_feature_pool_shape[0] + 1),
+                1000,
+                device=device,
+                dtype=bbox.dtype,
+            )
+            // self.config.image_feature_pool_shape[0]
+        )
+        visual_bbox = torch.stack(
+            [
+                visual_bbox_x[:-1].repeat(self.config.image_feature_pool_shape[0], 1),
+                visual_bbox_y[:-1].repeat(self.config.image_feature_pool_shape[1], 1).transpose(0, 1),
+                visual_bbox_x[1:].repeat(self.config.image_feature_pool_shape[0], 1),
+                visual_bbox_y[1:].repeat(self.config.image_feature_pool_shape[1], 1).transpose(0, 1),
+            ],
+            dim=-1,
+        ).view(-1, bbox.size(-1))
+        visual_bbox = visual_bbox.repeat(final_shape[0], 1, 1)
+        final_bbox = torch.cat([bbox, visual_bbox], dim=1)
+
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=device)
+
+        visual_attention_mask = torch.ones(visual_shape, device=device)
+        final_attention_mask = torch.cat([attention_mask, visual_attention_mask], dim=1)
+
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        if position_ids is None:
+            seq_length = input_shape[1]
+            position_ids = self.embeddings.position_ids[:, :seq_length]
+            position_ids = position_ids.expand_as(input_ids)
+
+        visual_position_ids = torch.arange(0, visual_shape[1], dtype=torch.long, device=device).repeat(
+            input_shape[0], 1
+        )
+        final_position_ids = torch.cat([position_ids, visual_position_ids], dim=1)
+
+        if bbox is None:
+            bbox = torch.zeros(tuple(list(input_shape) + [4]), dtype=torch.long, device=device)
+
+        text_layout_emb = self._calc_text_embeddings(
+            input_ids=input_ids,
+            bbox=bbox,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+        )
+
+        visual_emb = self._calc_img_embeddings(
+            image=image,
+            bbox=visual_bbox,
+            position_ids=visual_position_ids,
+        )
+        final_emb = torch.cat([text_layout_emb, visual_emb], dim=1)
+
+        extended_attention_mask = final_attention_mask.unsqueeze(1).unsqueeze(2)
+
+        extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+
+        if head_mask is not None:
+            if head_mask.dim() == 1:
+                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+                head_mask = head_mask.expand(self.config.num_hidden_layers, -1, -1, -1, -1)
+            elif head_mask.dim() == 2:
+                head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
+            head_mask = head_mask.to(dtype=next(self.parameters()).dtype)
+        else:
+            head_mask = [None] * self.config.num_hidden_layers
+
+        encoder_outputs = self.encoder(
+            final_emb,
+            extended_attention_mask,
+            bbox=final_bbox,
+            position_ids=final_position_ids,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output)
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            cross_attentions=encoder_outputs.cross_attentions,
+        )
+
+
+class LayoutLMv2ForTokenClassification(LayoutLMv2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.layoutlmv2 = LayoutLMv2Model(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.layoutlmv2.embeddings.word_embeddings
+
+    def forward(
+        self,
+        input_ids=None,
+        bbox=None,
+        image=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.layoutlmv2(
+            input_ids=input_ids,
+            bbox=bbox,
+            image=image,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        seq_length = input_ids.size(1)
+        sequence_output, image_output = outputs[0][:, :seq_length], outputs[0][:, seq_length:]
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)[active_loss]
+                active_labels = labels.view(-1)[active_loss]
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class LayoutLMv2ForRelationExtraction(LayoutLMv2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.layoutlmv2 = LayoutLMv2Model(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.extractor = REDecoder(config)
+        self.init_weights()
+
+    def forward(
+        self,
+        input_ids,
+        bbox,
+        labels=None,
+        image=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        entities=None,
+        relations=None,
+    ):
+        outputs = self.layoutlmv2(
+            input_ids=input_ids,
+            bbox=bbox,
+            image=image,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+        )
+
+        seq_length = input_ids.size(1)
+        sequence_output, image_output = outputs[0][:, :seq_length], outputs[0][:, seq_length:]
+        sequence_output = self.dropout(sequence_output)
+        loss, pred_relations = self.extractor(sequence_output, entities, relations)
+
+        return ReOutput(
+            loss=loss,
+            entities=entities,
+            relations=relations,
+            pred_relations=pred_relations,
+            hidden_states=outputs[0],
+        )

xdoc/fine_tuning/funsd/layoutlmft/models/layoutlmv2/tokenization_layoutlmv2.py

@@ -0,0 +1,47 @@
+# coding=utf-8
+from transformers.models.layoutlm.tokenization_layoutlm import LayoutLMTokenizer
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "microsoft/layoutlmv2-base-uncased": "https://huggingface.co/microsoft/layoutlmv2-base-uncased/resolve/main/vocab.txt",
+        "microsoft/layoutlmv2-large-uncased": "https://huggingface.co/microsoft/layoutlmv2-large-uncased/resolve/main/vocab.txt",
+    }
+}
+
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "microsoft/layoutlmv2-base-uncased": 512,
+    "microsoft/layoutlmv2-large-uncased": 512,
+}
+
+
+PRETRAINED_INIT_CONFIGURATION = {
+    "microsoft/layoutlmv2-base-uncased": {"do_lower_case": True},
+    "microsoft/layoutlmv2-large-uncased": {"do_lower_case": True},
+}
+
+
+class LayoutLMv2Tokenizer(LayoutLMTokenizer):
+    r"""
+    Constructs a LayoutLMv2 tokenizer.
+
+    :class:`~transformers.LayoutLMv2Tokenizer is identical to :class:`~transformers.BertTokenizer` and runs end-to-end
+    tokenization: punctuation splitting + wordpiece.
+
+    Refer to superclass :class:`~transformers.BertTokenizer` for usage examples and documentation concerning
+    parameters.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    pretrained_init_configuration = PRETRAINED_INIT_CONFIGURATION
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+
+    def __init__(self, model_max_length=512, **kwargs):
+        super().__init__(model_max_length=model_max_length, **kwargs)

xdoc/fine_tuning/funsd/layoutlmft/models/layoutlmv2/tokenization_layoutlmv2_fast.py

@@ -0,0 +1,51 @@
+# coding=utf-8
+from transformers.models.layoutlm.tokenization_layoutlm_fast import LayoutLMTokenizerFast
+from transformers.utils import logging
+
+from .tokenization_layoutlmv2 import LayoutLMv2Tokenizer
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "microsoft/layoutlmv2-base-uncased": "https://huggingface.co/microsoft/layoutlmv2-base-uncased/resolve/main/vocab.txt",
+        "microsoft/layoutlmv2-large-uncased": "https://huggingface.co/microsoft/layoutlmv2-large-uncased/resolve/main/vocab.txt",
+    },
+    "tokenizer_file": {
+        "microsoft/layoutlmv2-base-uncased": "https://huggingface.co/microsoft/layoutlmv2-base-uncased/resolve/main/tokenizer.json",
+        "microsoft/layoutlmv2-large-uncased": "https://huggingface.co/microsoft/layoutlmv2-large-uncased/resolve/main/tokenizer.json",
+    },
+}
+
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "microsoft/layoutlmv2-base-uncased": 512,
+    "microsoft/layoutlmv2-large-uncased": 512,
+}
+
+
+PRETRAINED_INIT_CONFIGURATION = {
+    "microsoft/layoutlmv2-base-uncased": {"do_lower_case": True},
+    "microsoft/layoutlmv2-large-uncased": {"do_lower_case": True},
+}
+
+
+class LayoutLMv2TokenizerFast(LayoutLMTokenizerFast):
+    r"""
+    Constructs a "Fast" LayoutLMv2Tokenizer.
+
+    Refer to superclass :class:`~transformers.BertTokenizerFast` for usage examples and documentation concerning
+    parameters.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    pretrained_init_configuration = PRETRAINED_INIT_CONFIGURATION
+    slow_tokenizer_class = LayoutLMv2Tokenizer
+
+    def __init__(self, model_max_length=512, **kwargs):
+        super().__init__(model_max_length=model_max_length, **kwargs)

xdoc/fine_tuning/funsd/layoutlmft/models/layoutxlm/__init__.py

@@ -0,0 +1,4 @@
+from .configuration_layoutxlm import LayoutXLMConfig
+from .modeling_layoutxlm import LayoutXLMForRelationExtraction, LayoutXLMForTokenClassification, LayoutXLMModel
+from .tokenization_layoutxlm import LayoutXLMTokenizer
+from .tokenization_layoutxlm_fast import LayoutXLMTokenizerFast

xdoc/fine_tuning/funsd/layoutlmft/models/layoutxlm/configuration_layoutxlm.py

@@ -0,0 +1,16 @@
+# coding=utf-8
+from transformers.utils import logging
+
+from ..layoutlmv2 import LayoutLMv2Config
+
+
+logger = logging.get_logger(__name__)
+
+LAYOUTXLM_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "layoutxlm-base": "https://huggingface.co/layoutxlm-base/resolve/main/config.json",
+    "layoutxlm-large": "https://huggingface.co/layoutxlm-large/resolve/main/config.json",
+}
+
+
+class LayoutXLMConfig(LayoutLMv2Config):
+    model_type = "layoutxlm"

xdoc/fine_tuning/funsd/layoutlmft/models/layoutxlm/modeling_layoutxlm.py

@@ -0,0 +1,25 @@
+# coding=utf-8
+from transformers.utils import logging
+
+from ..layoutlmv2 import LayoutLMv2ForRelationExtraction, LayoutLMv2ForTokenClassification, LayoutLMv2Model
+from .configuration_layoutxlm import LayoutXLMConfig
+
+
+logger = logging.get_logger(__name__)
+
+LAYOUTXLM_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "layoutxlm-base",
+    "layoutxlm-large",
+]
+
+
+class LayoutXLMModel(LayoutLMv2Model):
+    config_class = LayoutXLMConfig
+
+
+class LayoutXLMForTokenClassification(LayoutLMv2ForTokenClassification):
+    config_class = LayoutXLMConfig
+
+
+class LayoutXLMForRelationExtraction(LayoutLMv2ForRelationExtraction):
+    config_class = LayoutXLMConfig

xdoc/fine_tuning/funsd/layoutlmft/models/layoutxlm/tokenization_layoutxlm.py

@@ -0,0 +1,33 @@
+# coding=utf-8
+
+from transformers import XLMRobertaTokenizer
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+SPIECE_UNDERLINE = "▁"
+
+VOCAB_FILES_NAMES = {"vocab_file": "sentencepiece.bpe.model"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "layoutxlm-base": "https://huggingface.co/layoutxlm-base/resolve/main/sentencepiece.bpe.model",
+        "layoutxlm-large": "https://huggingface.co/layoutxlm-large/resolve/main/sentencepiece.bpe.model",
+    }
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "layoutxlm-base": 512,
+    "layoutxlm-large": 512,
+}
+
+
+class LayoutXLMTokenizer(XLMRobertaTokenizer):
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(self, model_max_length=512, **kwargs):
+        super().__init__(model_max_length=model_max_length, **kwargs)

xdoc/fine_tuning/funsd/layoutlmft/models/layoutxlm/tokenization_layoutxlm_fast.py

@@ -0,0 +1,43 @@
+# coding=utf-8
+from transformers import XLMRobertaTokenizerFast
+from transformers.file_utils import is_sentencepiece_available
+from transformers.utils import logging
+
+
+if is_sentencepiece_available():
+    from .tokenization_layoutxlm import LayoutXLMTokenizer
+else:
+    LayoutXLMTokenizer = None
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "sentencepiece.bpe.model", "tokenizer_file": "tokenizer.json"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "layoutxlm-base": "https://huggingface.co/layoutxlm-base/resolve/main/sentencepiece.bpe.model",
+        "layoutxlm-large": "https://huggingface.co/layoutxlm-large/resolve/main/sentencepiece.bpe.model",
+    },
+    "tokenizer_file": {
+        "layoutxlm-base": "https://huggingface.co/layoutxlm-base/resolve/main/tokenizer.json",
+        "layoutxlm-large": "https://huggingface.co/layoutxlm-large/resolve/main/tokenizer.json",
+    },
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "layoutxlm-base": 512,
+    "layoutxlm-large": 512,
+}
+
+
+class LayoutXLMTokenizerFast(XLMRobertaTokenizerFast):
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    model_input_names = ["input_ids", "attention_mask"]
+    slow_tokenizer_class = LayoutXLMTokenizer
+
+    def __init__(self, model_max_length=512, **kwargs):
+        super().__init__(model_max_length=model_max_length, **kwargs)

xdoc/fine_tuning/funsd/layoutlmft/models/model_args.py

@@ -0,0 +1,34 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    use_auth_token: bool = field(
+        default=False,
+        metadata={
+            "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
+            "with private models)."
+        },
+    )

xdoc/fine_tuning/funsd/layoutlmft/modules/decoders/re.py

@@ -0,0 +1,154 @@
+import copy
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+
+class BiaffineAttention(torch.nn.Module):
+    """Implements a biaffine attention operator for binary relation classification.
+
+    PyTorch implementation of the biaffine attention operator from "End-to-end neural relation
+    extraction using deep biaffine attention" (https://arxiv.org/abs/1812.11275) which can be used
+    as a classifier for binary relation classification.
+
+    Args:
+        in_features (int): The size of the feature dimension of the inputs.
+        out_features (int): The size of the feature dimension of the output.
+
+    Shape:
+        - x_1: `(N, *, in_features)` where `N` is the batch dimension and `*` means any number of
+          additional dimensisons.
+        - x_2: `(N, *, in_features)`, where `N` is the batch dimension and `*` means any number of
+          additional dimensions.
+        - Output: `(N, *, out_features)`, where `N` is the batch dimension and `*` means any number
+            of additional dimensions.
+
+    Examples:
+        >>> batch_size, in_features, out_features = 32, 100, 4
+        >>> biaffine_attention = BiaffineAttention(in_features, out_features)
+        >>> x_1 = torch.randn(batch_size, in_features)
+        >>> x_2 = torch.randn(batch_size, in_features)
+        >>> output = biaffine_attention(x_1, x_2)
+        >>> print(output.size())
+        torch.Size([32, 4])
+    """
+
+    def __init__(self, in_features, out_features):
+        super(BiaffineAttention, self).__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+
+        self.bilinear = torch.nn.Bilinear(in_features, in_features, out_features, bias=False)
+        self.linear = torch.nn.Linear(2 * in_features, out_features, bias=True)
+
+        self.reset_parameters()
+
+    def forward(self, x_1, x_2):
+        return self.bilinear(x_1, x_2) + self.linear(torch.cat((x_1, x_2), dim=-1))
+
+    def reset_parameters(self):
+        self.bilinear.reset_parameters()
+        self.linear.reset_parameters()
+
+
+class REDecoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.entity_emb = nn.Embedding(3, config.hidden_size, scale_grad_by_freq=True)
+        projection = nn.Sequential(
+            nn.Linear(config.hidden_size * 2, config.hidden_size),
+            nn.ReLU(),
+            nn.Dropout(config.hidden_dropout_prob),
+            nn.Linear(config.hidden_size, config.hidden_size // 2),
+            nn.ReLU(),
+            nn.Dropout(config.hidden_dropout_prob),
+        )
+        self.ffnn_head = copy.deepcopy(projection)
+        self.ffnn_tail = copy.deepcopy(projection)
+        self.rel_classifier = BiaffineAttention(config.hidden_size // 2, 2)
+        self.loss_fct = CrossEntropyLoss()
+
+    def build_relation(self, relations, entities):
+        batch_size = len(relations)
+        new_relations = []
+        for b in range(batch_size):
+            if len(entities[b]["start"]) <= 2:
+                entities[b] = {"end": [1, 1], "label": [0, 0], "start": [0, 0]}
+            all_possible_relations = set(
+                [
+                    (i, j)
+                    for i in range(len(entities[b]["label"]))
+                    for j in range(len(entities[b]["label"]))
+                    if entities[b]["label"][i] == 1 and entities[b]["label"][j] == 2
+                ]
+            )
+            if len(all_possible_relations) == 0:
+                all_possible_relations = set([(0, 1)])
+            positive_relations = set(list(zip(relations[b]["head"], relations[b]["tail"])))
+            negative_relations = all_possible_relations - positive_relations
+            positive_relations = set([i for i in positive_relations if i in all_possible_relations])
+            reordered_relations = list(positive_relations) + list(negative_relations)
+            relation_per_doc = {"head": [], "tail": [], "label": []}
+            relation_per_doc["head"] = [i[0] for i in reordered_relations]
+            relation_per_doc["tail"] = [i[1] for i in reordered_relations]
+            relation_per_doc["label"] = [1] * len(positive_relations) + [0] * (
+                len(reordered_relations) - len(positive_relations)
+            )
+            assert len(relation_per_doc["head"]) != 0
+            new_relations.append(relation_per_doc)
+        return new_relations, entities
+
+    def get_predicted_relations(self, logits, relations, entities):
+        pred_relations = []
+        for i, pred_label in enumerate(logits.argmax(-1)):
+            if pred_label != 1:
+                continue
+            rel = {}
+            rel["head_id"] = relations["head"][i]
+            rel["head"] = (entities["start"][rel["head_id"]], entities["end"][rel["head_id"]])
+            rel["head_type"] = entities["label"][rel["head_id"]]
+
+            rel["tail_id"] = relations["tail"][i]
+            rel["tail"] = (entities["start"][rel["tail_id"]], entities["end"][rel["tail_id"]])
+            rel["tail_type"] = entities["label"][rel["tail_id"]]
+            rel["type"] = 1
+            pred_relations.append(rel)
+        return pred_relations
+
+    def forward(self, hidden_states, entities, relations):
+        batch_size, max_n_words, context_dim = hidden_states.size()
+        device = hidden_states.device
+        relations, entities = self.build_relation(relations, entities)
+        loss = 0
+        all_pred_relations = []
+        for b in range(batch_size):
+            head_entities = torch.tensor(relations[b]["head"], device=device)
+            tail_entities = torch.tensor(relations[b]["tail"], device=device)
+            relation_labels = torch.tensor(relations[b]["label"], device=device)
+            entities_start_index = torch.tensor(entities[b]["start"], device=device)
+            entities_labels = torch.tensor(entities[b]["label"], device=device)
+            head_index = entities_start_index[head_entities]
+            head_label = entities_labels[head_entities]
+            head_label_repr = self.entity_emb(head_label)
+
+            tail_index = entities_start_index[tail_entities]
+            tail_label = entities_labels[tail_entities]
+            tail_label_repr = self.entity_emb(tail_label)
+
+            head_repr = torch.cat(
+                (hidden_states[b][head_index], head_label_repr),
+                dim=-1,
+            )
+            tail_repr = torch.cat(
+                (hidden_states[b][tail_index], tail_label_repr),
+                dim=-1,
+            )
+            heads = self.ffnn_head(head_repr)
+            tails = self.ffnn_tail(tail_repr)
+            logits = self.rel_classifier(heads, tails)
+            loss += self.loss_fct(logits, relation_labels)
+            pred_relations = self.get_predicted_relations(logits, relations[b], entities[b])
+            all_pred_relations.append(pred_relations)
+        return loss, all_pred_relations

xdoc/fine_tuning/funsd/layoutlmft/trainers/__init__.py

@@ -0,0 +1,2 @@
+from .funsd_trainer import FunsdTrainer
+from .xfun_trainer import XfunReTrainer, XfunSerTrainer

xdoc/fine_tuning/funsd/layoutlmft/trainers/funsd_trainer.py

@@ -0,0 +1,21 @@
+from typing import Any, Dict, Union
+
+import torch
+
+from transformers import Trainer
+
+
+class FunsdTrainer(Trainer):
+    def _prepare_inputs(self, inputs: Dict[str, Union[torch.Tensor, Any]]) -> Dict[str, Union[torch.Tensor, Any]]:
+        """
+        Prepare :obj:`inputs` before feeding them to the model, converting them to tensors if they are not already and
+        handling potential state.
+        """
+        for k, v in inputs.items():
+            if hasattr(v, "to") and hasattr(v, "device"):
+                inputs[k] = v.to(self.args.device)
+
+        if self.args.past_index >= 0 and self._past is not None:
+            inputs["mems"] = self._past
+
+        return inputs

xdoc/fine_tuning/funsd/layoutlmft/trainers/xfun_trainer.py

@@ -0,0 +1,197 @@
+import collections
+import time
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from packaging import version
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+
+from transformers.trainer_utils import EvalPrediction, PredictionOutput, speed_metrics
+from transformers.utils import logging
+
+from .funsd_trainer import FunsdTrainer
+
+
+if version.parse(torch.__version__) >= version.parse("1.6"):
+    _is_native_amp_available = True
+    from torch.cuda.amp import autocast
+
+logger = logging.get_logger(__name__)
+
+
+class XfunSerTrainer(FunsdTrainer):
+    pass
+
+
+class XfunReTrainer(FunsdTrainer):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.label_names.append("relations")
+
+    def prediction_step(
+        self,
+        model: nn.Module,
+        inputs: Dict[str, Union[torch.Tensor, Any]],
+        prediction_loss_only: bool,
+        ignore_keys: Optional[List[str]] = None,
+    ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        inputs = self._prepare_inputs(inputs)
+
+        with torch.no_grad():
+            if self.use_amp:
+                with autocast():
+                    outputs = model(**inputs)
+            else:
+                outputs = model(**inputs)
+        labels = tuple(inputs.get(name) for name in self.label_names)
+        return outputs, labels
+
+    def prediction_loop(
+        self,
+        dataloader: DataLoader,
+        description: str,
+        prediction_loss_only: Optional[bool] = None,
+        ignore_keys: Optional[List[str]] = None,
+        metric_key_prefix: str = "eval",
+    ) -> PredictionOutput:
+        """
+        Prediction/evaluation loop, shared by :obj:`Trainer.evaluate()` and :obj:`Trainer.predict()`.
+
+        Works both with or without labels.
+        """
+        if not isinstance(dataloader.dataset, collections.abc.Sized):
+            raise ValueError("dataset must implement __len__")
+        prediction_loss_only = (
+            prediction_loss_only if prediction_loss_only is not None else self.args.prediction_loss_only
+        )
+
+        if self.args.deepspeed and not self.args.do_train:
+            # no harm, but flagging to the user that deepspeed config is ignored for eval
+            # flagging only for when --do_train wasn't passed as only then it's redundant
+            logger.info("Detected the deepspeed argument but it will not be used for evaluation")
+
+        model = self._wrap_model(self.model, training=False)
+
+        # if full fp16 is wanted on eval and this ``evaluation`` or ``predict`` isn't called while
+        # ``train`` is running, half it first and then put on device
+        if not self.is_in_train and self.args.fp16_full_eval:
+            model = model.half().to(self.args.device)
+
+        batch_size = dataloader.batch_size
+        num_examples = self.num_examples(dataloader)
+        logger.info("***** Running %s *****", description)
+        logger.info("  Num examples = %d", num_examples)
+        logger.info("  Batch size = %d", batch_size)
+
+        model.eval()
+
+        self.callback_handler.eval_dataloader = dataloader
+
+        re_labels = None
+        pred_relations = None
+        entities = None
+        for step, inputs in enumerate(dataloader):
+            outputs, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
+            re_labels = labels[1] if re_labels is None else re_labels + labels[1]
+            pred_relations = (
+                outputs.pred_relations if pred_relations is None else pred_relations + outputs.pred_relations
+            )
+            entities = outputs.entities if entities is None else entities + outputs.entities
+
+            self.control = self.callback_handler.on_prediction_step(self.args, self.state, self.control)
+
+        gt_relations = []
+        for b in range(len(re_labels)):
+            rel_sent = []
+            for head, tail in zip(re_labels[b]["head"], re_labels[b]["tail"]):
+                rel = {}
+                rel["head_id"] = head
+                rel["head"] = (entities[b]["start"][rel["head_id"]], entities[b]["end"][rel["head_id"]])
+                rel["head_type"] = entities[b]["label"][rel["head_id"]]
+
+                rel["tail_id"] = tail
+                rel["tail"] = (entities[b]["start"][rel["tail_id"]], entities[b]["end"][rel["tail_id"]])
+                rel["tail_type"] = entities[b]["label"][rel["tail_id"]]
+
+                rel["type"] = 1
+
+                rel_sent.append(rel)
+
+            gt_relations.append(rel_sent)
+
+        re_metrics = self.compute_metrics(EvalPrediction(predictions=pred_relations, label_ids=gt_relations))
+
+        re_metrics = {
+            "precision": re_metrics["ALL"]["p"],
+            "recall": re_metrics["ALL"]["r"],
+            "f1": re_metrics["ALL"]["f1"],
+        }
+        re_metrics[f"{metric_key_prefix}_loss"] = outputs.loss.mean().item()
+
+        metrics = {}
+
+        # # Prefix all keys with metric_key_prefix + '_'
+        for key in list(re_metrics.keys()):
+            if not key.startswith(f"{metric_key_prefix}_"):
+                metrics[f"{metric_key_prefix}_{key}"] = re_metrics.pop(key)
+            else:
+                metrics[f"{key}"] = re_metrics.pop(key)
+
+        return metrics
+
+    def evaluate(
+        self,
+        eval_dataset: Optional[Dataset] = None,
+        ignore_keys: Optional[List[str]] = None,
+        metric_key_prefix: str = "eval",
+    ) -> Dict[str, float]:
+        """
+        Run evaluation and returns metrics.
+
+        The calling script will be responsible for providing a method to compute metrics, as they are task-dependent
+        (pass it to the init :obj:`compute_metrics` argument).
+
+        You can also subclass and override this method to inject custom behavior.
+
+        Args:
+            eval_dataset (:obj:`Dataset`, `optional`):
+                Pass a dataset if you wish to override :obj:`self.eval_dataset`. If it is an :obj:`datasets.Dataset`,
+                columns not accepted by the ``model.forward()`` method are automatically removed. It must implement the
+                :obj:`__len__` method.
+            ignore_keys (:obj:`Lst[str]`, `optional`):
+                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
+                gathering predictions.
+            metric_key_prefix (:obj:`str`, `optional`, defaults to :obj:`"eval"`):
+                An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
+                "eval_bleu" if the prefix is "eval" (default)
+
+        Returns:
+            A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The
+            dictionary also contains the epoch number which comes from the training state.
+        """
+        if eval_dataset is not None and not isinstance(eval_dataset, collections.abc.Sized):
+            raise ValueError("eval_dataset must implement __len__")
+
+        self.args.local_rank = -1
+        eval_dataloader = self.get_eval_dataloader(eval_dataset)
+        self.args.local_rank = torch.distributed.get_rank()
+
+        start_time = time.time()
+
+        metrics = self.prediction_loop(
+            eval_dataloader,
+            description="Evaluation",
+            # No point gathering the predictions if there are no metrics, otherwise we defer to
+            # self.args.prediction_loss_only
+            prediction_loss_only=True if self.compute_metrics is None else None,
+            ignore_keys=ignore_keys,
+            metric_key_prefix=metric_key_prefix,
+        )
+
+        n_samples = len(eval_dataset if eval_dataset is not None else self.eval_dataset)
+        metrics.update(speed_metrics(metric_key_prefix, start_time, n_samples))
+        self.log(metrics)
+        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, metrics)
+
+        return metrics

xdoc/fine_tuning/funsd/layoutlmft/utils.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple
+
+import torch
+
+from transformers.file_utils import ModelOutput
+
+
+@dataclass
+class ReOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    entities: Optional[Dict] = None
+    relations: Optional[Dict] = None
+    pred_relations: Optional[Dict] = None

xdoc/fine_tuning/funsd/model.py

@@ -0,0 +1,414 @@
+import logging
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers import BertConfig, BertModel, BertPreTrainedModel, RobertaConfig
+# from transformers.modeling_bert import BertLayerNorm, BertOnlyMLMHead
+
+logger = logging.getLogger(__name__)
+
+LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP = {}
+
+LAYOUTLMV1_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+
+
+class Layoutlmv1Config(RobertaConfig):
+    pretrained_config_archive_map = LAYOUTLMV1_PRETRAINED_CONFIG_ARCHIVE_MAP 
+    model_type = "bert"
+
+    def __init__(self, max_2d_position_embeddings=1024, add_linear=False, **kwargs):
+        super().__init__(**kwargs)
+        pass
+
+
+class Layoutlmv1Embeddings(nn.Module):
+    def __init__(self, config):
+        super(Layoutlmv1Embeddings, self).__init__()
+
+        self.config = config
+
+        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
+        )
+        config.max_2d_position_embeddings = 1024
+        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
+        )
+        self.token_type_embeddings = nn.Embedding(
+            config.type_vocab_size, config.hidden_size
+        )
+
+        self.LayerNorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        self.doc_linear1 = nn.Linear(config.hidden_size, config.hidden_size)
+        self.doc_linear2 = nn.Linear(config.hidden_size, config.hidden_size)
+        self.doc_linear3 = nn.Linear(config.hidden_size, config.hidden_size)
+        self.doc_linear4 = nn.Linear(config.hidden_size, config.hidden_size)
+
+        self.relu = nn.ReLU()
+
+    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)
+
+        words_embeddings = self.word_embeddings(input_ids)
+        position_embeddings = self.position_embeddings(position_ids)
+        token_type_embeddings = self.token_type_embeddings(token_type_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]
+        )
+
+        
+        temp_embeddings = self.doc_linear2(self.relu(self.doc_linear1(
+                    left_position_embeddings
+                    + upper_position_embeddings
+                    + right_position_embeddings
+                    + lower_position_embeddings
+                    + h_position_embeddings
+                    + w_position_embeddings
+                )))
+
+        embeddings = (
+            words_embeddings
+            + position_embeddings
+            + temp_embeddings
+            + token_type_embeddings
+        )
+
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class Layoutlmv1Model(BertModel):
+
+    config_class = Layoutlmv1Config
+    pretrained_model_archive_map = LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "bert"
+
+    def __init__(self, config):
+        super(Layoutlmv1Model, self).__init__(config)
+        self.embeddings = Layoutlmv1Embeddings(config)
+        self.init_weights()
+
+    def forward(
+        self,
+        input_ids,
+        bbox,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+    ):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+
+        # We create a 3D attention mask from a 2D tensor mask.
+        # Sizes are [batch_size, 1, 1, to_seq_length]
+        # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
+        # this attention mask is more simple than the triangular masking of causal attention
+        # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
+        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) 
+
+        # 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=torch.float32
+            # dtype=next(self.parameters()).dtype # this will trigger error when using high version torch
+        )  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        if head_mask is not None: 
+            if head_mask.dim() == 1:
+                head_mask = (
+                    head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+                )
+                head_mask = head_mask.expand(
+                    self.config.num_hidden_layers, -1, -1, -1, -1
+                )
+            elif head_mask.dim() == 2:
+                head_mask = (
+                    head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
+                )  # We can specify head_mask for each layer
+            head_mask = head_mask.to(
+                dtype=next(self.parameters()).dtype
+            )  # switch to fload if need + fp16 compatibility
+        else:
+            head_mask = [None] * self.config.num_hidden_layers
+
+        embedding_output = self.embeddings(
+            input_ids, bbox, position_ids=position_ids, token_type_ids=token_type_ids
+        )
+        encoder_outputs = self.encoder(
+            embedding_output, extended_attention_mask, head_mask=head_mask
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output)
+
+        outputs = (sequence_output, pooled_output) + encoder_outputs[
+            1:
+        ]  # add hidden_states and attentions if they are here
+        return outputs  # sequence_output, pooled_output, (hidden_states), (attentions)
+
+
+class Layoutlmv1ForTokenClassification(BertPreTrainedModel):
+    config_class = Layoutlmv1Config
+    pretrained_model_archive_map = LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "bert"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.roberta = Layoutlmv1Model(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    def forward(
+        self,
+        input_ids,
+        bbox,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+    ):
+
+        outputs = self.roberta(
+            input_ids=input_ids,
+            bbox=bbox,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+        )
+
+        sequence_output = outputs[0]
+
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        outputs = (logits,) + outputs[
+            2:
+        ]  # add hidden states and attention if they are here
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # Only keep active parts of the loss
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)[active_loss]
+                active_labels = labels.view(-1)[active_loss]
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), scores, (hidden_states), (attentions)
+
+
+class Layoutlmv1ForMaskedLM(BertPreTrainedModel):
+    config_class = Layoutlmv1Config
+    pretrained_model_archive_map = LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "bert"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = Layoutlmv1Model(config)
+        self.cls = BertOnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.bert.embeddings.word_embeddings
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def forward(
+        self,
+        input_ids,
+        bbox,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        masked_lm_labels=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        lm_labels=None,
+    ):
+
+        outputs = self.layoutlm(
+            input_ids,
+            bbox,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+
+        outputs = (prediction_scores,) + outputs[
+            2:
+        ]  # Add hidden states and attention if they are here
+
+        # Although this may seem awkward, BertForMaskedLM supports two scenarios:
+        # 1. If a tensor that contains the indices of masked labels is provided,
+        #    the cross-entropy is the MLM cross-entropy that measures the likelihood
+        #    of predictions for masked words.
+        # 2. If `lm_labels` is provided we are in a causal scenario where we
+        #    try to predict the next token for each input in the decoder.
+        if masked_lm_labels is not None:
+            loss_fct = CrossEntropyLoss()
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size),
+                masked_lm_labels.view(-1),
+            )
+            outputs = (masked_lm_loss,) + outputs
+        return (
+            outputs
+        )  # (masked_lm_loss), (ltr_lm_loss), prediction_scores, (hidden_states), (attentions)
+
+
+
+class Layoutlmv1ForQuestionAnswering(BertPreTrainedModel):
+    config_class = Layoutlmv1Config
+    pretrained_model_archive_map = LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "bert"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.bert = Layoutlmv1Model(config)
+        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    def forward(
+        self,
+        input_ids,
+        bbox,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        # inputs_embeds=None,
+        start_positions=None,
+        end_positions=None,
+        # output_attentions=None,
+        # output_hidden_states=None,
+        # return_dict=None,
+    ):
+        # import numpy as np
+        # torch.set_printoptions(threshold=np.inf)
+        # print(bbox[0])
+        # exit(0)
+        r"""
+        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        """
+        # return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.bert(
+            input_ids=input_ids,
+            bbox=bbox,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1)
+        end_logits = end_logits.squeeze(-1)
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions.clamp_(0, ignored_index)
+            end_positions.clamp_(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+
+        output = (start_logits, end_logits) + outputs[2:]
+
+
+        return ((total_loss,) + output) if total_loss is not None else output

xdoc/fine_tuning/funsd/requirements.txt

@@ -0,0 +1,7 @@
+torchvision==0.8 
+transformers==4.5.1
+scikit-learn==1.0.2
+datasets==2.6.1
+pillow==5.2.0
+torch
+seqeval

xdoc/fine_tuning/funsd/run_funsd.py

@@ -0,0 +1,363 @@
+#!/usr/bin/env python
+# coding=utf-8
+import logging
+import os
+import sys
+from dataclasses import dataclass, field
+from typing import Optional
+
+import numpy as np
+from datasets import ClassLabel, load_dataset, load_metric
+
+import layoutlmft.data.datasets.funsd
+import transformers
+from layoutlmft.data import DataCollatorForKeyValueExtraction
+from layoutlmft.data.data_args import DataTrainingArguments
+from layoutlmft.models.model_args import ModelArguments
+from layoutlmft.trainers import FunsdTrainer as Trainer
+from transformers import (
+    AutoConfig,
+    AutoModelForTokenClassification,
+    AutoTokenizer,
+    HfArgumentParser,
+    PreTrainedTokenizerFast,
+    TrainingArguments,
+    set_seed,
+    RobertaConfig
+)
+import torch
+from model import Layoutlmv1ForTokenClassification
+from transformers.trainer_utils import get_last_checkpoint, is_main_process
+from transformers.utils import check_min_version
+
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version("4.5.0")
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    training_args.per_device_train_batch_size = 16
+    training_args.num_train_epochs=100.0
+
+    # Detecting last checkpoint.
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        elif last_checkpoint is not None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+
+    # Setup logging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)
+
+    # Log on each process the small summary:
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+    # Set the verbosity to info of the Transformers logger (on main process only):
+    if is_main_process(training_args.local_rank):
+        transformers.utils.logging.set_verbosity_info()
+        transformers.utils.logging.enable_default_handler()
+        transformers.utils.logging.enable_explicit_format()
+    logger.info(f"Training/evaluation parameters {training_args}")
+
+    set_seed(training_args.seed) 
+
+    datasets = load_dataset(os.path.abspath(layoutlmft.data.datasets.funsd.__file__))
+    print(datasets)
+
+    if training_args.do_train:
+        column_names = datasets["train"].column_names
+        features = datasets["train"].features
+    else:
+        column_names = datasets["test"].column_names
+        features = datasets["test"].features
+    text_column_name = "tokens" if "tokens" in column_names else column_names[0]
+    label_column_name = (
+        f"{data_args.task_name}_tags" if f"{data_args.task_name}_tags" in column_names else column_names[1]
+    )
+
+    remove_columns = column_names
+
+    def get_label_list(labels):
+        unique_labels = set()
+        for label in labels:
+            unique_labels = unique_labels | set(label)
+        label_list = list(unique_labels)
+        label_list.sort()
+        return label_list
+
+    if isinstance(features[label_column_name].feature, ClassLabel):
+        label_list = features[label_column_name].feature.names
+        # No need to convert the labels since they are already ints.
+        label_to_id = {i: i for i in range(len(label_list))}
+    else:
+        label_list = get_label_list(datasets["train"][label_column_name])
+        label_to_id = {l: i for i, l in enumerate(label_list)}
+    num_labels = len(label_list)
+
+    config = RobertaConfig.from_pretrained(
+        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
+        num_labels=num_labels,
+        finetuning_task=data_args.task_name,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        use_fast=True,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+        add_prefix_space=True
+    )
+    # model = AutoModelForTokenClassification.from_pretrained(
+    model = Layoutlmv1ForTokenClassification.from_pretrained(
+        model_args.model_name_or_path,
+        from_tf=bool(".ckpt" in model_args.model_name_or_path),
+        config=config,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+
+    # Tokenizer check: this script requires a fast tokenizer.
+    if not isinstance(tokenizer, PreTrainedTokenizerFast):
+        raise ValueError(
+            "This example script only works for models that have a fast tokenizer. Checkout the big table of models "
+            "at https://huggingface.co/transformers/index.html#bigtable to find the model types that meet this "
+            "requirement"
+        )
+
+    # Preprocessing the dataset
+    # Padding strategy
+    padding = "max_length" if data_args.pad_to_max_length else False
+
+    # Tokenize all texts and align the labels with them.
+    def tokenize_and_align_labels(examples):
+        tokenized_inputs = tokenizer(
+            examples[text_column_name],
+            padding=padding,
+            truncation=True,
+            return_overflowing_tokens=True,
+            # We use this argument because the texts in our dataset are lists of words (with a label for each word).
+            is_split_into_words=True,
+        )
+
+        labels = []
+        bboxes = []
+        
+        images = []
+        for batch_index in range(len(tokenized_inputs["input_ids"])):
+            word_ids = tokenized_inputs.word_ids(batch_index=batch_index)
+            org_batch_index = tokenized_inputs["overflow_to_sample_mapping"][batch_index]
+
+            label = examples[label_column_name][org_batch_index]
+            bbox = examples["bboxes"][org_batch_index]
+            image = examples["image"][org_batch_index]
+            previous_word_idx = None
+            label_ids = []
+            bbox_inputs = []
+            for word_idx in word_ids:
+                # Special tokens have a word id that is None. We set the label to -100 so they are automatically
+                # ignored in the loss function.
+                if word_idx is None:
+                    label_ids.append(-100)
+                    bbox_inputs.append([0, 0, 0, 0])
+                # We set the label for the first token of each word.
+                elif word_idx != previous_word_idx:
+                    label_ids.append(label_to_id[label[word_idx]])
+                    bbox_inputs.append(bbox[word_idx])
+                # For the other tokens in a word, we set the label to either the current label or -100, depending on
+                # the label_all_tokens flag.
+                else:
+                    label_ids.append(label_to_id[label[word_idx]] if data_args.label_all_tokens else -100)
+                    bbox_inputs.append(bbox[word_idx])
+                previous_word_idx = word_idx
+            labels.append(label_ids)
+            bboxes.append(bbox_inputs)
+            images.append(image)
+
+        tokenized_inputs["labels"] = labels
+        tokenized_inputs["bbox"] = bboxes
+        tokenized_inputs["image"] = images
+        return tokenized_inputs
+
+    if training_args.do_train:
+        if "train" not in datasets:
+            raise ValueError("--do_train requires a train dataset")
+        train_dataset = datasets["train"]
+        if data_args.max_train_samples is not None:
+            train_dataset = train_dataset.select(range(data_args.max_train_samples))
+        train_dataset = train_dataset.map(
+            tokenize_and_align_labels,
+            batched=True,
+            remove_columns=remove_columns,
+            num_proc=data_args.preprocessing_num_workers,
+            load_from_cache_file=not data_args.overwrite_cache,
+        )
+
+    if training_args.do_eval:
+        if "test" not in datasets:
+            raise ValueError("--do_eval requires a validation dataset")
+        eval_dataset = datasets["test"]
+        if data_args.max_val_samples is not None:
+            eval_dataset = eval_dataset.select(range(data_args.max_val_samples))
+        eval_dataset = eval_dataset.map(
+            tokenize_and_align_labels,
+            batched=True,
+            remove_columns=remove_columns,
+            num_proc=data_args.preprocessing_num_workers,
+            load_from_cache_file=not data_args.overwrite_cache,
+        )
+
+    if training_args.do_predict:
+        if "test" not in datasets:
+            raise ValueError("--do_predict requires a test dataset")
+        test_dataset = datasets["test"]
+        if data_args.max_test_samples is not None:
+            test_dataset = test_dataset.select(range(data_args.max_test_samples))
+        test_dataset = test_dataset.map(
+            tokenize_and_align_labels,
+            batched=True,
+            remove_columns=remove_columns,
+            num_proc=data_args.preprocessing_num_workers,
+            load_from_cache_file=not data_args.overwrite_cache,
+        )
+
+    # Data collator
+    data_collator = DataCollatorForKeyValueExtraction(
+        tokenizer,
+        pad_to_multiple_of=8 if training_args.fp16 else None,
+        padding=padding,
+        max_length=512,
+    )
+
+    # Metrics
+    metric = load_metric("seqeval")
+
+    def compute_metrics(p):
+        predictions, labels = p
+        predictions = np.argmax(predictions, axis=2)
+
+        # Remove ignored index (special tokens)
+        true_predictions = [
+            [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
+            for prediction, label in zip(predictions, labels)
+        ]
+        true_labels = [
+            [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
+            for prediction, label in zip(predictions, labels)
+        ]
+
+        results = metric.compute(predictions=true_predictions, references=true_labels)
+        if data_args.return_entity_level_metrics:
+            # Unpack nested dictionaries
+            final_results = {}
+            for key, value in results.items():
+                if isinstance(value, dict):
+                    for n, v in value.items():
+                        final_results[f"{key}_{n}"] = v
+                else:
+                    final_results[key] = value
+            return final_results
+        else:
+            return {
+                "precision": results["overall_precision"],
+                "recall": results["overall_recall"],
+                "f1": results["overall_f1"],
+                "accuracy": results["overall_accuracy"],
+            }
+
+    # Initialize our Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset if training_args.do_train else None,
+        eval_dataset=eval_dataset if training_args.do_eval else None,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+    )
+
+    # Training
+    if training_args.do_train:
+        checkpoint = last_checkpoint if last_checkpoint else None
+        train_result = trainer.train(resume_from_checkpoint=checkpoint)
+        metrics = train_result.metrics
+        trainer.save_model()  # Saves the tokenizer too for easy upload
+
+        max_train_samples = (
+            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
+        )
+        metrics["train_samples"] = min(max_train_samples, len(train_dataset))
+
+        trainer.log_metrics("train", metrics)
+        trainer.save_metrics("train", metrics)
+        trainer.save_state()
+
+    # Evaluation
+    if training_args.do_eval:
+        logger.info("*** Evaluate ***")
+
+        metrics = trainer.evaluate()
+
+        max_val_samples = data_args.max_val_samples if data_args.max_val_samples is not None else len(eval_dataset)
+        metrics["eval_samples"] = min(max_val_samples, len(eval_dataset))
+
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+    # Predict
+    if training_args.do_predict:
+        logger.info("*** Predict ***")
+
+        predictions, labels, metrics = trainer.predict(test_dataset)
+        predictions = np.argmax(predictions, axis=2)
+
+        # Remove ignored index (special tokens)
+        true_predictions = [
+            [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
+            for prediction, label in zip(predictions, labels)
+        ]
+
+        trainer.log_metrics("test", metrics)
+        trainer.save_metrics("test", metrics)
+
+        # Save predictions
+        output_test_predictions_file = os.path.join(training_args.output_dir, "test_predictions.txt")
+        if trainer.is_world_process_zero():
+            with open(output_test_predictions_file, "w") as writer:
+                for prediction in true_predictions:
+                    writer.write(" ".join(prediction) + "\n")
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()

xdoc/fine_tuning/funsd/run_funsd.sh

@@ -0,0 +1,10 @@
+CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 --master_port 5678 run_funsd.py \
+        --model_name_or_path /path/to/xdoc-pretrain-roberta-1M \
+        --output_dir camera_ready_funsd_1M \
+        --do_train \
+        --do_eval \
+        --max_steps 1000 \
+        --warmup_ratio 0.1 \
+        --fp16 \
+        --overwrite_output_dir \
+        --seed 42

xdoc/fine_tuning/squad/requirements.txt

@@ -0,0 +1,6 @@
+sklearn
+transformers==4.23.1
+datasets==2.6.1
+numpy==1.21.6
+torch
+IPython

xdoc/fine_tuning/squad/run_squad.py

@@ -0,0 +1,658 @@
+import logging
+import os
+os.environ['DISABLE_MLFLOW_INTEGRATION'] = 'True'
+import sys
+from dataclasses import dataclass, field
+from typing import Optional
+
+import datasets
+from datasets import load_dataset, load_metric
+
+import transformers
+from trainer_qa import QuestionAnsweringTrainer
+from transformers import (
+    AutoConfig,
+    AutoModelForQuestionAnswering,
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    EvalPrediction,
+    HfArgumentParser,
+    PreTrainedTokenizerFast,
+    TrainingArguments,
+    default_data_collator,
+    set_seed,
+)
+from transformers.trainer_utils import get_last_checkpoint
+from transformers.utils import check_min_version
+from transformers.utils.versions import require_version
+from utils_qa import postprocess_qa_predictions
+
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version("4.20.0.dev0")
+
+require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt")
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to directory to store the pretrained models downloaded from huggingface.co"},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    use_auth_token: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Will use the token generated when running `transformers-cli login` (necessary to use this script "
+                "with private models)."
+            )
+        },
+    )
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    dataset_name: Optional[str] = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
+    validation_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
+    )
+    test_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input test data file to evaluate the perplexity on (a text file)."},
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    max_seq_length: int = field(
+        default=384,
+        metadata={
+            "help": (
+                "The maximum total input sequence length after tokenization. Sequences longer "
+                "than this will be truncated, sequences shorter will be padded."
+            )
+        },
+    )
+    pad_to_max_length: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to pad all samples to `max_seq_length`. If False, will pad the samples dynamically when"
+                " batching to the maximum length in the batch (which can be faster on GPU but will be slower on TPU)."
+            )
+        },
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_predict_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of prediction examples to this "
+                "value if set."
+            )
+        },
+    )
+    version_2_with_negative: bool = field(
+        default=False, metadata={"help": "If true, some of the examples do not have an answer."}
+    )
+    null_score_diff_threshold: float = field(
+        default=0.0,
+        metadata={
+            "help": (
+                "The threshold used to select the null answer: if the best answer has a score that is less than "
+                "the score of the null answer minus this threshold, the null answer is selected for this example. "
+                "Only useful when `version_2_with_negative=True`."
+            )
+        },
+    )
+    doc_stride: int = field(
+        default=128,
+        metadata={"help": "When splitting up a long document into chunks, how much stride to take between chunks."},
+    )
+    n_best_size: int = field(
+        default=20,
+        metadata={"help": "The total number of n-best predictions to generate when looking for an answer."},
+    )
+    max_answer_length: int = field(
+        default=30,
+        metadata={
+            "help": (
+                "The maximum length of an answer that can be generated. This is needed because the start "
+                "and end predictions are not conditioned on one another."
+            )
+        },
+    )
+
+    def __post_init__(self):
+        if (
+            self.dataset_name is None
+            and self.train_file is None
+            and self.validation_file is None
+            and self.test_file is None
+        ):
+            raise ValueError("Need either a dataset name or a training/validation file/test_file.")
+        else:
+            if self.train_file is not None:
+                extension = self.train_file.split(".")[-1]
+                assert extension in ["csv", "json"], "`train_file` should be a csv or a json file."
+            if self.validation_file is not None:
+                extension = self.validation_file.split(".")[-1]
+                assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
+            if self.test_file is not None:
+                extension = self.test_file.split(".")[-1]
+                assert extension in ["csv", "json"], "`test_file` should be a csv or a json file."
+
+
+def main():
+    # See all possible arguments in src/transformers/training_args.py
+    # or by passing the --help flag to this script.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # Setup logging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+
+    log_level = training_args.get_process_log_level()
+    logger.setLevel(log_level)
+    datasets.utils.logging.set_verbosity(log_level)
+    transformers.utils.logging.set_verbosity(log_level)
+    transformers.utils.logging.enable_default_handler()
+    transformers.utils.logging.enable_explicit_format()
+
+    # Log on each process the small summary:
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+    logger.info(f"Training/evaluation parameters {training_args}")
+
+    # Detecting last checkpoint.
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+
+    # Set seed before initializing model.
+    # set_seed(training_args.seed)
+    set_seed(training_args.seed)
+
+    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+    #
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
+    #
+    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+    # download the dataset.
+    if data_args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        raw_datasets = load_dataset(
+            data_args.dataset_name,
+            data_args.dataset_config_name,
+            cache_dir=model_args.cache_dir,
+            use_auth_token=True if model_args.use_auth_token else None,
+        )
+    else:
+        data_files = {}
+        if data_args.train_file is not None:
+            data_files["train"] = data_args.train_file
+            extension = data_args.train_file.split(".")[-1]
+
+        if data_args.validation_file is not None:
+            data_files["validation"] = data_args.validation_file
+            extension = data_args.validation_file.split(".")[-1]
+        if data_args.test_file is not None:
+            data_files["test"] = data_args.test_file
+            extension = data_args.test_file.split(".")[-1]
+        raw_datasets = load_dataset(
+            extension,
+            data_files=data_files,
+            field="data",
+            cache_dir=model_args.cache_dir,
+            use_auth_token=True if model_args.use_auth_token else None,
+        )
+    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    # Load pretrained model and tokenizer
+    #
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+    config = AutoConfig.from_pretrained(
+        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        use_fast=True,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+    model = AutoModelForQuestionAnswering.from_pretrained(
+        model_args.model_name_or_path,
+        from_tf=bool(".ckpt" in model_args.model_name_or_path),
+        config=config,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+
+    # Tokenizer check: this script requires a fast tokenizer.
+    if not isinstance(tokenizer, PreTrainedTokenizerFast):
+        raise ValueError(
+            "This example script only works for models that have a fast tokenizer. Checkout the big table of models at"
+            " https://huggingface.co/transformers/index.html#supported-frameworks to find the model types that meet"
+            " this requirement"
+        )
+
+    # Preprocessing the datasets.
+    # Preprocessing is slighlty different for training and evaluation.
+    if training_args.do_train:
+        column_names = raw_datasets["train"].column_names
+    elif training_args.do_eval:
+        column_names = raw_datasets["validation"].column_names
+    else:
+        column_names = raw_datasets["test"].column_names
+    question_column_name = "question" if "question" in column_names else column_names[0]
+    context_column_name = "context" if "context" in column_names else column_names[1]
+    answer_column_name = "answers" if "answers" in column_names else column_names[2]
+
+    # Padding side determines if we do (question|context) or (context|question).
+    pad_on_right = tokenizer.padding_side == "right"
+
+    if data_args.max_seq_length > tokenizer.model_max_length:
+        logger.warning(
+            f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
+            f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
+        )
+    max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)
+
+    # Training preprocessing
+    def prepare_train_features(examples):
+        # Some of the questions have lots of whitespace on the left, which is not useful and will make the
+        # truncation of the context fail (the tokenized question will take a lots of space). So we remove that
+        # left whitespace
+        examples[question_column_name] = [q.lstrip() for q in examples[question_column_name]]
+
+        # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results
+        # in one example possible giving several features when a context is long, each of those features having a
+        # context that overlaps a bit the context of the previous feature.
+        tokenized_examples = tokenizer(
+            examples[question_column_name if pad_on_right else context_column_name],
+            examples[context_column_name if pad_on_right else question_column_name],
+            truncation="only_second" if pad_on_right else "only_first",
+            max_length=max_seq_length,
+            stride=data_args.doc_stride,
+            return_overflowing_tokens=True,
+            return_offsets_mapping=True,
+            padding="max_length" if data_args.pad_to_max_length else False,
+        )
+
+        # Since one example might give us several features if it has a long context, we need a map from a feature to
+        # its corresponding example. This key gives us just that.
+        sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
+        # The offset mappings will give us a map from token to character position in the original context. This will
+        # help us compute the start_positions and end_positions.
+        offset_mapping = tokenized_examples.pop("offset_mapping")
+
+        # Let's label those examples!
+        tokenized_examples["start_positions"] = []
+        tokenized_examples["end_positions"] = []
+
+        for i, offsets in enumerate(offset_mapping):
+            # We will label impossible answers with the index of the CLS token.
+            input_ids = tokenized_examples["input_ids"][i]
+            cls_index = input_ids.index(tokenizer.cls_token_id)
+
+            # Grab the sequence corresponding to that example (to know what is the context and what is the question).
+            sequence_ids = tokenized_examples.sequence_ids(i)
+
+            # One example can give several spans, this is the index of the example containing this span of text.
+            sample_index = sample_mapping[i]
+            answers = examples[answer_column_name][sample_index]
+            # If no answers are given, set the cls_index as answer.
+            if len(answers["answer_start"]) == 0:
+                tokenized_examples["start_positions"].append(cls_index)
+                tokenized_examples["end_positions"].append(cls_index)
+            else:
+                # Start/end character index of the answer in the text.
+                start_char = answers["answer_start"][0]
+                end_char = start_char + len(answers["text"][0])
+
+                # Start token index of the current span in the text.
+                token_start_index = 0
+                while sequence_ids[token_start_index] != (1 if pad_on_right else 0):
+                    token_start_index += 1
+
+                # End token index of the current span in the text.
+                token_end_index = len(input_ids) - 1
+                while sequence_ids[token_end_index] != (1 if pad_on_right else 0):
+                    token_end_index -= 1
+
+                # Detect if the answer is out of the span (in which case this feature is labeled with the CLS index).
+                if not (offsets[token_start_index][0] <= start_char and offsets[token_end_index][1] >= end_char):
+                    tokenized_examples["start_positions"].append(cls_index)
+                    tokenized_examples["end_positions"].append(cls_index)
+                else:
+                    # Otherwise move the token_start_index and token_end_index to the two ends of the answer.
+                    # Note: we could go after the last offset if the answer is the last word (edge case).
+                    while token_start_index < len(offsets) and offsets[token_start_index][0] <= start_char:
+                        token_start_index += 1
+                    tokenized_examples["start_positions"].append(token_start_index - 1)
+                    while offsets[token_end_index][1] >= end_char:
+                        token_end_index -= 1
+                    tokenized_examples["end_positions"].append(token_end_index + 1)
+
+        return tokenized_examples
+
+    if training_args.do_train:
+        if "train" not in raw_datasets:
+            raise ValueError("--do_train requires a train dataset")
+        train_dataset = raw_datasets["train"]
+        if data_args.max_train_samples is not None:
+            # We will select sample from whole data if argument is specified
+            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
+            train_dataset = train_dataset.select(range(max_train_samples))
+        # Create train feature from dataset
+        with training_args.main_process_first(desc="train dataset map pre-processing"):
+            train_dataset = train_dataset.map(
+                prepare_train_features,
+                batched=True,
+                num_proc=data_args.preprocessing_num_workers,
+                remove_columns=column_names,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc="Running tokenizer on train dataset",
+            )
+        if data_args.max_train_samples is not None:
+            # Number of samples might increase during Feature Creation, We select only specified max samples
+            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
+            train_dataset = train_dataset.select(range(max_train_samples))
+
+    # Validation preprocessing
+    def prepare_validation_features(examples):
+        # Some of the questions have lots of whitespace on the left, which is not useful and will make the
+        # truncation of the context fail (the tokenized question will take a lots of space). So we remove that
+        # left whitespace
+        examples[question_column_name] = [q.lstrip() for q in examples[question_column_name]]
+
+        # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results
+        # in one example possible giving several features when a context is long, each of those features having a
+        # context that overlaps a bit the context of the previous feature.
+        tokenized_examples = tokenizer(
+            examples[question_column_name if pad_on_right else context_column_name],
+            examples[context_column_name if pad_on_right else question_column_name],
+            truncation="only_second" if pad_on_right else "only_first",
+            max_length=max_seq_length,
+            stride=data_args.doc_stride,
+            return_overflowing_tokens=True,
+            return_offsets_mapping=True,
+            padding="max_length" if data_args.pad_to_max_length else False,
+        )
+
+        # Since one example might give us several features if it has a long context, we need a map from a feature to
+        # its corresponding example. This key gives us just that.
+        sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
+
+        # For evaluation, we will need to convert our predictions to substrings of the context, so we keep the
+        # corresponding example_id and we will store the offset mappings.
+        tokenized_examples["example_id"] = []
+
+        for i in range(len(tokenized_examples["input_ids"])):
+            # Grab the sequence corresponding to that example (to know what is the context and what is the question).
+            sequence_ids = tokenized_examples.sequence_ids(i)
+            context_index = 1 if pad_on_right else 0
+
+            # One example can give several spans, this is the index of the example containing this span of text.
+            sample_index = sample_mapping[i]
+            tokenized_examples["example_id"].append(examples["id"][sample_index])
+
+            # Set to None the offset_mapping that are not part of the context so it's easy to determine if a token
+            # position is part of the context or not.
+            tokenized_examples["offset_mapping"][i] = [
+                (o if sequence_ids[k] == context_index else None)
+                for k, o in enumerate(tokenized_examples["offset_mapping"][i])
+            ]
+
+        return tokenized_examples
+
+    if training_args.do_eval:
+        if "validation" not in raw_datasets:
+            raise ValueError("--do_eval requires a validation dataset")
+        eval_examples = raw_datasets["validation"]
+        if data_args.max_eval_samples is not None:
+            # We will select sample from whole data
+            max_eval_samples = min(len(eval_examples), data_args.max_eval_samples)
+            eval_examples = eval_examples.select(range(max_eval_samples))
+        # Validation Feature Creation
+        with training_args.main_process_first(desc="validation dataset map pre-processing"):
+            eval_dataset = eval_examples.map(
+                prepare_validation_features,
+                batched=True,
+                num_proc=data_args.preprocessing_num_workers,
+                remove_columns=column_names,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc="Running tokenizer on validation dataset",
+            )
+        if data_args.max_eval_samples is not None:
+            # During Feature creation dataset samples might increase, we will select required samples again
+            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
+            eval_dataset = eval_dataset.select(range(max_eval_samples))
+
+    if training_args.do_predict:
+        if "test" not in raw_datasets:
+            raise ValueError("--do_predict requires a test dataset")
+        predict_examples = raw_datasets["test"]
+        if data_args.max_predict_samples is not None:
+            # We will select sample from whole data
+            predict_examples = predict_examples.select(range(data_args.max_predict_samples))
+        # Predict Feature Creation
+        with training_args.main_process_first(desc="prediction dataset map pre-processing"):
+            predict_dataset = predict_examples.map(
+                prepare_validation_features,
+                batched=True,
+                num_proc=data_args.preprocessing_num_workers,
+                remove_columns=column_names,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc="Running tokenizer on prediction dataset",
+            )
+        if data_args.max_predict_samples is not None:
+            # During Feature creation dataset samples might increase, we will select required samples again
+            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
+            predict_dataset = predict_dataset.select(range(max_predict_samples))
+
+    # Data collator
+    # We have already padded to max length if the corresponding flag is True, otherwise we need to pad in the data
+    # collator.
+    data_collator = (
+        default_data_collator
+        if data_args.pad_to_max_length
+        else DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if training_args.fp16 else None)
+    )
+
+    # Post-processing:
+    def post_processing_function(examples, features, predictions, stage="eval"):
+        # Post-processing: we match the start logits and end logits to answers in the original context.
+        predictions = postprocess_qa_predictions(
+            examples=examples,
+            features=features,
+            predictions=predictions,
+            version_2_with_negative=data_args.version_2_with_negative,
+            n_best_size=data_args.n_best_size,
+            max_answer_length=data_args.max_answer_length,
+            null_score_diff_threshold=data_args.null_score_diff_threshold,
+            output_dir=training_args.output_dir,
+            log_level=log_level,
+            prefix=stage,
+        )
+        # Format the result to the format the metric expects.
+        if data_args.version_2_with_negative:
+            formatted_predictions = [
+                {"id": k, "prediction_text": v, "no_answer_probability": 0.0} for k, v in predictions.items()
+            ]
+        else:
+            formatted_predictions = [{"id": k, "prediction_text": v} for k, v in predictions.items()]
+
+        references = [{"id": ex["id"], "answers": ex[answer_column_name]} for ex in examples]
+        return EvalPrediction(predictions=formatted_predictions, label_ids=references)
+
+    metric = load_metric("squad_v2" if data_args.version_2_with_negative else "squad")
+
+    def compute_metrics(p: EvalPrediction):
+        return metric.compute(predictions=p.predictions, references=p.label_ids)
+
+    # Initialize our Trainer
+    trainer = QuestionAnsweringTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset if training_args.do_train else None,
+        eval_dataset=eval_dataset if training_args.do_eval else None,
+        eval_examples=eval_examples if training_args.do_eval else None,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        post_process_function=post_processing_function,
+        compute_metrics=compute_metrics,
+    )
+
+    # Training
+    if training_args.do_train:
+        checkpoint = None
+        if training_args.resume_from_checkpoint is not None:
+            checkpoint = training_args.resume_from_checkpoint
+        elif last_checkpoint is not None:
+            checkpoint = last_checkpoint
+        train_result = trainer.train(resume_from_checkpoint=checkpoint)
+        trainer.save_model()  # Saves the tokenizer too for easy upload
+
+        metrics = train_result.metrics
+        max_train_samples = (
+            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
+        )
+        metrics["train_samples"] = min(max_train_samples, len(train_dataset))
+
+        trainer.log_metrics("train", metrics)
+        trainer.save_metrics("train", metrics)
+        trainer.save_state()
+
+    # Evaluation
+    if training_args.do_eval:
+        logger.info("*** Evaluate ***")
+        metrics = trainer.evaluate()
+
+        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
+        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
+
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+    # Prediction
+    if training_args.do_predict:
+        logger.info("*** Predict ***")
+        results = trainer.predict(predict_dataset, predict_examples)
+        metrics = results.metrics
+
+        max_predict_samples = (
+            data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
+        )
+        metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))
+
+        trainer.log_metrics("predict", metrics)
+        trainer.save_metrics("predict", metrics)
+
+    kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "question-answering"}
+    if data_args.dataset_name is not None:
+        kwargs["dataset_tags"] = data_args.dataset_name
+        if data_args.dataset_config_name is not None:
+            kwargs["dataset_args"] = data_args.dataset_config_name
+            kwargs["dataset"] = f"{data_args.dataset_name} {data_args.dataset_config_name}"
+        else:
+            kwargs["dataset"] = data_args.dataset_name
+
+    if training_args.push_to_hub:
+        trainer.push_to_hub(**kwargs)
+    else:
+        trainer.create_model_card(**kwargs)
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()

xdoc/fine_tuning/squad/run_squad_v1.sh

@@ -0,0 +1,12 @@
+CUDA_VISIBLE_DEVICES=0 python run_squad.py \
+  --model_name_or_path /path/to/xdoc-pretrain-roberta-1M \
+  --dataset_name squad \
+  --do_train \
+  --do_eval \
+  --per_device_train_batch_size 16 \
+  --learning_rate 3e-5 \
+  --num_train_epochs 2 \
+  --max_seq_length 384 \
+  --doc_stride 128 \
+  --output_dir ./squadv1.1_result \
+  --overwrite_output_dir

xdoc/fine_tuning/squad/run_squad_v2.sh

@@ -0,0 +1,13 @@
+CUDA_VISIBLE_DEVICES=0 python run_squad.py \
+  --model_name_or_path /path/to/xdoc-pretrain-roberta-1M \
+  --dataset_name squad_v2 \
+  --do_train \
+  --do_eval \
+  --version_2_with_negative \
+  --per_device_train_batch_size 16 \
+  --learning_rate 3e-5 \
+  --num_train_epochs 4 \
+  --max_seq_length 384 \
+  --doc_stride 128 \
+  --output_dir ./squadv2.0_result \
+  --overwrite_output_dir

xdoc/fine_tuning/squad/trainer_qa.py

@@ -0,0 +1,87 @@
+from transformers import Trainer, is_torch_tpu_available
+from transformers.trainer_utils import PredictionOutput
+
+
+if is_torch_tpu_available():
+    import torch_xla.core.xla_model as xm
+    import torch_xla.debug.metrics as met
+
+
+class QuestionAnsweringTrainer(Trainer):
+    def __init__(self, *args, eval_examples=None, post_process_function=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.eval_examples = eval_examples
+        self.post_process_function = post_process_function
+
+    def evaluate(self, eval_dataset=None, eval_examples=None, ignore_keys=None, metric_key_prefix: str = "eval"):
+        eval_dataset = self.eval_dataset if eval_dataset is None else eval_dataset
+        eval_dataloader = self.get_eval_dataloader(eval_dataset)
+        eval_examples = self.eval_examples if eval_examples is None else eval_examples
+
+        # Temporarily disable metric computation, we will do it in the loop here.
+        compute_metrics = self.compute_metrics
+        self.compute_metrics = None
+        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
+        try:
+            output = eval_loop(
+                eval_dataloader,
+                description="Evaluation",
+                # No point gathering the predictions if there are no metrics, otherwise we defer to
+                # self.args.prediction_loss_only
+                prediction_loss_only=True if compute_metrics is None else None,
+                ignore_keys=ignore_keys,
+            )
+        finally:
+            self.compute_metrics = compute_metrics
+
+        if self.post_process_function is not None and self.compute_metrics is not None:
+            eval_preds = self.post_process_function(eval_examples, eval_dataset, output.predictions)
+            metrics = self.compute_metrics(eval_preds)
+
+            # Prefix all keys with metric_key_prefix + '_'
+            for key in list(metrics.keys()):
+                if not key.startswith(f"{metric_key_prefix}_"):
+                    metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)
+
+            self.log(metrics)
+        else:
+            metrics = {}
+
+        if self.args.tpu_metrics_debug or self.args.debug:
+            # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
+            xm.master_print(met.metrics_report())
+
+        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, metrics)
+        return metrics
+
+    def predict(self, predict_dataset, predict_examples, ignore_keys=None, metric_key_prefix: str = "test"):
+        predict_dataloader = self.get_test_dataloader(predict_dataset)
+
+        # Temporarily disable metric computation, we will do it in the loop here.
+        compute_metrics = self.compute_metrics
+        self.compute_metrics = None
+        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
+        try:
+            output = eval_loop(
+                predict_dataloader,
+                description="Prediction",
+                # No point gathering the predictions if there are no metrics, otherwise we defer to
+                # self.args.prediction_loss_only
+                prediction_loss_only=True if compute_metrics is None else None,
+                ignore_keys=ignore_keys,
+            )
+        finally:
+            self.compute_metrics = compute_metrics
+
+        if self.post_process_function is None or self.compute_metrics is None:
+            return output
+
+        predictions = self.post_process_function(predict_examples, predict_dataset, output.predictions, "predict")
+        metrics = self.compute_metrics(predictions)
+
+        # Prefix all keys with metric_key_prefix + '_'
+        for key in list(metrics.keys()):
+            if not key.startswith(f"{metric_key_prefix}_"):
+                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)
+
+        return PredictionOutput(predictions=predictions.predictions, label_ids=predictions.label_ids, metrics=metrics)

xdoc/fine_tuning/squad/trainer_seq2seq_qa.py

@@ -0,0 +1,103 @@
+from typing import Dict, List, Optional
+
+from torch.utils.data import Dataset
+
+from transformers import Seq2SeqTrainer, is_torch_tpu_available
+from transformers.trainer_utils import PredictionOutput
+
+
+if is_torch_tpu_available():
+    import torch_xla.core.xla_model as xm
+    import torch_xla.debug.metrics as met
+
+
+class QuestionAnsweringSeq2SeqTrainer(Seq2SeqTrainer):
+    def __init__(self, *args, eval_examples=None, post_process_function=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.eval_examples = eval_examples
+        self.post_process_function = post_process_function
+
+    # def evaluate(self, eval_dataset=None, eval_examples=None, ignore_keys=None, metric_key_prefix: str = "eval"):
+    def evaluate(
+        self,
+        eval_dataset: Optional[Dataset] = None,
+        eval_examples=None,
+        ignore_keys: Optional[List[str]] = None,
+        metric_key_prefix: str = "eval",
+        max_length: Optional[int] = None,
+        num_beams: Optional[int] = None,
+    ) -> Dict[str, float]:
+        self._max_length = max_length if max_length is not None else self.args.generation_max_length
+        self._num_beams = num_beams if num_beams is not None else self.args.generation_num_beams
+
+        eval_dataset = self.eval_dataset if eval_dataset is None else eval_dataset
+        eval_dataloader = self.get_eval_dataloader(eval_dataset)
+        eval_examples = self.eval_examples if eval_examples is None else eval_examples
+
+        # Temporarily disable metric computation, we will do it in the loop here.
+        compute_metrics = self.compute_metrics
+        self.compute_metrics = None
+        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
+        try:
+            output = eval_loop(
+                eval_dataloader,
+                description="Evaluation",
+                # No point gathering the predictions if there are no metrics, otherwise we defer to
+                # self.args.prediction_loss_only
+                prediction_loss_only=True if compute_metrics is None else None,
+                ignore_keys=ignore_keys,
+            )
+        finally:
+            self.compute_metrics = compute_metrics
+
+        if self.post_process_function is not None and self.compute_metrics is not None:
+            eval_preds = self.post_process_function(eval_examples, eval_dataset, output)
+            metrics = self.compute_metrics(eval_preds)
+
+            # Prefix all keys with metric_key_prefix + '_'
+            for key in list(metrics.keys()):
+                if not key.startswith(f"{metric_key_prefix}_"):
+                    metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)
+
+            self.log(metrics)
+        else:
+            metrics = {}
+
+        if self.args.tpu_metrics_debug or self.args.debug:
+            # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
+            xm.master_print(met.metrics_report())
+
+        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, metrics)
+        return metrics
+
+    def predict(self, predict_dataset, predict_examples, ignore_keys=None, metric_key_prefix: str = "test"):
+        predict_dataloader = self.get_test_dataloader(predict_dataset)
+
+        # Temporarily disable metric computation, we will do it in the loop here.
+        compute_metrics = self.compute_metrics
+        self.compute_metrics = None
+        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
+        try:
+            output = eval_loop(
+                predict_dataloader,
+                description="Prediction",
+                # No point gathering the predictions if there are no metrics, otherwise we defer to
+                # self.args.prediction_loss_only
+                prediction_loss_only=True if compute_metrics is None else None,
+                ignore_keys=ignore_keys,
+            )
+        finally:
+            self.compute_metrics = compute_metrics
+
+        if self.post_process_function is None or self.compute_metrics is None:
+            return output
+
+        predictions = self.post_process_function(predict_examples, predict_dataset, output.predictions, "predict")
+        metrics = self.compute_metrics(predictions)
+
+        # Prefix all keys with metric_key_prefix + '_'
+        for key in list(metrics.keys()):
+            if not key.startswith(f"{metric_key_prefix}_"):
+                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)
+
+        return PredictionOutput(predictions=predictions.predictions, label_ids=predictions.label_ids, metrics=metrics)

xdoc/fine_tuning/squad/utils_qa.py

@@ -0,0 +1,426 @@
+import collections
+import json
+import logging
+import os
+from typing import Optional, Tuple
+
+import numpy as np
+from tqdm.auto import tqdm
+
+
+logger = logging.getLogger(__name__)
+
+
+def postprocess_qa_predictions(
+    examples,
+    features,
+    predictions: Tuple[np.ndarray, np.ndarray],
+    version_2_with_negative: bool = False,
+    n_best_size: int = 20,
+    max_answer_length: int = 30,
+    null_score_diff_threshold: float = 0.0,
+    output_dir: Optional[str] = None,
+    prefix: Optional[str] = None,
+    log_level: Optional[int] = logging.WARNING,
+):
+    """
+    Post-processes the predictions of a question-answering model to convert them to answers that are substrings of the
+    original contexts. This is the base postprocessing functions for models that only return start and end logits.
+
+    Args:
+        examples: The non-preprocessed dataset (see the main script for more information).
+        features: The processed dataset (see the main script for more information).
+        predictions (:obj:`Tuple[np.ndarray, np.ndarray]`):
+            The predictions of the model: two arrays containing the start logits and the end logits respectively. Its
+            first dimension must match the number of elements of :obj:`features`.
+        version_2_with_negative (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether or not the underlying dataset contains examples with no answers.
+        n_best_size (:obj:`int`, `optional`, defaults to 20):
+            The total number of n-best predictions to generate when looking for an answer.
+        max_answer_length (:obj:`int`, `optional`, defaults to 30):
+            The maximum length of an answer that can be generated. This is needed because the start and end predictions
+            are not conditioned on one another.
+        null_score_diff_threshold (:obj:`float`, `optional`, defaults to 0):
+            The threshold used to select the null answer: if the best answer has a score that is less than the score of
+            the null answer minus this threshold, the null answer is selected for this example (note that the score of
+            the null answer for an example giving several features is the minimum of the scores for the null answer on
+            each feature: all features must be aligned on the fact they `want` to predict a null answer).
+
+            Only useful when :obj:`version_2_with_negative` is :obj:`True`.
+        output_dir (:obj:`str`, `optional`):
+            If provided, the dictionaries of predictions, n_best predictions (with their scores and logits) and, if
+            :obj:`version_2_with_negative=True`, the dictionary of the scores differences between best and null
+            answers, are saved in `output_dir`.
+        prefix (:obj:`str`, `optional`):
+            If provided, the dictionaries mentioned above are saved with `prefix` added to their names.
+        log_level (:obj:`int`, `optional`, defaults to ``logging.WARNING``):
+            ``logging`` log level (e.g., ``logging.WARNING``)
+    """
+    if len(predictions) != 2:
+        raise ValueError("`predictions` should be a tuple with two elements (start_logits, end_logits).")
+    all_start_logits, all_end_logits = predictions
+
+    if len(predictions[0]) != len(features):
+        raise ValueError(f"Got {len(predictions[0])} predictions and {len(features)} features.")
+
+    # Build a map example to its corresponding features.
+    example_id_to_index = {k: i for i, k in enumerate(examples["id"])}
+    features_per_example = collections.defaultdict(list)
+    for i, feature in enumerate(features):
+        features_per_example[example_id_to_index[feature["example_id"]]].append(i)
+
+    # The dictionaries we have to fill.
+    all_predictions = collections.OrderedDict()
+    all_nbest_json = collections.OrderedDict()
+    if version_2_with_negative:
+        scores_diff_json = collections.OrderedDict()
+
+    # Logging.
+    logger.setLevel(log_level)
+    logger.info(f"Post-processing {len(examples)} example predictions split into {len(features)} features.")
+
+    # Let's loop over all the examples!
+    for example_index, example in enumerate(tqdm(examples)):
+        # Those are the indices of the features associated to the current example.
+        feature_indices = features_per_example[example_index]
+
+        min_null_prediction = None
+        prelim_predictions = []
+
+        # Looping through all the features associated to the current example.
+        for feature_index in feature_indices:
+            # We grab the predictions of the model for this feature.
+            start_logits = all_start_logits[feature_index]
+            end_logits = all_end_logits[feature_index]
+            # This is what will allow us to map some the positions in our logits to span of texts in the original
+            # context.
+            offset_mapping = features[feature_index]["offset_mapping"]
+            # Optional `token_is_max_context`, if provided we will remove answers that do not have the maximum context
+            # available in the current feature.
+            token_is_max_context = features[feature_index].get("token_is_max_context", None)
+
+            # Update minimum null prediction.
+            feature_null_score = start_logits[0] + end_logits[0]
+            if min_null_prediction is None or min_null_prediction["score"] > feature_null_score:
+                min_null_prediction = {
+                    "offsets": (0, 0),
+                    "score": feature_null_score,
+                    "start_logit": start_logits[0],
+                    "end_logit": end_logits[0],
+                }
+
+            # Go through all possibilities for the `n_best_size` greater start and end logits.
+            start_indexes = np.argsort(start_logits)[-1 : -n_best_size - 1 : -1].tolist()
+            end_indexes = np.argsort(end_logits)[-1 : -n_best_size - 1 : -1].tolist()
+            for start_index in start_indexes:
+                for end_index in end_indexes:
+                    # Don't consider out-of-scope answers, either because the indices are out of bounds or correspond
+                    # to part of the input_ids that are not in the context.
+                    if (
+                        start_index >= len(offset_mapping)
+                        or end_index >= len(offset_mapping)
+                        or offset_mapping[start_index] is None
+                        or len(offset_mapping[start_index]) < 2
+                        or offset_mapping[end_index] is None
+                        or len(offset_mapping[end_index]) < 2
+                    ):
+                        continue
+                    # Don't consider answers with a length that is either < 0 or > max_answer_length.
+                    if end_index < start_index or end_index - start_index + 1 > max_answer_length:
+                        continue
+                    # Don't consider answer that don't have the maximum context available (if such information is
+                    # provided).
+                    if token_is_max_context is not None and not token_is_max_context.get(str(start_index), False):
+                        continue
+
+                    prelim_predictions.append(
+                        {
+                            "offsets": (offset_mapping[start_index][0], offset_mapping[end_index][1]),
+                            "score": start_logits[start_index] + end_logits[end_index],
+                            "start_logit": start_logits[start_index],
+                            "end_logit": end_logits[end_index],
+                        }
+                    )
+        if version_2_with_negative and min_null_prediction is not None:
+            # Add the minimum null prediction
+            prelim_predictions.append(min_null_prediction)
+            null_score = min_null_prediction["score"]
+
+        # Only keep the best `n_best_size` predictions.
+        predictions = sorted(prelim_predictions, key=lambda x: x["score"], reverse=True)[:n_best_size]
+
+        # Add back the minimum null prediction if it was removed because of its low score.
+        if (
+            version_2_with_negative
+            and min_null_prediction is not None
+            and not any(p["offsets"] == (0, 0) for p in predictions)
+        ):
+            predictions.append(min_null_prediction)
+
+        # Use the offsets to gather the answer text in the original context.
+        context = example["context"]
+        for pred in predictions:
+            offsets = pred.pop("offsets")
+            pred["text"] = context[offsets[0] : offsets[1]]
+
+        # In the very rare edge case we have not a single non-null prediction, we create a fake prediction to avoid
+        # failure.
+        if len(predictions) == 0 or (len(predictions) == 1 and predictions[0]["text"] == ""):
+            predictions.insert(0, {"text": "empty", "start_logit": 0.0, "end_logit": 0.0, "score": 0.0})
+
+        # Compute the softmax of all scores (we do it with numpy to stay independent from torch/tf in this file, using
+        # the LogSumExp trick).
+        scores = np.array([pred.pop("score") for pred in predictions])
+        exp_scores = np.exp(scores - np.max(scores))
+        probs = exp_scores / exp_scores.sum()
+
+        # Include the probabilities in our predictions.
+        for prob, pred in zip(probs, predictions):
+            pred["probability"] = prob
+
+        # Pick the best prediction. If the null answer is not possible, this is easy.
+        if not version_2_with_negative:
+            all_predictions[example["id"]] = predictions[0]["text"]
+        else:
+            # Otherwise we first need to find the best non-empty prediction.
+            i = 0
+            while predictions[i]["text"] == "":
+                i += 1
+            best_non_null_pred = predictions[i]
+
+            # Then we compare to the null prediction using the threshold.
+            score_diff = null_score - best_non_null_pred["start_logit"] - best_non_null_pred["end_logit"]
+            scores_diff_json[example["id"]] = float(score_diff)  # To be JSON-serializable.
+            if score_diff > null_score_diff_threshold:
+                all_predictions[example["id"]] = ""
+            else:
+                all_predictions[example["id"]] = best_non_null_pred["text"]
+
+        # Make `predictions` JSON-serializable by casting np.float back to float.
+        all_nbest_json[example["id"]] = [
+            {k: (float(v) if isinstance(v, (np.float16, np.float32, np.float64)) else v) for k, v in pred.items()}
+            for pred in predictions
+        ]
+
+    # If we have an output_dir, let's save all those dicts.
+    if output_dir is not None:
+        if not os.path.isdir(output_dir):
+            raise EnvironmentError(f"{output_dir} is not a directory.")
+
+        prediction_file = os.path.join(
+            output_dir, "predictions.json" if prefix is None else f"{prefix}_predictions.json"
+        )
+        nbest_file = os.path.join(
+            output_dir, "nbest_predictions.json" if prefix is None else f"{prefix}_nbest_predictions.json"
+        )
+        if version_2_with_negative:
+            null_odds_file = os.path.join(
+                output_dir, "null_odds.json" if prefix is None else f"{prefix}_null_odds.json"
+            )
+
+        logger.info(f"Saving predictions to {prediction_file}.")
+        with open(prediction_file, "w") as writer:
+            writer.write(json.dumps(all_predictions, indent=4) + "\n")
+        logger.info(f"Saving nbest_preds to {nbest_file}.")
+        with open(nbest_file, "w") as writer:
+            writer.write(json.dumps(all_nbest_json, indent=4) + "\n")
+        if version_2_with_negative:
+            logger.info(f"Saving null_odds to {null_odds_file}.")
+            with open(null_odds_file, "w") as writer:
+                writer.write(json.dumps(scores_diff_json, indent=4) + "\n")
+
+    return all_predictions
+
+
+def postprocess_qa_predictions_with_beam_search(
+    examples,
+    features,
+    predictions: Tuple[np.ndarray, np.ndarray],
+    version_2_with_negative: bool = False,
+    n_best_size: int = 20,
+    max_answer_length: int = 30,
+    start_n_top: int = 5,
+    end_n_top: int = 5,
+    output_dir: Optional[str] = None,
+    prefix: Optional[str] = None,
+    log_level: Optional[int] = logging.WARNING,
+):
+    """
+    Post-processes the predictions of a question-answering model with beam search to convert them to answers that are substrings of the
+    original contexts. This is the postprocessing functions for models that return start and end logits, indices, as well as
+    cls token predictions.
+
+    Args:
+        examples: The non-preprocessed dataset (see the main script for more information).
+        features: The processed dataset (see the main script for more information).
+        predictions (:obj:`Tuple[np.ndarray, np.ndarray]`):
+            The predictions of the model: two arrays containing the start logits and the end logits respectively. Its
+            first dimension must match the number of elements of :obj:`features`.
+        version_2_with_negative (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether or not the underlying dataset contains examples with no answers.
+        n_best_size (:obj:`int`, `optional`, defaults to 20):
+            The total number of n-best predictions to generate when looking for an answer.
+        max_answer_length (:obj:`int`, `optional`, defaults to 30):
+            The maximum length of an answer that can be generated. This is needed because the start and end predictions
+            are not conditioned on one another.
+        start_n_top (:obj:`int`, `optional`, defaults to 5):
+            The number of top start logits too keep when searching for the :obj:`n_best_size` predictions.
+        end_n_top (:obj:`int`, `optional`, defaults to 5):
+            The number of top end logits too keep when searching for the :obj:`n_best_size` predictions.
+        output_dir (:obj:`str`, `optional`):
+            If provided, the dictionaries of predictions, n_best predictions (with their scores and logits) and, if
+            :obj:`version_2_with_negative=True`, the dictionary of the scores differences between best and null
+            answers, are saved in `output_dir`.
+        prefix (:obj:`str`, `optional`):
+            If provided, the dictionaries mentioned above are saved with `prefix` added to their names.
+        log_level (:obj:`int`, `optional`, defaults to ``logging.WARNING``):
+            ``logging`` log level (e.g., ``logging.WARNING``)
+    """
+    if len(predictions) != 5:
+        raise ValueError("`predictions` should be a tuple with five elements.")
+    start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits = predictions
+
+    if len(predictions[0]) != len(features):
+        raise ValueError(f"Got {len(predictions[0])} predictions and {len(features)} features.")
+
+    # Build a map example to its corresponding features.
+    example_id_to_index = {k: i for i, k in enumerate(examples["id"])}
+    features_per_example = collections.defaultdict(list)
+    for i, feature in enumerate(features):
+        features_per_example[example_id_to_index[feature["example_id"]]].append(i)
+
+    # The dictionaries we have to fill.
+    all_predictions = collections.OrderedDict()
+    all_nbest_json = collections.OrderedDict()
+    scores_diff_json = collections.OrderedDict() if version_2_with_negative else None
+
+    # Logging.
+    logger.setLevel(log_level)
+    logger.info(f"Post-processing {len(examples)} example predictions split into {len(features)} features.")
+
+    # Let's loop over all the examples!
+    for example_index, example in enumerate(tqdm(examples)):
+        # Those are the indices of the features associated to the current example.
+        feature_indices = features_per_example[example_index]
+
+        min_null_score = None
+        prelim_predictions = []
+
+        # Looping through all the features associated to the current example.
+        for feature_index in feature_indices:
+            # We grab the predictions of the model for this feature.
+            start_log_prob = start_top_log_probs[feature_index]
+            start_indexes = start_top_index[feature_index]
+            end_log_prob = end_top_log_probs[feature_index]
+            end_indexes = end_top_index[feature_index]
+            feature_null_score = cls_logits[feature_index]
+            # This is what will allow us to map some the positions in our logits to span of texts in the original
+            # context.
+            offset_mapping = features[feature_index]["offset_mapping"]
+            # Optional `token_is_max_context`, if provided we will remove answers that do not have the maximum context
+            # available in the current feature.
+            token_is_max_context = features[feature_index].get("token_is_max_context", None)
+
+            # Update minimum null prediction
+            if min_null_score is None or feature_null_score < min_null_score:
+                min_null_score = feature_null_score
+
+            # Go through all possibilities for the `n_start_top`/`n_end_top` greater start and end logits.
+            for i in range(start_n_top):
+                for j in range(end_n_top):
+                    start_index = int(start_indexes[i])
+                    j_index = i * end_n_top + j
+                    end_index = int(end_indexes[j_index])
+                    # Don't consider out-of-scope answers (last part of the test should be unnecessary because of the
+                    # p_mask but let's not take any risk)
+                    if (
+                        start_index >= len(offset_mapping)
+                        or end_index >= len(offset_mapping)
+                        or offset_mapping[start_index] is None
+                        or len(offset_mapping[start_index]) < 2
+                        or offset_mapping[end_index] is None
+                        or len(offset_mapping[end_index]) < 2
+                    ):
+                        continue
+
+                    # Don't consider answers with a length negative or > max_answer_length.
+                    if end_index < start_index or end_index - start_index + 1 > max_answer_length:
+                        continue
+                    # Don't consider answer that don't have the maximum context available (if such information is
+                    # provided).
+                    if token_is_max_context is not None and not token_is_max_context.get(str(start_index), False):
+                        continue
+                    prelim_predictions.append(
+                        {
+                            "offsets": (offset_mapping[start_index][0], offset_mapping[end_index][1]),
+                            "score": start_log_prob[i] + end_log_prob[j_index],
+                            "start_log_prob": start_log_prob[i],
+                            "end_log_prob": end_log_prob[j_index],
+                        }
+                    )
+
+        # Only keep the best `n_best_size` predictions.
+        predictions = sorted(prelim_predictions, key=lambda x: x["score"], reverse=True)[:n_best_size]
+
+        # Use the offsets to gather the answer text in the original context.
+        context = example["context"]
+        for pred in predictions:
+            offsets = pred.pop("offsets")
+            pred["text"] = context[offsets[0] : offsets[1]]
+
+        # In the very rare edge case we have not a single non-null prediction, we create a fake prediction to avoid
+        # failure.
+        if len(predictions) == 0:
+            # Without predictions min_null_score is going to be None and None will cause an exception later
+            min_null_score = -2e-6
+            predictions.insert(0, {"text": "", "start_logit": -1e-6, "end_logit": -1e-6, "score": min_null_score})
+
+        # Compute the softmax of all scores (we do it with numpy to stay independent from torch/tf in this file, using
+        # the LogSumExp trick).
+        scores = np.array([pred.pop("score") for pred in predictions])
+        exp_scores = np.exp(scores - np.max(scores))
+        probs = exp_scores / exp_scores.sum()
+
+        # Include the probabilities in our predictions.
+        for prob, pred in zip(probs, predictions):
+            pred["probability"] = prob
+
+        # Pick the best prediction and set the probability for the null answer.
+        all_predictions[example["id"]] = predictions[0]["text"]
+        if version_2_with_negative:
+            scores_diff_json[example["id"]] = float(min_null_score)
+
+        # Make `predictions` JSON-serializable by casting np.float back to float.
+        all_nbest_json[example["id"]] = [
+            {k: (float(v) if isinstance(v, (np.float16, np.float32, np.float64)) else v) for k, v in pred.items()}
+            for pred in predictions
+        ]
+
+    # If we have an output_dir, let's save all those dicts.
+    if output_dir is not None:
+        if not os.path.isdir(output_dir):
+            raise EnvironmentError(f"{output_dir} is not a directory.")
+
+        prediction_file = os.path.join(
+            output_dir, "predictions.json" if prefix is None else f"{prefix}_predictions.json"
+        )
+        nbest_file = os.path.join(
+            output_dir, "nbest_predictions.json" if prefix is None else f"{prefix}_nbest_predictions.json"
+        )
+        if version_2_with_negative:
+            null_odds_file = os.path.join(
+                output_dir, "null_odds.json" if prefix is None else f"{prefix}_null_odds.json"
+            )
+
+        logger.info(f"Saving predictions to {prediction_file}.")
+        with open(prediction_file, "w") as writer:
+            writer.write(json.dumps(all_predictions, indent=4) + "\n")
+        logger.info(f"Saving nbest_preds to {nbest_file}.")
+        with open(nbest_file, "w") as writer:
+            writer.write(json.dumps(all_nbest_json, indent=4) + "\n")
+        if version_2_with_negative:
+            logger.info(f"Saving null_odds to {null_odds_file}.")
+            with open(null_odds_file, "w") as writer:
+                writer.write(json.dumps(scores_diff_json, indent=4) + "\n")
+
+    return all_predictions, scores_diff_json

xdoc/fine_tuning/websrc/args.py

@@ -0,0 +1,65 @@
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--exp_name", default='your_exp_name', type=str)
+parser.add_argument("--seed", default=42, type=int)
+parser.add_argument("--output_dir", default='.', type=str) 
+parser.add_argument("--overwrite_output_dir", default=True)
+parser.add_argument("--model_name_or_path", type=str, default='/path/to/xdoc-pretrain-roberta-1M')
+parser.add_argument("--cache_dir", type=str, default='./cache')
+parser.add_argument("--scratch", action='store_true')
+parser.add_argument("--model_type", default='roberta', choices=['bert', 'roberta'])
+
+parser.add_argument("--do_train", default=False, type=bool)
+parser.add_argument("--do_eval", default=True, type=bool)
+parser.add_argument("--do_test", default=False, type=bool)
+
+parser.add_argument("--overwrite_cache", default=False, type=bool)
+parser.add_argument("--train_file", default='train.json', type=str)
+parser.add_argument("--dev_file", default='val.json', type=str)
+parser.add_argument("--test_file", default='test.json', type=str)
+parser.add_argument("--doc_stride", default=128, type=int)
+parser.add_argument("--threads", default=1, type=int)
+parser.add_argument("--max_query_length", default=64, type=int)
+parser.add_argument("--local_rank", default=-1, type=int)
+parser.add_argument("--pad_img_input", default=False, type=bool)
+parser.add_argument("--fix_visual", default=False, type=bool)
+parser.add_argument("--num_workers", default=8, type=int)
+
+# dataset (args to load websrc)
+parser.add_argument("--web_train_file", default='/path/to/WebSRC/websrc1.0_train_.json', type=str)
+parser.add_argument("--web_eval_file", default='/path/to/WebSRC/websrc1.0_dev_.json', type=str)
+parser.add_argument("--web_root_dir", default='/path/to/WebSRC', type=str)
+parser.add_argument("--root_dir", default='/path/to/WebSRC', type=str)
+
+parser.add_argument("--n_best_size", default=20, type=int)
+parser.add_argument("--max_answer_length", default=30, type=int)
+parser.add_argument("--do_lower_case", default=True, type=bool)
+
+parser.add_argument("--web_num_features", default=0, type=int)
+parser.add_argument("--web_save_features", default=True, type=bool)
+parser.add_argument("--verbose_logging", default=True, type=bool)
+parser.add_argument("--embedding_mode", choices=['html','box','html+box'], default='html', type=str)
+parser.add_argument("--dataloader_shuffle", default=True, type=bool)
+
+# train
+parser.add_argument("--batch_per_gpu", default=16, type=int)
+parser.add_argument("--epoch", default=5, type=int)
+parser.add_argument("--warmup_ratio", default=0.1, type=float)
+parser.add_argument("--weight_decay", default=0, type=float)
+parser.add_argument("--fp16", default=False)
+parser.add_argument("--fp16_opt_level", default='O1', type=str)
+parser.add_argument("--learning_rate", default=5e-5, type=float)
+parser.add_argument("--adam_epsilon", default=1e-8, type=float) 
+parser.add_argument("--max_grad_norm", default=1, type=float) 
+
+parser.add_argument("--log_step", default=50, type=int) 
+parser.add_argument("--save_step", default=10000, type=int)
+
+parser.add_argument("--add_linear", default=True, type=bool) 
+parser.add_argument("--accumulation", default=1, type=int)
+
+# mlm
+parser.add_argument("--mlm_probability", default=0.15, type=float)
+
+args = parser.parse_args()

xdoc/fine_tuning/websrc/model.py

@@ -0,0 +1,717 @@
+import logging
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers import BertConfig, BertModel, BertPreTrainedModel, RobertaConfig
+# from transformers.modeling_bert import BertLayerNorm, BertOnlyMLMHead
+from transformers.models.bert.modeling_bert import BertOnlyMLMHead
+BertLayerNorm = torch.nn.LayerNorm
+logger = logging.getLogger(__name__)
+
+LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP = {}
+
+LAYOUTLMV1_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+
+
+class Layoutlmv1Config_roberta(RobertaConfig):
+    pretrained_config_archive_map = LAYOUTLMV1_PRETRAINED_CONFIG_ARCHIVE_MAP 
+    model_type = "bert"
+
+    def __init__(self, max_2d_position_embeddings=1024, add_linear=False, **kwargs):
+        super().__init__(**kwargs)
+        self.max_2d_position_embeddings = max_2d_position_embeddings
+        self.add_linear = add_linear # determine whether to add an additional mapping
+
+
+
+class Layoutlmv1Config(BertConfig):
+    pretrained_config_archive_map = LAYOUTLMV1_PRETRAINED_CONFIG_ARCHIVE_MAP 
+    model_type = "bert"
+
+    def __init__(self, max_2d_position_embeddings=1024, add_linear=False, **kwargs):
+        super().__init__(**kwargs)
+        self.max_2d_position_embeddings = max_2d_position_embeddings
+        self.add_linear = add_linear # determine whether to add an additional mapping
+
+
+
+class WebConfig:
+    max_depth = 50
+    xpath_unit_hidden_size = 32
+    hidden_size = 768
+    hidden_dropout_prob = 0.1
+    layer_norm_eps = 1e-12
+    max_xpath_tag_unit_embeddings = 256
+    max_xpath_subs_unit_embeddings = 1024
+
+
+
+
+class XPathEmbeddings(nn.Module):
+    """Construct the embddings from xpath -- tag and subscript"""
+
+    # we drop tree-id in this version, as its info can be covered by xpath
+
+    def __init__(self, config):
+        super(XPathEmbeddings, self).__init__()
+        config = WebConfig()
+        self.max_depth = config.max_depth
+
+        self.xpath_unitseq2_embeddings = nn.Linear(
+            config.xpath_unit_hidden_size * self.max_depth, config.hidden_size)
+
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        self.activation = nn.ReLU()
+        self.xpath_unitseq2_inner = nn.Linear(config.xpath_unit_hidden_size * self.max_depth, 4 * config.hidden_size)
+        self.inner2emb = nn.Linear(4 * config.hidden_size, config.hidden_size)
+
+        self.xpath_tag_sub_embeddings = nn.ModuleList(
+            [nn.Embedding(config.max_xpath_tag_unit_embeddings, config.xpath_unit_hidden_size) for _ in
+             range(self.max_depth)])
+
+        self.xpath_subs_sub_embeddings = nn.ModuleList(
+            [nn.Embedding(config.max_xpath_subs_unit_embeddings, config.xpath_unit_hidden_size) for _ in
+             range(self.max_depth)])
+
+    def forward(self,
+                xpath_tags_seq=None,
+                xpath_subs_seq=None):
+        xpath_tags_embeddings = []
+        xpath_subs_embeddings = []
+
+        for i in range(self.max_depth):
+            xpath_tags_embeddings.append(self.xpath_tag_sub_embeddings[i](xpath_tags_seq[:, :, i]))
+            xpath_subs_embeddings.append(self.xpath_subs_sub_embeddings[i](xpath_subs_seq[:, :, i]))
+
+        xpath_tags_embeddings = torch.cat(xpath_tags_embeddings, dim=-1)
+        xpath_subs_embeddings = torch.cat(xpath_subs_embeddings, dim=-1)
+
+        xpath_embeddings = xpath_tags_embeddings + xpath_subs_embeddings
+
+        xpath_embeddings = self.inner2emb(
+            self.dropout(self.activation(self.xpath_unitseq2_inner(xpath_embeddings))))
+
+        return xpath_embeddings
+
+
+class Layoutlmv1Embeddings(nn.Module):
+    def __init__(self, config):
+        super(Layoutlmv1Embeddings, self).__init__()
+        self.config = config
+        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
+        )
+        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
+        )
+        self.token_type_embeddings = nn.Embedding(
+            config.type_vocab_size, config.hidden_size
+        )
+
+        # for web extension
+        self.xpath_embeddings = XPathEmbeddings(config)
+
+        # 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)
+
+        self.doc_linear1 = nn.Linear(config.hidden_size, config.hidden_size)
+        self.doc_linear2 = nn.Linear(config.hidden_size, config.hidden_size)
+
+        self.web_linear1 = nn.Linear(config.hidden_size, config.hidden_size)
+        self.web_linear2 = nn.Linear(config.hidden_size, config.hidden_size)
+        self.web_linear3 = nn.Linear(config.hidden_size, config.hidden_size)
+        self.web_linear4 = nn.Linear(config.hidden_size, config.hidden_size)
+
+        self.relu = nn.ReLU()
+
+    def forward(
+        self,
+        input_ids,
+        bbox=None,
+        xpath_tags_seq=None,
+        xpath_subs_seq=None,
+        token_type_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+        embedding_mode=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)
+
+        words_embeddings = self.word_embeddings(input_ids)
+        position_embeddings = self.position_embeddings(position_ids)
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        if embedding_mode != None and embedding_mode == 'box' : # doc entry
+           
+            bbox = torch.clamp(bbox, 0, self.config.max_2d_position_embeddings-1)
+            
+            left_position_embeddings = self.x_position_embeddings(bbox[:, :, 0])
+            upper_position_embeddings = self.y_position_embeddings(bbox[:, :, 1])
+           
+            embeddings = (
+                words_embeddings
+                + position_embeddings
+                + left_position_embeddings
+                + upper_position_embeddings
+                # + right_position_embeddings
+                # + lower_position_embeddings
+                # + h_position_embeddings
+                # + w_position_embeddings
+                + token_type_embeddings
+            )
+        elif embedding_mode != None and embedding_mode == 'html+box' : # doc entry
+          
+            bbox = torch.clamp(bbox, 0, self.config.max_2d_position_embeddings-1)
+            
+            left_position_embeddings = self.x_position_embeddings(bbox[:, :, 0])
+            upper_position_embeddings = self.y_position_embeddings(bbox[:, :, 1])
+            xpath_embeddings = self.xpath_embeddings(xpath_tags_seq, xpath_subs_seq)
+
+            embeddings = (
+                words_embeddings
+                + position_embeddings
+                + left_position_embeddings
+                + upper_position_embeddings
+                + xpath_embeddings
+                # + right_position_embeddings
+                # + lower_position_embeddings
+                # + h_position_embeddings
+                # + w_position_embeddings
+                + token_type_embeddings
+            )
+        else: # web entry
+            if not self.config.add_linear:
+                xpath_embeddings = self.xpath_embeddings(xpath_tags_seq, xpath_subs_seq)
+                embeddings = (
+                    words_embeddings
+                    + position_embeddings
+                    + token_type_embeddings
+                    + xpath_embeddings
+                )
+            else:
+                xpath_embeddings = self.xpath_embeddings(xpath_tags_seq, xpath_subs_seq)
+
+                temp_embeddings = self.web_linear2(self.relu(self.web_linear1(
+                    xpath_embeddings
+                )))
+                embeddings = (
+                    words_embeddings
+                    + position_embeddings
+                    + token_type_embeddings
+                    + temp_embeddings
+                )
+
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class Layoutlmv1Model(BertModel):
+
+    config_class = Layoutlmv1Config
+    pretrained_model_archive_map = LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "bert"
+
+    def __init__(self, config):
+        super(Layoutlmv1Model, self).__init__(config)
+        self.embeddings = Layoutlmv1Embeddings(config)
+        self.init_weights()
+
+    def forward(
+        self,
+        input_ids,
+        bbox=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        xpath_tags_seq=None,
+        xpath_subs_seq=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        embedding_mode=None,
+    ):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+
+        # We create a 3D attention mask from a 2D tensor mask.
+        # Sizes are [batch_size, 1, 1, to_seq_length]
+        # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
+        # this attention mask is more simple than the triangular masking of causal attention
+        # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
+        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) 
+        # 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=torch.float32
+            # dtype=next(self.parameters()).dtype # this will trigger error when using high version torch
+        )  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        if head_mask is not None: 
+            if head_mask.dim() == 1:
+                head_mask = (
+                    head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+                )
+                head_mask = head_mask.expand(
+                    self.config.num_hidden_layers, -1, -1, -1, -1
+                )
+            elif head_mask.dim() == 2:
+                head_mask = (
+                    head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
+                )  # We can specify head_mask for each layer
+            head_mask = head_mask.to(
+                dtype=next(self.parameters()).dtype
+            )  # switch to fload if need + fp16 compatibility
+        else:
+            head_mask = [None] * self.config.num_hidden_layers
+
+        embedding_output = self.embeddings(
+            input_ids, bbox=bbox, xpath_tags_seq=xpath_tags_seq, xpath_subs_seq=xpath_subs_seq, position_ids=position_ids, token_type_ids=token_type_ids, embedding_mode=embedding_mode
+        )
+        encoder_outputs = self.encoder(
+            embedding_output, extended_attention_mask, head_mask=head_mask
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output)
+
+        outputs = (sequence_output, pooled_output) + encoder_outputs[
+            1:
+        ]  # add hidden_states and attentions if they are here
+        return outputs  # sequence_output, pooled_output, (hidden_states), (attentions)
+
+
+class Layoutlmv1ForTokenClassification(BertPreTrainedModel):
+    config_class = Layoutlmv1Config
+    pretrained_model_archive_map = LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "bert"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.bert = Layoutlmv1Model(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    def forward(
+        self,
+        input_ids,
+        bbox=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+    ):
+
+        outputs = self.bert(
+            input_ids=input_ids,
+            bbox=bbox,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+        )
+
+        sequence_output = outputs[0]
+
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        outputs = (logits,) + outputs[
+            2:
+        ]  # add hidden states and attention if they are here
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # Only keep active parts of the loss
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)[active_loss]
+                active_labels = labels.view(-1)[active_loss]
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), scores, (hidden_states), (attentions)
+
+
+class Layoutlmv1ForMaskedLM(BertPreTrainedModel):
+    config_class = Layoutlmv1Config
+    pretrained_model_archive_map = LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "bert"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = Layoutlmv1Model(config)
+        self.cls = BertOnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.bert.embeddings.word_embeddings
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def forward(
+        self,
+        input_ids,
+        bbox=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        masked_lm_labels=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        lm_labels=None,
+        xpath_tags_seq=None,
+        xpath_subs_seq=None,
+    ):
+
+        outputs = self.bert(
+            input_ids,
+            bbox,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            xpath_tags_seq=xpath_tags_seq,
+            xpath_subs_seq=xpath_subs_seq,
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+
+        outputs = (prediction_scores,) + outputs[
+            2:
+        ]  # Add hidden states and attention if they are here
+
+        # Although this may seem awkward, BertForMaskedLM supports two scenarios:
+        # 1. If a tensor that contains the indices of masked labels is provided,
+        #    the cross-entropy is the MLM cross-entropy that measures the likelihood
+        #    of predictions for masked words.
+        # 2. If `lm_labels` is provided we are in a causal scenario where we
+        #    try to predict the next token for each input in the decoder.
+        if masked_lm_labels is not None:
+            loss_fct = CrossEntropyLoss()
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size),
+                masked_lm_labels.view(-1),
+            )
+            outputs = (masked_lm_loss,) + outputs
+        return (
+            outputs
+        )  # (masked_lm_loss), (ltr_lm_loss), prediction_scores, (hidden_states), (attentions)
+
+
+class Layoutlmv1ForMaskedLM_roberta(BertPreTrainedModel):
+    config_class = Layoutlmv1Config
+    pretrained_model_archive_map = LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "bert"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.roberta = Layoutlmv1Model(config)
+        self.cls = BertOnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.roberta.embeddings.word_embeddings
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def forward(
+        self,
+        input_ids,
+        bbox=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        masked_lm_labels=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        lm_labels=None,
+        xpath_tags_seq=None,
+        xpath_subs_seq=None,
+    ):
+
+        outputs = self.roberta(
+            input_ids,
+            bbox,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            xpath_tags_seq=xpath_tags_seq,
+            xpath_subs_seq=xpath_subs_seq,
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+
+        outputs = (prediction_scores,) + outputs[
+            2:
+        ]  # Add hidden states and attention if they are here
+
+        # Although this may seem awkward, BertForMaskedLM supports two scenarios:
+        # 1. If a tensor that contains the indices of masked labels is provided,
+        #    the cross-entropy is the MLM cross-entropy that measures the likelihood
+        #    of predictions for masked words.
+        # 2. If `lm_labels` is provided we are in a causal scenario where we
+        #    try to predict the next token for each input in the decoder.
+        if masked_lm_labels is not None:
+            loss_fct = CrossEntropyLoss()
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size),
+                masked_lm_labels.view(-1),
+            )
+            outputs = (masked_lm_loss,) + outputs
+        return (
+            outputs
+        )  # (masked_lm_loss), (ltr_lm_loss), prediction_scores, (hidden_states), (attentions)
+
+
+
+class Layoutlmv1ForQuestionAnswering(BertPreTrainedModel):
+    config_class = Layoutlmv1Config
+    pretrained_model_archive_map = LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "bert"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.bert = Layoutlmv1Model(config) 
+        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    def forward(
+        self,
+        input_ids,
+        bbox=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        # inputs_embeds=None,
+        start_positions=None,
+        end_positions=None,
+        # output_attentions=None,
+        # output_hidden_states=None,
+        # return_dict=None,
+        xpath_tags_seq=None,
+        xpath_subs_seq=None,
+        embedding_mode=None,
+    ):
+        r"""
+        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        """
+        # return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.bert(
+            input_ids=input_ids,
+            bbox=bbox,
+            xpath_tags_seq=xpath_tags_seq,
+            xpath_subs_seq=xpath_subs_seq,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            embedding_mode=embedding_mode
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1)
+        end_logits = end_logits.squeeze(-1)
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions.clamp_(0, ignored_index)
+            end_positions.clamp_(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        # if not return_dict:
+        #     output = (start_logits, end_logits) + outputs[2:]
+        #     return ((total_loss,) + output) if total_loss is not None else output
+        #
+        # return QuestionAnsweringModelOutput(
+        #     loss=total_loss,
+        #     start_logits=start_logits,
+        #     end_logits=end_logits,
+        #     hidden_states=outputs.hidden_states,
+        #     attentions=outputs.attentions,
+        # )
+
+        output = (start_logits, end_logits) + outputs[2:]
+        return ((total_loss,) + output) if total_loss is not None else output
+
+
+
+class Layoutlmv1ForQuestionAnswering_roberta(BertPreTrainedModel):
+    config_class = Layoutlmv1Config
+    pretrained_model_archive_map = LAYOUTLMV1_PRETRAINED_MODEL_ARCHIVE_MAP
+    base_model_prefix = "bert"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.roberta = Layoutlmv1Model(config)
+        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    def forward(
+        self,
+        input_ids,
+        bbox=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        # inputs_embeds=None,
+        start_positions=None,
+        end_positions=None,
+        # output_attentions=None,
+        # output_hidden_states=None,
+        # return_dict=None,
+        xpath_tags_seq=None,
+        xpath_subs_seq=None,
+        embedding_mode=None,
+    ):
+        r"""
+        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        """
+        # return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.roberta(
+            input_ids=input_ids,
+            bbox=bbox,
+            xpath_tags_seq=xpath_tags_seq,
+            xpath_subs_seq=xpath_subs_seq,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            embedding_mode=embedding_mode
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1)
+        end_logits = end_logits.squeeze(-1)
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions.clamp_(0, ignored_index)
+            end_positions.clamp_(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        # if not return_dict:
+        #     output = (start_logits, end_logits) + outputs[2:]
+        #     return ((total_loss,) + output) if total_loss is not None else output
+        #
+        # return QuestionAnsweringModelOutput(
+        #     loss=total_loss,
+        #     start_logits=start_logits,
+        #     end_logits=end_logits,
+        #     hidden_states=outputs.hidden_states,
+        #     attentions=outputs.attentions,
+        # )
+
+        output = (start_logits, end_logits) + outputs[2:]
+        return ((total_loss,) + output) if total_loss is not None else output

xdoc/fine_tuning/websrc/requirements.txt

@@ -0,0 +1,10 @@
+torch
+transformers==4.5.1
+torchvision==0.2.1
+tensorboard
+pillow==5.2.0
+textdistance
+h5py
+bs4
+networkx
+lxml

xdoc/fine_tuning/websrc/run_websrc.py

@@ -0,0 +1,92 @@
+import os
+import sys
+sys.path.append(os.getcwd())
+import torch
+import torch.nn as nn
+import shutil
+import logging
+import torch.distributed as dist
+
+
+from transformers import (
+    BertTokenizer,
+    RobertaTokenizer
+)
+
+from args import args
+from model import (  
+    Layoutlmv1ForQuestionAnswering,
+    Layoutlmv1Config,
+    Layoutlmv1Config_roberta,
+    Layoutlmv1ForQuestionAnswering_roberta
+)
+from util import set_seed, set_exp_folder, check_screen
+from trainer import train, evaluate # choose a specific train function
+# from data.datasets.docvqa import DocvqaDataset
+from websrc import get_websrc_dataset
+
+def main(args):
+
+    set_seed(args)
+    set_exp_folder(args)
+
+    # Set up logger
+    logging.basicConfig(filename="{}/output/{}/log.txt".format(args.output_dir, args.exp_name), level=logging.INFO,
+                        format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S')
+    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
+    logging.info('Args '+str(args))
+
+    # Get config, model, and tokenizer
+
+    if args.model_type == 'bert':
+        config_class, model_class, tokenizer_class = Layoutlmv1Config, Layoutlmv1ForQuestionAnswering, BertTokenizer
+    elif args.model_type == 'roberta':
+        config_class, model_class, tokenizer_class = Layoutlmv1Config_roberta, Layoutlmv1ForQuestionAnswering_roberta, RobertaTokenizer
+
+    config = config_class.from_pretrained(
+                args.model_name_or_path, cache_dir=args.cache_dir
+            )
+    config.add_linear = args.add_linear
+
+    tokenizer = tokenizer_class.from_pretrained(
+                args.model_name_or_path, cache_dir=args.cache_dir
+            )
+
+
+    model = model_class.from_pretrained(
+            args.model_name_or_path,
+            from_tf=bool(".ckpt" in args.model_name_or_path),
+            config=config,
+            cache_dir=args.cache_dir,
+        )
+
+    parameters = sum(p.numel() for p in model.parameters())
+    print("Total params: %.2fM" % (parameters/1e6))
+
+
+    ## Start training
+    if args.do_train:
+   
+        dataset_web = get_websrc_dataset(args, tokenizer)
+      
+        logging.info(f'Web dataset is successfully loaded. Length : {len(dataset_web)}')
+        train(args, dataset_web, model, tokenizer)
+
+    # ## Start evaluating
+    # if args.do_eval:
+
+    logging.info('Start evaluating')
+    dataset_web, examples, features = get_websrc_dataset(args, tokenizer, evaluate=True, output_examples=True)
+    logging.info(f'[Eval] Web dataset is successfully loaded. Length : {len(dataset_web)}')
+    evaluate(args, dataset_web, examples, features, model, tokenizer)
+
+
+    ## Start testing
+    if args.do_test:
+        pass
+
+
+if __name__ == '__main__':
+    main(args)
+
+   

xdoc/fine_tuning/websrc/trainer.py

@@ -0,0 +1,873 @@
+from genericpath import exists
+import os
+import torch.nn as nn
+import torch
+import logging
+from tqdm import tqdm, trange
+import timeit
+import collections
+import json
+import math
+from bs4 import BeautifulSoup
+from copy import deepcopy
+import string
+import re
+
+
+from torch.utils.tensorboard import SummaryWriter
+from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler)
+from transformers import (
+    BasicTokenizer,
+)
+
+from transformers import (
+    AdamW,
+    get_linear_schedule_with_warmup,
+)
+
+
+def reorganize_batch_web(args, batch_web):
+    dic = {}
+    dic['input_ids'] = batch_web[0].cuda()
+    dic['attention_mask'] = batch_web[1].cuda()
+    dic['token_type_ids'] = batch_web[2].cuda()
+    dic['xpath_tags_seq'] = batch_web[3].cuda()
+    dic['xpath_subs_seq'] = batch_web[4].cuda()
+    dic['start_positions'] = batch_web[5].cuda()
+    dic['end_positions'] = batch_web[6].cuda()
+    if 'box' in args.embedding_mode:
+        dic['bbox'] = batch_web[7].cuda() # new added
+    dic['embedding_mode'] = args.embedding_mode
+    return dic
+
+
+
+
+def train(args, dataset_web, model, tokenizer):
+    # torch.cuda.set_device(args.local_rank)
+
+    # Log when executing on clusters
+    try:
+        from azureml.core.run import Run
+        aml_run = Run.get_context()
+    except:
+        aml_run = None
+        
+    # Open tensorboard
+    writer = SummaryWriter(f'{args.output_dir}/output/{args.exp_name}')
+
+    # Count batch
+    gpu_nums = torch.cuda.device_count()
+    batch = args.batch_per_gpu * gpu_nums
+
+
+    dataloader_web = DataLoader(
+        dataset_web, batch_size=batch, num_workers=args.num_workers, pin_memory=False, shuffle=True,
+        
+    )
+
+    # Get warmup steps
+    total_step = args.epoch * len(dataloader_web)
+    warmup_steps = int(args.warmup_ratio * total_step)
+
+    # Prepare optimizers
+    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
+    )
+    scheduler = get_linear_schedule_with_warmup(
+        optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_step
+    )
+
+    # Transfer the parameters to cuda
+    model = model.cuda()
+
+
+    # Prepare fp16
+    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."
+            )
+        model, optimizer = amp.initialize(
+            model, optimizer, opt_level=args.fp16_opt_level
+        )
+        logging.info('Successfully load fp16 mode')
+
+    # Parallel or Distribute
+    if gpu_nums > 1:
+        model = torch.nn.DataParallel(model)
+
+
+    # Record some training info
+    logging.info("***** Running training *****")
+    # logging.info("  Num examples in dataset_doc = %d", len(dataset_doc))
+    logging.info("  Num examples in dataset_web = %d", len(dataset_web))
+    # logging.info("  Num steps for each epoch for doc = %d", len(dataloader_doc))
+    logging.info("  Num steps for each epoch for web = %d", len(dataloader_web))
+    logging.info("  Num Epochs = %d", args.epoch)
+    logging.info(
+        "  Instantaneous batch size per GPU = %d", args.batch_per_gpu
+    )
+    logging.info("  Total optimization steps = %d", total_step)
+
+    # Start training
+    model.zero_grad()
+    train_iterator = trange(
+        0,
+        int(args.epoch),
+        desc="Epoch",
+    )
+
+    global_step = 0
+    for now_epoch, _ in enumerate(tqdm(train_iterator, desc="Iteration")): # tqdm for epoch
+
+        # epoch_iterator_doc = iter(dataloader_doc)
+        epoch_iterator_web = iter(dataloader_web)
+
+        min_step = len(epoch_iterator_web)
+
+        for now_step in tqdm(range(min_step), desc="Iteration"): # tqdm for step
+
+            # batch_doc = epoch_iterator_doc.next()
+            batch_web = epoch_iterator_web.next()
+            batch_web = reorganize_batch_web(args, batch_web)
+
+            model.train()
+
+            # loss_doc = model(**batch_doc)[0]
+            loss_web = model(**batch_web)[0]
+            loss = loss_web
+
+            if gpu_nums > 1:
+                loss = loss.mean()
+                # loss_doc = loss_doc.mean()
+                loss_web = loss_web.mean()
+            
+            if args.fp16:
+                with amp.scale_loss(loss, optimizer) as scaled_loss:
+                    scaled_loss.backward()
+            else:
+                loss.backward()
+
+            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
+                )
+
+            if global_step % args.accumulation == 0:
+                optimizer.step()
+                model.zero_grad()
+            scheduler.step()
+
+            global_step += 1
+           
+            
+            if global_step % args.log_step == 0:
+                logging.info(f'epoch: {now_epoch} | step: {now_step+1} | total_step: {global_step} | loss: {loss} | lr: {scheduler.get_lr()[0]}')
+                writer.add_scalar('loss', loss, global_step//args.log_step)
+                # writer.add_scalar('loss_doc', loss_doc, global_step//args.log_step)
+                writer.add_scalar('loss_web', loss_web, global_step//args.log_step)
+                writer.add_scalar('lr', scheduler.get_lr()[0], global_step//args.log_step)
+                if aml_run is not None:
+                    aml_run.log('loss', loss.item())
+                    # aml_run.log('loss_doc', loss_doc.item())
+                    aml_run.log('loss_web', loss_web.item())
+                    aml_run.log('lr', scheduler.get_lr()[0])
+
+            if global_step % args.save_step == 0:
+                # Save model checkpoint
+                output_dir = os.path.join(args.output_dir, 'output', args.exp_name, f'step-{global_step}')
+                os.makedirs(output_dir, exist_ok=True)
+                
+                model_to_save = (
+                    model.module if hasattr(model, "module") else model
+                )  # Take care of distributed/parallel training
+                model_to_save.save_pretrained(output_dir)
+                tokenizer.save_pretrained(output_dir)
+
+                torch.save(args, os.path.join(output_dir, "training_args.bin"))
+                logging.info("Saving model checkpoint to %s", output_dir)
+
+                torch.save(
+                    optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt")
+                )
+                torch.save(
+                    scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt")
+                )
+                logging.info(
+                    "Saving optimizer and scheduler states to %s", output_dir
+                )
+
+            if global_step % 1000 == 0:
+                # eval
+                print('Start eval!')
+                from data.datasets.websrc import get_websrc_dataset
+                dataset_web, examples, features = get_websrc_dataset(args, tokenizer, evaluate=True, output_examples=True)
+                evaluate(args, dataset_web, examples, features, model, tokenizer, global_step)
+
+
+         
+
+
+
+RawResult = collections.namedtuple("RawResult",
+                                   ["unique_id", "start_logits", "end_logits"])
+
+def to_list(tensor):
+    return tensor.detach().cpu().tolist()
+
+def _get_best_indexes(logits, n_best_size):
+    """Get the n-best logits from a list."""
+    index_and_score = sorted(enumerate(logits), key=lambda x: x[1], reverse=True)
+
+    best_indexes = []
+    for i in range(len(index_and_score)):
+        if i >= n_best_size:
+            break
+        best_indexes.append(index_and_score[i][0])
+    return best_indexes
+
+
+def _get_final_text(pred_text, orig_text, do_lower_case, verbose_logging=False):
+    def _strip_spaces(text):
+        ns_chars = []
+        ns_to_s_map = collections.OrderedDict()
+        for (i, c) in enumerate(text):
+            if c == " ":
+                continue
+            ns_to_s_map[len(ns_chars)] = i
+            ns_chars.append(c)
+        ns_text = "".join(ns_chars)
+        return ns_text, ns_to_s_map
+
+    # We first tokenize `orig_text`, strip whitespace from the result
+    # and `pred_text`, and check if they are the same length. If they are
+    # NOT the same length, the heuristic has failed. If they are the same
+    # length, we assume the characters are one-to-one aligned.
+    tokenizer = BasicTokenizer(do_lower_case=do_lower_case)
+
+    tok_text = " ".join(tokenizer.tokenize(orig_text))
+
+    start_position = tok_text.find(pred_text)
+    if start_position == -1:
+        # if verbose_logging:
+        #     logging.info(
+        #         "Unable to find text: '%s' in '%s'" % (pred_text, orig_text))
+        return orig_text
+    end_position = start_position + len(pred_text) - 1
+
+    (orig_ns_text, orig_ns_to_s_map) = _strip_spaces(orig_text)
+    (tok_ns_text, tok_ns_to_s_map) = _strip_spaces(tok_text)
+
+    if len(orig_ns_text) != len(tok_ns_text):
+        if verbose_logging:
+            logging.info("Length not equal after stripping spaces: '%s' vs '%s'",
+                        orig_ns_text, tok_ns_text)
+        return orig_text
+
+    # We then project the characters in `pred_text` back to `orig_text` using
+    # the character-to-character alignment.
+    tok_s_to_ns_map = {}
+    for (i, tok_index) in tok_ns_to_s_map.items():
+        tok_s_to_ns_map[tok_index] = i
+
+    orig_start_position = None
+    if start_position in tok_s_to_ns_map:
+        ns_start_position = tok_s_to_ns_map[start_position]
+        if ns_start_position in orig_ns_to_s_map:
+            orig_start_position = orig_ns_to_s_map[ns_start_position]
+
+    if orig_start_position is None:
+        if verbose_logging:
+            logging.info("Couldn't map start position")
+        return orig_text
+
+    orig_end_position = None
+    if end_position in tok_s_to_ns_map:
+        ns_end_position = tok_s_to_ns_map[end_position]
+        if ns_end_position in orig_ns_to_s_map:
+            orig_end_position = orig_ns_to_s_map[ns_end_position]
+
+    if orig_end_position is None:
+        if verbose_logging:
+            logging.info("Couldn't map end position")
+        return orig_text
+
+    output_text = orig_text[orig_start_position:(orig_end_position + 1)]
+    return output_text
+
+
+
+def _compute_softmax(scores):
+    """Compute softmax probability over raw logits."""
+    if not scores:
+        return []
+
+    max_score = None
+    for score in scores:
+        if max_score is None or score > max_score:
+            max_score = score
+
+    exp_scores = []
+    total_sum = 0.0
+    for score in scores:
+        x = math.exp(score - max_score)
+        exp_scores.append(x)
+        total_sum += x
+
+    probs = []
+    for score in exp_scores:
+        probs.append(score / total_sum)
+    return probs
+
+
+
+class EvalOpts:
+    r"""
+    The options which the matrix evaluation process needs.
+    
+    Arguments:
+        data_file (str): the SQuAD-style json file of the dataset in evaluation.
+        root_dir (str): the root directory of the raw WebSRC dataset, which contains the HTML files.
+        pred_file (str): the prediction file which contain the best predicted answer text of each question from the
+                         model.
+        tag_pred_file (str): the prediction file which contain the best predicted answer tag id of each question from
+                             the model.
+        result_file (str): the file to write down the matrix evaluation results of each question.
+        out_file (str): the file to write down the final matrix evaluation results of the whole dataset.
+    """
+    def __init__(self, data_file, root_dir, pred_file, tag_pred_file, result_file='', out_file=""):
+        self.data_file = data_file
+        self.root_dir = root_dir
+        self.pred_file = pred_file
+        self.tag_pred_file = tag_pred_file
+        self.result_file = result_file
+        self.out_file = out_file
+
+
+def write_predictions(all_examples, all_features, all_results, n_best_size, max_answer_length, do_lower_case,
+                      output_prediction_file, output_tag_prediction_file,
+                      output_nbest_file, verbose_logging, tokenizer):
+    r"""
+    Compute and write down the final results, including the n best results.
+
+    Arguments:
+        all_examples (list[SRCExample]): all the SRC Example of the dataset; note that we only need it to provide the
+                                         mapping from example index to the question-answers id.
+        all_features (list[InputFeatures]): all the features for the input doc spans.
+        all_results (list[RawResult]): all the results from the models.
+        n_best_size (int): the number of the n best buffer and the final n best result saved.
+        max_answer_length (int): constrain the model to predict the answer no longer than it.
+        do_lower_case (bool): whether the model distinguish upper and lower case of the letters.
+        output_prediction_file (str): the file which the best answer text predictions will be written to.
+        output_tag_prediction_file (str): the file which the best answer tag predictions will be written to.
+        output_nbest_file (str): the file which the n best answer predictions including text, tag, and probabilities
+                                 will be written to.
+        verbose_logging (bool): if true, all of the warnings related to data processing will be printed.
+    """
+    logging.info("Writing predictions to: %s" % output_prediction_file)
+    logging.info("Writing nbest to: %s" % output_nbest_file)
+
+    example_index_to_features = collections.defaultdict(list)
+    for feature in all_features:
+        example_index_to_features[feature.example_index].append(feature)
+
+    unique_id_to_result = {}
+    for result in all_results:
+        unique_id_to_result[result.unique_id] = result
+
+    _PrelimPrediction = collections.namedtuple(  # pylint: disable=invalid-name
+        "PrelimPrediction",
+        ["feature_index", "start_index", "end_index", "start_logit", "end_logit", "tag_ids"])
+
+    all_predictions = collections.OrderedDict()
+    all_tag_predictions = collections.OrderedDict()
+    all_nbest_json = collections.OrderedDict()
+
+    for (example_index, example) in enumerate(all_examples):
+        features = example_index_to_features[example_index]
+
+        prelim_predictions = []
+
+        for (feature_index, feature) in enumerate(features): 
+            result = unique_id_to_result[feature.unique_id]
+            start_indexes = _get_best_indexes(result.start_logits, n_best_size)
+            end_indexes = _get_best_indexes(result.end_logits, n_best_size)
+            # if we could have irrelevant answers, get the min score of irrelevant
+            for start_index in start_indexes:
+                for end_index in end_indexes:
+                    # We could hypothetically create invalid predictions, e.g., predict
+                    # that the start of the span is in the question. We throw out all
+                    # invalid predictions.
+                    if start_index >= len(feature.tokens):
+                        continue
+                    if end_index >= len(feature.tokens):
+                        continue
+                    if start_index not in feature.token_to_orig_map:
+                        continue
+                    if end_index not in feature.token_to_orig_map:
+                        continue
+                    if not feature.token_is_max_context.get(start_index, False):
+                        continue
+                    if end_index < start_index:
+                        continue
+                    length = end_index - start_index + 1
+                    if length > max_answer_length:
+                        continue
+                    tag_ids = set(feature.token_to_tag_index[start_index: end_index + 1])
+                    prelim_predictions.append(
+                        _PrelimPrediction(
+                            feature_index=feature_index,
+                            start_index=start_index,
+                            end_index=end_index,
+                            start_logit=result.start_logits[start_index],
+                            end_logit=result.end_logits[end_index],
+                            tag_ids=list(tag_ids)))
+        prelim_predictions = sorted(
+            prelim_predictions,
+            key=lambda x: (x.start_logit + x.end_logit),
+            reverse=True)
+
+        _NbestPrediction = collections.namedtuple(  # pylint: disable=invalid-name
+            "NbestPrediction", ["text", "start_logit", "end_logit", "tag_ids"])
+
+        seen_predictions = {}
+        nbest = []
+        for pred in prelim_predictions:
+            if len(nbest) >= n_best_size:
+                break
+            feature = features[pred.feature_index]
+            if pred.start_index > 0:  # this is a non-null prediction
+                tok_tokens = feature.tokens[pred.start_index:(pred.end_index + 1)]
+                orig_doc_start = feature.token_to_orig_map[pred.start_index]
+                orig_doc_end = feature.token_to_orig_map[pred.end_index]
+                orig_tokens = example.doc_tokens[orig_doc_start:(orig_doc_end + 1)]
+                tok_text = " ".join(tok_tokens)
+
+                # De-tokenize WordPieces that have been split off.
+                tok_text = tok_text.replace(" ##", "")
+                tok_text = tok_text.replace("##", "")
+
+                # Clean whitespace
+                tok_text = tok_text.strip()
+                tok_text = " ".join(tok_text.split())
+                orig_text = " ".join(orig_tokens)
+
+                final_text = _get_final_text(tok_text, orig_text, do_lower_case, verbose_logging)
+                if final_text in seen_predictions:
+                    continue
+
+                seen_predictions[final_text] = True
+            else:
+                final_text = ""
+                seen_predictions[final_text] = True
+
+            nbest.append(
+                _NbestPrediction(
+                    text=final_text,
+                    start_logit=pred.start_logit,
+                    end_logit=pred.end_logit,
+                    tag_ids=pred.tag_ids))
+
+        # In very rare edge cases we could have no valid predictions. So we
+        # just create a nonce prediction in this case to avoid failure.
+        if not nbest:
+            nbest.append(
+                _NbestPrediction(text="empty", start_logit=0.0, end_logit=0.0, tag_ids=[-1]))
+
+        assert len(nbest) >= 1
+
+        total_scores = []
+        best_non_null_entry = None
+        for entry in nbest:
+            total_scores.append(entry.start_logit + entry.end_logit)
+            if not best_non_null_entry:
+                if entry.text:
+                    best_non_null_entry = entry
+
+        probs = _compute_softmax(total_scores)
+
+        nbest_json = []
+        for (i, entry) in enumerate(nbest):
+            output = collections.OrderedDict()
+            output["text"] = entry.text
+            output["probability"] = probs[i]
+            output["start_logit"] = entry.start_logit
+            output["end_logit"] = entry.end_logit
+            output["tag_ids"] = entry.tag_ids
+            nbest_json.append(output)
+
+        assert len(nbest_json) >= 1
+
+        best = nbest_json[0]["text"].split()
+        best = ' '.join([w for w in best
+                         if (w[0] != '<' or w[-1] != '>')
+                         and w != "<end-of-node>"
+                         and w != tokenizer.sep_token
+                         and w != tokenizer.cls_token])
+        all_predictions[example.qas_id] = best
+        all_tag_predictions[example.qas_id] = nbest_json[0]["tag_ids"]
+        all_nbest_json[example.qas_id] = nbest_json
+
+    with open(output_prediction_file, "w+") as writer:
+        writer.write(json.dumps(all_predictions, indent=4) + "\n")
+
+    with open(output_nbest_file, "w+") as writer:
+        writer.write(json.dumps(all_nbest_json, indent=4) + "\n")
+
+    with open(output_tag_prediction_file, 'w+') as writer:
+        writer.write(json.dumps(all_tag_predictions, indent=4) + '\n')
+    return
+
+
+
+def make_qid_to_has_ans(dataset):
+    r"""
+    Pick all the questions which has answer in the dataset and return the list.
+    """
+    qid_to_has_ans = {}
+    for domain in dataset:
+        for w in domain['websites']:
+            for qa in w['qas']:
+                qid_to_has_ans[qa['id']] = bool(qa['answers'])
+    return qid_to_has_ans
+
+
+
+def normalize_answer(s):
+    """Lower text and remove punctuation, articles and extra whitespace."""
+
+    def remove_articles(text):
+        regex = re.compile(r'\b(a|an|the)\b', re.UNICODE)
+        return re.sub(regex, ' ', text)
+
+    def white_space_fix(text):
+        return ' '.join(text.split())
+
+    def remove_punc(text):
+        exclude = set(string.punctuation)
+        return ''.join(ch for ch in text if ch not in exclude)
+
+    def lower(text):
+        return text.lower()
+
+    return white_space_fix(remove_articles(remove_punc(lower(s))))
+
+
+def compute_exact(a_gold, a_pred):
+    r"""
+    Calculate the exact match.
+    """
+    if normalize_answer(a_gold) == normalize_answer(a_pred):
+        return 1
+    return 0
+
+
+
+def get_raw_scores(dataset, preds, tag_preds, root_dir):
+    r"""
+    Calculate all the three matrix (exact match, f1, POS) for each question.
+
+    Arguments:
+        dataset (dict): the dataset in use.
+        preds (dict): the answer text prediction for each question in the dataset.
+        tag_preds (dict): the answer tags prediction for each question in the dataset.
+        root_dir (str): the base directory for the html files.
+
+    Returns:
+        tuple(dict, dict, dict): exact match, f1, pos scores for each question.
+    """
+    exact_scores = {}
+    f1_scores = {}
+    pos_scores = {}
+    for websites in dataset:
+        for w in websites['websites']:
+            f = os.path.join(root_dir, websites['domain'], w['page_id'][0:2], 'processed_data',
+                             w['page_id'] + '.html')
+            for qa in w['qas']:
+                qid = qa['id']
+                gold_answers = [a['text'] for a in qa['answers']
+                                if normalize_answer(a['text'])]
+                gold_tag_answers = [a['element_id'] for a in qa['answers']]
+                additional_tag_information = [a['answer_start'] for a in qa['answers']]
+                if not gold_answers:
+                    # For unanswerable questions, only correct answer is empty string
+                    gold_answers = ['']
+                if qid not in preds:
+                    print('Missing prediction for %s' % qid)
+                    continue
+                a_pred, t_pred = preds[qid], tag_preds[qid]
+                # Take max over all gold answers
+                exact_scores[qid] = max(compute_exact(a, a_pred) for a in gold_answers)
+                f1_scores[qid] = max(compute_f1(a, a_pred) for a in gold_answers)
+                pos_scores[qid] = max(compute_pos(f, t, a, t_pred)
+                                      for t, a in zip(gold_tag_answers, additional_tag_information))
+    return exact_scores, f1_scores, pos_scores
+
+def get_tokens(s):
+    r"""
+    Get the word list in the input.
+    """
+    if not s:
+        return []
+    return normalize_answer(s).split()
+
+
+
+def compute_f1(a_gold, a_pred):
+    r"""
+    Calculate the f1 score.
+    """
+    gold_toks = get_tokens(a_gold)
+    pred_toks = get_tokens(a_pred)
+    common = collections.Counter(gold_toks) & collections.Counter(pred_toks)
+    num_same = sum(common.values())
+    if len(gold_toks) == 0 or len(pred_toks) == 0:
+        # If either is no-answer, then F1 is 1 if they agree, 0 otherwise
+        return int(gold_toks == pred_toks)
+    if num_same == 0:
+        return 0
+    precision = 1.0 * num_same / len(pred_toks)
+    recall = 1.0 * num_same / len(gold_toks)
+    f1 = (2 * precision * recall) / (precision + recall)
+    return f1
+
+
+def compute_pos(f, t_gold, addition, t_pred):
+    r"""
+    Calculate the POS score.
+
+    Arguments:
+        f (str): the html file on which the question is based.
+        t_gold (int): the gold answer tag id provided by the dataset (the value correspond to the key element_id).
+        addition (int): the addition information used for yes/no question provided by the dataset (the value
+                        corresponding to the key answer_start).
+        t_pred (list[int]): the tag ids of the tags corresponding the each word in the predicted answer.
+    Returns:
+        float: the POS score.
+    """
+    h = BeautifulSoup(open(f), "lxml")
+    p_gold, e_gold = set(), h.find(tid=t_gold)
+    if e_gold is None:
+        if len(t_pred) != 1:
+            return 0
+        else:
+            t = t_pred[0]
+            e_pred, e_prev = h.find(tid=t), h.find(tid=t-1)
+            if (e_pred is not None) or (addition == 1 and e_prev is not None) or\
+                    (addition == 0 and e_prev is None):
+                return 0
+            else:
+                return 1
+    else:
+        p_gold.add(e_gold['tid'])
+        for e in e_gold.parents:
+            if int(e['tid']) < 2:
+                break
+            p_gold.add(e['tid'])
+        p = None
+        for t in t_pred:
+            p_pred, e_pred = set(), h.find(tid=t)
+            if e_pred is not None:
+                p_pred.add(e_pred['tid'])
+                if e_pred.name != 'html':
+                    for e in e_pred.parents:
+                        if int(e['tid']) < 2:
+                            break
+                        p_pred.add(e['tid'])
+            else:
+                p_pred.add(str(t))
+            if p is None:
+                p = p_pred
+            else:
+                p = p & p_pred
+        return len(p_gold & p) / len(p_gold | p)
+
+
+
+
+def make_pages_list(dataset):
+    r"""
+    Record all the pages which appears in the dataset and return the list.
+    """
+    pages_list = []
+    last_page = None
+    for domain in dataset:
+        for w in domain['websites']:
+            for qa in w['qas']:
+                if last_page != qa['id'][:4]:
+                    last_page = qa['id'][:4]
+                    pages_list.append(last_page)
+    return pages_list
+
+
+
+def make_eval_dict(exact_scores, f1_scores, pos_scores, qid_list=None):
+    r"""
+    Make the dictionary to show the evaluation results.
+    """
+    if qid_list is None:
+        total = len(exact_scores)
+        return collections.OrderedDict([
+            ('exact', 100.0 * sum(exact_scores.values()) / total),
+            ('f1', 100.0 * sum(f1_scores.values()) / total),
+            ('pos', 100.0 * sum(pos_scores.values()) / total),
+            ('total', total),
+        ])
+    else:
+        total = len(qid_list)
+        if total == 0:
+            return collections.OrderedDict([
+                ('exact', 0),
+                ('f1', 0),
+                ('pos', 0),
+                ('total', 0),
+            ])
+        return collections.OrderedDict([
+            ('exact', 100.0 * sum(exact_scores[k] for k in qid_list) / total),
+            ('f1', 100.0 * sum(f1_scores[k] for k in qid_list) / total),
+            ('pos', 100.0 * sum(pos_scores[k] for k in qid_list) / total),
+            ('total', total),
+        ])
+
+
+def merge_eval(main_eval, new_eval, prefix):
+    for k in new_eval:
+        main_eval['%s_%s' % (prefix, k)] = new_eval[k]
+
+
+def evaluate_on_squad(opts):
+    with open(opts.data_file) as f:
+        dataset_json = json.load(f)
+        dataset = dataset_json['data']
+    if isinstance(opts.pred_file, str):
+        with open(opts.pred_file) as f:
+            preds = json.load(f)
+    else:
+        preds = opts.pred_file
+    if isinstance(opts.tag_pred_file, str):
+        with open(opts.tag_pred_file) as f:
+            tag_preds = json.load(f)
+    else:
+        tag_preds = opts.tag_pred_file
+    qid_to_has_ans = make_qid_to_has_ans(dataset)
+    has_ans_qids = [k for k, v in qid_to_has_ans.items() if v]
+    no_ans_qids = [k for k, v in qid_to_has_ans.items() if not v]
+    exact, f1, pos = get_raw_scores(dataset, preds, tag_preds, opts.root_dir)
+    out_eval = make_eval_dict(exact, f1, pos)
+    if has_ans_qids:
+        has_ans_eval = make_eval_dict(exact, f1, pos, qid_list=has_ans_qids)
+        merge_eval(out_eval, has_ans_eval, 'HasAns')
+    if no_ans_qids:
+        no_ans_eval = make_eval_dict(exact, f1, pos, qid_list=no_ans_qids)
+        merge_eval(out_eval, no_ans_eval, 'NoAns')
+    print(json.dumps(out_eval, indent=2))
+    pages_list, write_eval = make_pages_list(dataset), deepcopy(out_eval)
+    for p in pages_list:
+        pages_ans_qids = [k for k, _ in qid_to_has_ans.items() if p in k]
+        page_eval = make_eval_dict(exact, f1, pos, qid_list=pages_ans_qids)
+        merge_eval(write_eval, page_eval, p)
+    if opts.result_file:
+        with open(opts.result_file, 'w') as f:
+            w = {}
+            for k, v in qid_to_has_ans.items():
+                w[k] = {'exact': exact[k], 'f1': f1[k], 'pos': pos[k]}
+            json.dump(w, f)
+    if opts.out_file:
+        with open(opts.out_file, 'w') as f:
+            json.dump(write_eval, f)
+    print('****** result ******')
+    print(out_eval)
+    return out_eval
+
+
+
+def evaluate(args, dataset_web, examples, features, model, tokenizer, step=0):
+
+    gpu_nums = torch.cuda.device_count()
+    batch = args.batch_per_gpu * gpu_nums
+
+    eval_sampler = SequentialSampler(dataset_web)
+    eval_dataloader = DataLoader(dataset_web, sampler=eval_sampler, batch_size=batch, num_workers=8)
+
+    # Eval!
+    logging.info("***** Running evaluation *****")
+    logging.info("  Num examples = %d", len(dataset_web))
+    logging.info("  Batch size = %d", batch)
+
+    model = model.cuda()
+
+    all_results = []
+    start_time = timeit.default_timer()
+    for batch in tqdm(eval_dataloader, desc="Evaluating"):
+        model.eval()
+        batch = tuple(t.cuda() for t in batch)
+        with torch.no_grad():
+            inputs = {'input_ids': batch[0],
+                      'attention_mask': batch[1],
+                      'token_type_ids': batch[2],
+                      'xpath_tags_seq': batch[4],
+                      'xpath_subs_seq': batch[5],
+                      }
+            feature_indices = batch[3]
+            outputs = model(**inputs)
+
+        for i, feature_index in enumerate(feature_indices):
+            eval_feature = features[feature_index.item()]
+            unique_id = int(eval_feature.unique_id)
+            result = RawResult(unique_id=unique_id,
+                               start_logits=to_list(outputs[0][i]),
+                               end_logits=to_list(outputs[1][i]))
+            all_results.append(result)
+
+    eval_time = timeit.default_timer() - start_time
+    logging.info("  Evaluation done in total %f secs (%f sec per example)", eval_time, eval_time / len(dataset_web))
+
+    # Compute predictions
+    # output_dir = os.path.join(args.output_dir, 'output', args.exp_name, f'step-{global_step}')
+
+    output_prediction_file = os.path.join(args.output_dir,"output", args.exp_name, f"predictions_{step}.json")
+    output_tag_prediction_file = os.path.join(args.output_dir,"output", args.exp_name, f"tag_predictions_{step}.json")
+    output_nbest_file = os.path.join(args.output_dir,"output", args.exp_name, f"nbest_predictions_{step}.json")
+    output_result_file = os.path.join(args.output_dir,"output", args.exp_name, f"qas_eval_results_{step}.json")
+    output_file = os.path.join(args.output_dir,"output", args.exp_name, f"eval_matrix_results_{step}")
+
+    write_predictions(examples, features, all_results, args.n_best_size, args.max_answer_length, args.do_lower_case,
+                      output_prediction_file, output_tag_prediction_file, output_nbest_file, args.verbose_logging,
+                      tokenizer)
+
+    # Evaluate
+    evaluate_options = EvalOpts(data_file=args.web_eval_file,
+                                root_dir=args.root_dir,
+                                pred_file=output_prediction_file,
+                                tag_pred_file=output_tag_prediction_file,
+                                result_file=output_result_file,
+                                out_file=output_file)
+    results = evaluate_on_squad(evaluate_options)
+    return results

xdoc/fine_tuning/websrc/util.py

@@ -0,0 +1,64 @@
+import random
+import numpy as np
+import torch
+import os
+import shutil
+# import logging
+import sys
+
+
+# def set_logging(args):
+#     '''
+#     Set logger for recording
+#     '''
+#     logging.basicConfig(filename="./output/{}/log.txt".format(args.exp_name), level=logging.INFO,
+#                         format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S')
+#     logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
+#     logging.info(str(args))
+
+
+
+def set_seed(args):
+    '''
+    Set seed for reproducibility
+    '''
+    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)
+
+
+def set_exp_folder(args):
+    '''
+    Create a folder to store experimental results e.g., checkpoints or log
+    '''
+    os.makedirs(os.path.join(args.output_dir, 'output'), exist_ok=True) 
+
+    if os.path.exists(os.path.join('output', args.exp_name)):
+        if not args.overwrite_output_dir: 
+            assert False, 'The exp_name is already used. Please modify the experiment name or use --overwrite_output_dir'
+        else:
+            print('Remove original directories.')
+            shutil.rmtree(os.path.join('output', args.exp_name))
+            print('Remove successfully.')
+    
+    os.makedirs(os.path.join(args.output_dir, 'output', args.exp_name), exist_ok=True)
+    exp_path = os.path.join(args.output_dir, 'output', args.exp_name)
+    print(f'Path [{exp_path}] has been created')
+
+
+def check_screen():
+    '''
+    Check whether the experiment is in screen
+    '''
+    text = os.popen('echo $STY').readlines()
+    string = ''
+    for line in text:
+        string += line
+    if len(string.strip()) == 0:
+        print("**** Attention Please! The code is not executed in Screen! ****")
+    else:
+        print(f'**** Screen Name : {string} ****')
+
+    

xdoc/fine_tuning/websrc/web_tag_utils.py

@@ -0,0 +1,32 @@
+tags_dict = {'a': 0, 'abbr': 1, 'acronym': 2, 'address': 3, 'altGlyph': 4, 'altGlyphDef': 5, 'altGlyphItem': 6,
+             'animate': 7, 'animateColor': 8, 'animateMotion': 9, 'animateTransform': 10, 'applet': 11, 'area': 12,
+             'article': 13, 'aside': 14, 'audio': 15, 'b': 16, 'base': 17, 'basefont': 18, 'bdi': 19, 'bdo': 20,
+             'bgsound': 21, 'big': 22, 'blink': 23, 'blockquote': 24, 'body': 25, 'br': 26, 'button': 27, 'canvas': 28,
+             'caption': 29, 'center': 30, 'circle': 31, 'cite': 32, 'clipPath': 33, 'code': 34, 'col': 35,
+             'colgroup': 36, 'color-profile': 37, 'content': 38, 'cursor': 39, 'data': 40, 'datalist': 41, 'dd': 42,
+             'defs': 43, 'del': 44, 'desc': 45, 'details': 46, 'dfn': 47, 'dialog': 48, 'dir': 49, 'div': 50, 'dl': 51,
+             'dt': 52, 'ellipse': 53, 'em': 54, 'embed': 55, 'feBlend': 56, 'feColorMatrix': 57,
+             'feComponentTransfer': 58, 'feComposite': 59, 'feConvolveMatrix': 60, 'feDiffuseLighting': 61,
+             'feDisplacementMap': 62, 'feDistantLight': 63, 'feFlood': 64, 'feFuncA': 65, 'feFuncB': 66, 'feFuncG': 67,
+             'feFuncR': 68, 'feGaussianBlur': 69, 'feImage': 70, 'feMerge': 71, 'feMergeNode': 72, 'feMorphology': 73,
+             'feOffset': 74, 'fePointLight': 75, 'feSpecularLighting': 76, 'feSpotLight': 77, 'feTile': 78,
+             'feTurbulence': 79, 'fieldset': 80, 'figcaption': 81, 'figure': 82, 'filter': 83, 'font-face-format': 84,
+             'font-face-name': 85, 'font-face-src': 86, 'font-face-uri': 87, 'font-face': 88, 'font': 89, 'footer': 90,
+             'foreignObject': 91, 'form': 92, 'frame': 93, 'frameset': 94, 'g': 95, 'glyph': 96, 'glyphRef': 97,
+             'h1': 98, 'h2': 99, 'h3': 100, 'h4': 101, 'h5': 102, 'h6': 103, 'head': 104, 'header': 105, 'hgroup': 106,
+             'hkern': 107, 'hr': 108, 'html': 109, 'i': 110, 'iframe': 111, 'image': 112, 'img': 113, 'input': 114,
+             'ins': 115, 'kbd': 116, 'keygen': 117, 'label': 118, 'legend': 119, 'li': 120, 'line': 121,
+             'linearGradient': 122, 'link': 123, 'main': 124, 'map': 125, 'mark': 126, 'marker': 127, 'marquee': 128,
+             'mask': 129, 'math': 130, 'menu': 131, 'menuitem': 132, 'meta': 133, 'metadata': 134, 'meter': 135,
+             'missing-glyph': 136, 'mpath': 137, 'nav': 138, 'nobr': 139, 'noembed': 140, 'noframes': 141,
+             'noscript': 142, 'object': 143, 'ol': 144, 'optgroup': 145, 'option': 146, 'output': 147, 'p': 148,
+             'param': 149, 'path': 150, 'pattern': 151, 'picture': 152, 'plaintext': 153, 'polygon': 154,
+             'polyline': 155, 'portal': 156, 'pre': 157, 'progress': 158, 'q': 159, 'radialGradient': 160, 'rb': 161,
+             'rect': 162, 'rp': 163, 'rt': 164, 'rtc': 165, 'ruby': 166, 's': 167, 'samp': 168, 'script': 169,
+             'section': 170, 'select': 171, 'set': 172, 'shadow': 173, 'slot': 174, 'small': 175, 'source': 176,
+             'spacer': 177, 'span': 178, 'stop': 179, 'strike': 180, 'strong': 181, 'style': 182, 'sub': 183,
+             'summary': 184, 'sup': 185, 'svg': 186, 'switch': 187, 'symbol': 188, 'table': 189, 'tbody': 190,
+             'td': 191, 'template': 192, 'text': 193, 'textPath': 194, 'textarea': 195, 'tfoot': 196, 'th': 197,
+             'thead': 198, 'time': 199, 'title': 200, 'tr': 201, 'track': 202, 'tref': 203, 'tspan': 204, 'tt': 205,
+             'u': 206, 'ul': 207, 'use': 208, 'var': 209, 'video': 210, 'view': 211, 'vkern': 212, 'wbr': 213,
+             'xmp': 214}