paddlepaddle--paddlenlp/paddlenlp/taskflow/utils.py

# coding:utf-8
# Copyright (c) 2021  PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import contextlib
import copy
import csv
import json
import math
import os
import pickle
import re
import traceback
import warnings
from collections import OrderedDict, namedtuple
from dataclasses import dataclass
from datetime import datetime
from functools import cmp_to_key
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import paddle
import paddle.nn.functional as F
import six
from paddle.dataset.common import md5file
from PIL import Image

from ..transformers.tokenizer_utils_base import PaddingStrategy, PretrainedTokenizerBase
from ..utils.downloader import DownloaderCheck, get_path_from_url
from ..utils.image_utils import (
    Bbox,
    DecodeImage,
    NormalizeImage,
    PadBatch,
    Permute,
    ResizeImage,
    check,
    img2base64,
    two_dimension_sort_layout,
)
from ..utils.log import logger

DOC_FORMAT = r"""
    Examples:
        .. code-block:: python
              """
DOWNLOAD_CHECK = False


def download_file(save_dir, filename, url, md5=None):
    """
    Download the file from the url to specified directory.
    Check md5 value when the file is exists, if the md5 value is the same as the existed file, just use
    the older file, if not, will download the file from the url.

    Args:
        save_dir(string): The specified directory saving the file.
        filename(string): The specified filename saving the file.
        url(string): The url downling the file.
        md5(string, optional): The md5 value that checking the version downloaded.
    """
    fullname = os.path.join(save_dir, filename)
    if os.path.exists(fullname):
        if md5 and (not md5file(fullname) == md5):
            logger.info("Updating {} from {}".format(filename, url))
            logger.disable()
            get_path_from_url(url, save_dir, md5)
    else:
        logger.info("Downloading {} from {}".format(filename, url))
        logger.disable()
        get_path_from_url(url, save_dir, md5)
    logger.enable()
    return fullname


def download_check(task):
    """
    Check the resource status in the specified task.

    Args:
        task(string): The name of specified task.
    """
    logger.disable()
    global DOWNLOAD_CHECK
    if not DOWNLOAD_CHECK:
        DOWNLOAD_CHECK = True
        checker = DownloaderCheck(task)
        checker.start()
        checker.join()
    logger.enable()


def add_docstrings(*docstr):
    """
    The function that add the doc string to doc of class.
    """

    def docstring_decorator(fn):
        fn.__doc__ = fn.__doc__ + "".join(DOC_FORMAT) + "".join(docstr)
        return fn

    return docstring_decorator


@contextlib.contextmanager
def static_mode_guard():
    paddle.enable_static()
    yield
    paddle.disable_static()


@contextlib.contextmanager
def dygraph_mode_guard():
    paddle.disable_static()
    yield


def cut_chinese_sent(para):
    """
    Cut the Chinese sentences more precisely, reference to "https://blog.csdn.net/blmoistawinde/article/details/82379256".
    """
    para = re.sub(r"([。！？\?])([^”’])", r"\1\n\2", para)
    para = re.sub(r"(\.{6})([^”’])", r"\1\n\2", para)
    para = re.sub(r"(\…{2})([^”’])", r"\1\n\2", para)
    para = re.sub(r"([。！？\?][”’])([^，。！？\?])", r"\1\n\2", para)
    para = para.rstrip()
    return para.split("\n")


class TermTreeNode(object):
    """Definition of term node. All members are protected, to keep rigorism of data struct.

    Args:
        sid (str): term id of node.
        term (str): term, common name of this term.
        base (str): `cb` indicates concept base, `eb` indicates entity base.
        term_type (Optional[str], optional): type of this term, constructs hierarchical of `term` node. Defaults to None.
        hyper (Optional[str], optional): parent type of a `type` node. Defaults to None.
        node_type (str, optional): type statement of node, `type` or `term`. Defaults to "term".
        alias (Optional[List[str]], optional): alias of this term. Defaults to None.
        alias_ext (Optional[List[str]], optional): extended alias of this term, CANNOT be used in matching.
            Defaults to None.
        sub_type (Optional[List[str]], optional): grouped by some term. Defaults to None.
        sub_term (Optional[List[str]], optional): some lower term. Defaults to None.
        data (Optional[Dict[str, Any]], optional): to sore full information of a term. Defaults to None.

    """

    def __init__(
        self,
        sid: str,
        term: str,
        base: str,
        node_type: str = "term",
        term_type: Optional[str] = None,
        hyper: Optional[str] = None,
        level: Optional[int] = None,
        alias: Optional[List[str]] = None,
        alias_ext: Optional[List[str]] = None,
        sub_type: Optional[List[str]] = None,
        sub_term: Optional[List[str]] = None,
        data: Optional[Dict[str, Any]] = None,
    ):
        self._sid = sid
        self._term = term
        self._base = base
        self._term_type = term_type
        self._hyper = hyper
        self._sub_term = sub_term if sub_term is not None else []
        self._sub_type = sub_type if sub_type is not None else []
        self._alias = alias if alias is not None else []
        self._alias_ext = alias_ext if alias_ext is not None else []
        self._data = data
        self._level = level
        self._node_type = node_type
        self._sons = set()

    def __str__(self):
        if self._data is not None:
            return json.dumps(self._data, ensure_ascii=False)
        else:
            res = {
                "termid": self._sid,
                "term": self._term,
                "src": self._base,
                "alias": self._alias,
                "alias_ext": self._alias_ext,
                "termtype": self._term_type,
                "subterms": self._sub_term,
                "subtype": self._sub_type,
                "links": [],
            }
            return json.dumps(res, ensure_ascii=False)

    @property
    def sid(self):
        return self._sid

    @property
    def term(self):
        return self._term

    @property
    def base(self):
        return self._base

    @property
    def alias(self):
        return self._alias

    @property
    def alias_ext(self):
        return self._alias_ext

    @property
    def termtype(self):
        return self._term_type

    @property
    def subtype(self):
        return self._sub_type

    @property
    def subterm(self):
        return self._sub_term

    @property
    def hyper(self):
        return self._hyper

    @property
    def level(self):
        return self._level

    @property
    def sons(self):
        return self._sons

    @property
    def node_type(self):
        return self._node_type

    def add_son(self, son_name):
        self._sons.add(son_name)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]):
        """Build a node from dictionary data.

        Args:
            data (Dict[str, Any]): Dictionary data contain all k-v data.

        Returns:
            [type]: TermTree node object.
        """
        return cls(
            sid=data["termid"],
            term=data["term"],
            base=data["src"],
            term_type=data["termtype"],
            sub_type=data["subtype"],
            sub_term=data["subterms"],
            alias=data["alias"],
            alias_ext=data["alias_ext"],
            data=data,
        )

    @classmethod
    def from_json(cls, json_str: str):
        """Build a node from JSON string.

        Args:
            json_str (str): JSON string formatted by TermTree data.

        Returns:
            [type]: TermTree node object.
        """
        dict_data = json.loads(json_str)
        return cls.from_dict(dict_data)


class TermTree(object):
    """TermTree class."""

    def __init__(self):
        self._nodes: Dict[str, TermTreeNode] = {}
        self._root = TermTreeNode(sid="root", term="root", base="cb", node_type="root", level=0)
        self._nodes["root"] = self.root
        self._index = {}

    def __build_sons(self):
        for node in self._nodes:
            self.__build_son(self._nodes[node])

    def __getitem__(self, item):
        return self._nodes[item]

    def __contains__(self, item):
        return item in self._nodes

    def __iter__(self):
        return self._nodes.__iter__()

    @property
    def root(self):
        return self._root

    def __load_type(self, file_path: str):
        with open(file_path, "rt", newline="", encoding="utf8") as csvfile:
            file_handler = csv.DictReader(csvfile, delimiter="\t")
            for row in file_handler:
                if row["type-1"] not in self:
                    self.add_type(type_name=row["type-1"], hyper_type="root")
                if row["type-2"] != "" and row["type-2"] not in self:
                    self.add_type(type_name=row["type-2"], hyper_type=row["type-1"])
                if row["type-3"] != "" and row["type-3"] not in self:
                    self.add_type(type_name=row["type-3"], hyper_type=row["type-2"])

    def __judge_term_node(self, node: TermTreeNode) -> bool:
        if node.termtype not in self:
            raise ValueError(f"Term type of new node {node.termtype} does not exists.")
        if node.sid in self:
            warnings.warn(f"{node.sid} exists, will be replaced by new node.")

    def add_term(
        self,
        term: Optional[str] = None,
        base: Optional[str] = None,
        term_type: Optional[str] = None,
        sub_type: Optional[List[str]] = None,
        sub_term: Optional[List[str]] = None,
        alias: Optional[List[str]] = None,
        alias_ext: Optional[List[str]] = None,
        data: Optional[Dict[str, Any]] = None,
    ):
        """Add a term into TermTree.

        Args:
            term (str): common name of name.
            base (str): term is concept or entity.
            term_type (str): term type of this term
            sub_type (Optional[List[str]], optional): sub type of this term, must exists in TermTree. Defaults to None.
            sub_terms (Optional[List[str]], optional): sub terms of this term. Defaults to None.
            alias (Optional[List[str]], optional): alias of this term. Defaults to None.
            alias_ext (Optional[List[str]], optional): . Defaults to None.
            data (Optional[Dict[str, Any]], optional): [description]. Defaults to None.
        """
        if data is not None:
            new_node = TermTreeNode.from_dict(data)
        else:
            new_node = TermTreeNode(
                sid=f"{term_type}_{base}_{term}",
                term=term,
                base=base,
                term_type=term_type,
                sub_term=sub_term,
                sub_type=sub_type,
                alias=alias,
                alias_ext=alias_ext,
                node_type="term",
            )
        self.__judge_term_node(new_node)
        self._nodes[new_node.sid] = new_node
        self.__build_index(new_node)

    def add_type(self, type_name, hyper_type):
        if type_name in self._nodes:
            raise ValueError(f"Term Type {type_name} exists.")
        if hyper_type not in self._nodes:
            raise ValueError(f"Hyper type {hyper_type} does not exist, please add it first.")
        if self._nodes[hyper_type].level == 3:
            raise ValueError(
                "Term type schema must be 3-LEVEL, 3rd level type node should not be a parent of type node."
            )
        self._nodes[type_name] = TermTreeNode(
            sid=type_name,
            term=type_name,
            base=None,
            hyper=hyper_type,
            node_type="type",
            level=self._nodes[hyper_type].level + 1,
        )
        self.__build_index(self._nodes[type_name])

    def __load_file(self, file_path: str):
        with open(file_path, encoding="utf-8") as fp:
            for line in fp:
                data = json.loads(line)
                self.add_term(data=data)

    def __build_son(self, node: TermTreeNode):
        """Build sons of a node

        Args:
            node (TermTreeNode): son node.
        """
        type_node = None
        if node.termtype is not None:
            type_node = self._nodes[node.termtype]
        elif node.hyper is not None:
            type_node = self._nodes[node.hyper]
        if type_node is not None:
            type_node.add_son(node.sid)
        for sub_type in node.subtype:
            sub_type_node = self._nodes[sub_type]
            sub_type_node.add_son(node.sid)

    def build_son(self, node: str):
        self.__build_son(self[node])

    def __build_index(self, node: TermTreeNode):
        if node.term not in self._index:
            self._index[node.term] = []
        self._index[node.term].append(node.sid)
        for alia in node.alias:
            if alia not in self._index:
                self._index[alia] = []
            self._index[alia].append(node.sid)

    def __judge_hyper(self, source_id, target_id) -> bool:
        queue = [source_id]
        visited_node = {source_id}
        while len(queue) > 0:
            cur_id = queue.pop(0)
            if cur_id == target_id:
                return True
            cur_node = self._nodes[cur_id]
            edge = []
            if cur_node.hyper is not None:
                edge.append(cur_node.hyper)
            if cur_node.termtype is not None:
                edge.append(cur_node.termtype)
            edge.extend(cur_node.subtype)
            for next_id in edge:
                if next_id not in visited_node:
                    queue.append(next_id)
                    visited_node.add(next_id)
        return False

    def find_term(self, term: str, term_type: Optional[str] = None) -> Tuple[bool, Union[List[str], None]]:
        """Find a term in Term Tree. If term not exists, return None.
        If `term_type` is not None, will find term with this type.

        Args:
            term (str): term to look up.
            term_type (Optional[str], optional): find term in this term_type. Defaults to None.

        Returns:
            Union[None, List[str]]: [description]
        """
        if term not in self._index:
            return False, None
        else:
            if term_type is None:
                return True, self._index[term]
            else:
                out = []
                for term_id in self._index[term]:
                    if self.__judge_hyper(term_id, term_type) is True:
                        out.append(term_id)
                if len(out) > 0:
                    return True, out
                else:
                    return False, None

    def build_from_dir(self, term_schema_path, term_data_path, linking=True):
        """Build TermTree from a directory which should contain type schema and term data.

        Args:
            dir ([type]): [description]
        """
        self.__load_type(term_schema_path)
        if linking:
            self.__load_file(term_data_path)
            self.__build_sons()

    @classmethod
    def from_dir(cls, term_schema_path, term_data_path, linking) -> "TermTree":
        """Build TermTree from a directory which should contain type schema and term data.

        Args:
            source_dir ([type]): [description]

        Returns:
            TermTree: [description]
        """
        term_tree = cls()
        term_tree.build_from_dir(term_schema_path, term_data_path, linking)
        return term_tree

    def __dfs(self, cur_id: str, depth: int, path: Dict[str, str], writer: csv.DictWriter):
        cur_node = self._nodes[cur_id]
        if cur_node.node_type == "term":
            return
        if depth > 0:
            path[f"type-{depth}"] = cur_id
        if path["type-1"] != "":
            writer.writerow(path)
        for son in cur_node.sons:
            self.__dfs(son, depth + 1, path, writer)
        if depth > 0:
            path[f"type-{depth}"] = ""

    def save(self, save_dir):
        """Save term tree to directory `save_dir`

        Args:
            save_dir ([type]): Directory.
        """
        if os.path.exists(save_dir) is False:
            os.makedirs(save_dir, exist_ok=True)
        out_path = {}
        for i in range(1, 3):
            out_path[f"type-{i}"] = ""
        with open(f"{save_dir}/termtree_type.csv", "wt", encoding="utf-8", newline="") as fp:
            fieldnames = ["type-1", "type-2", "type-3"]
            csv_writer = csv.DictWriter(fp, delimiter="\t", fieldnames=fieldnames)
            csv_writer.writeheader()
            self.__dfs("root", 0, out_path, csv_writer)
        with open(f"{save_dir}/termtree_data", "w", encoding="utf-8", newline="") as fp:
            for nid in self:
                node = self[nid]
                if node.node_type == "term":
                    print(node, file=fp)


def levenstein_distance(s1: str, s2: str) -> int:
    """Calculate minimal Levenstein distance between s1 and s2.

    Args:
        s1 (str): string
        s2 (str): string

    Returns:
        int: the minimal distance.
    """
    m, n = len(s1) + 1, len(s2) + 1

    # Initialize
    dp = [[0] * n for i in range(m)]
    dp[0][0] = 0
    for i in range(1, m):
        dp[i][0] = dp[i - 1][0] + 1
    for j in range(1, n):
        dp[0][j] = dp[0][j - 1] + 1

    for i in range(1, m):
        for j in range(1, n):
            if s1[i - 1] != s2[j - 1]:
                dp[i][j] = min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1]) + 1
            else:
                dp[i][j] = dp[i - 1][j - 1]
    return dp[m - 1][n - 1]


class BurkhardKellerNode(object):
    """Node implementation for BK-Tree. A BK-Tree node stores the information of current word, and its approximate words calculated by levenstein distance.

    Args:
        word (str): word of current node.
    """

    def __init__(self, word: str):
        self.word = word
        self.next = {}


class BurkhardKellerTree(object):
    """Implementation of BK-Tree"""

    def __init__(self):
        self.root = None
        self.nodes = {}

    def __add(self, cur_node: BurkhardKellerNode, word: str):
        """Insert a word into current tree. If tree is empty, set this word to root.

        Args:
            word (str): word to be inserted.
        """
        if self.root is None:
            self.root = BurkhardKellerNode(word)
            return
        if word in self.nodes:
            return
        dist = levenstein_distance(word, cur_node.word)
        if dist not in cur_node.next:
            self.nodes[word] = cur_node.next[dist] = BurkhardKellerNode(word)
        else:
            self.__add(cur_node.next[dist], word)

    def add(self, word: str):
        """Insert a word into current tree. If tree is empty, set this word to root.

        Args:
            word (str): word to be inserted.
        """
        return self.__add(self.root, word)

    def __search_similar_word(self, cur_node: BurkhardKellerNode, s: str, threshold: int = 2) -> List[str]:
        res = []
        if cur_node is None:
            return res
        dist = levenstein_distance(cur_node.word, s)
        if dist <= threshold:
            res.append((cur_node.word, dist))
        start = max(dist - threshold, 1)
        while start < dist + threshold:
            tmp_res = self.__search_similar_word(cur_node.next.get(start, None), s)[:]
            res.extend(tmp_res)
            start += 1
        return res

    def search_similar_word(self, word: str) -> List[str]:
        """Search the most similar (minimal levenstain distance) word between `s`.

        Args:
            s (str): target word

        Returns:
            List[str]: similar words.
        """
        res = self.__search_similar_word(self.root, word)

        def max_prefix(s1: str, s2: str) -> int:
            res = 0
            length = min(len(s1), len(s2))
            for i in range(length):
                if s1[i] == s2[i]:
                    res += 1
                else:
                    break
            return res

        res.sort(key=lambda d: (d[1], -max_prefix(d[0], word)))
        return res


class TriedTree(object):
    """Implementation of TriedTree"""

    def __init__(self):
        self.tree = {}

    def add_word(self, word):
        """add single word into TriedTree"""
        self.tree[word] = len(word)
        for i in range(1, len(word)):
            wfrag = word[:i]
            self.tree[wfrag] = self.tree.get(wfrag, None)

    def search(self, content):
        """Backward maximum matching

        Args:
            content (str): string to be searched
        Returns:
            List[Tuple]: list of maximum matching words, each element represents
                the starting and ending position of the matching string.
        """
        result = []
        length = len(content)
        for start in range(length):
            for end in range(start + 1, length + 1):
                pos = self.tree.get(content[start:end], -1)
                if pos == -1:
                    break
                if pos and (len(result) == 0 or end > result[-1][1]):
                    result.append((start, end))
        return result


class Customization(object):
    """
    User intervention based on Aho-Corasick automaton
    """

    def __init__(self):
        self.dictitem = {}
        self.ac = None

    def load_customization(self, filename, sep=None):
        """Load the custom vocab"""
        self.ac = TriedTree()
        with open(filename, "r", encoding="utf8") as f:
            for line in f:
                if sep is None:
                    words = line.strip().split()

                if len(words) == 0:
                    continue

                phrase = ""
                tags = []
                offset = []
                for word in words:
                    if word.rfind("/") < 1:
                        phrase += word
                        tags.append("")
                    else:
                        phrase += word[: word.rfind("/")]
                        tags.append(word[word.rfind("/") + 1 :])
                    offset.append(len(phrase))

                if len(phrase) < 2 and tags[0] == "":
                    continue

                self.dictitem[phrase] = (tags, offset)
                self.ac.add_word(phrase)

    def parse_customization(self, query, lac_tags, prefix=False):
        """Use custom vocab to modify the lac results"""
        if not self.ac:
            logger.warning("customization dict is not load")
            return
        ac_res = self.ac.search(query)

        for begin, end in ac_res:
            phrase = query[begin:end]
            index = begin

            tags, offsets = self.dictitem[phrase]

            if prefix:
                for tag, offset in zip(tags, offsets):
                    while index < begin + offset:
                        if len(tag) == 0:
                            lac_tags[index] = "I" + lac_tags[index][1:]
                        else:
                            lac_tags[index] = "I-" + tag
                        index += 1
                lac_tags[begin] = "B" + lac_tags[begin][1:]
                for offset in offsets:
                    index = begin + offset
                    if index < len(lac_tags):
                        lac_tags[index] = "B" + lac_tags[index][1:]
            else:
                for tag, offset in zip(tags, offsets):
                    while index < begin + offset:
                        if len(tag) == 0:
                            lac_tags[index] = lac_tags[index][:-1] + "I"
                        else:
                            lac_tags[index] = tag + "-I"
                        index += 1
                lac_tags[begin] = lac_tags[begin][:-1] + "B"
                for offset in offsets:
                    index = begin + offset
                    if index < len(lac_tags):
                        lac_tags[index] = lac_tags[index][:-1] + "B"


class SchemaTree(object):
    """
    Implementation of SchemaTree
    """

    def __init__(self, name="root", children=None):
        self.name = name
        self.children = []
        self.prefix = None
        self.parent_relations = None
        self.parent = None
        if children is not None:
            for child in children:
                self.add_child(child)

    def __repr__(self):
        return self.name

    def add_child(self, node):
        assert isinstance(node, SchemaTree), "The children of a node should be an instance of SchemaTree."
        self.children.append(node)


def get_id_and_prob(span_set, offset_mapping):
    """
    Return text id and probability of predicted spans

    Args:
        span_set (set): set of predicted spans.
        offset_mapping (list[int]): list of pair preserving the
                index of start and end char in original text pair (prompt + text) for each token.
    Returns:
        sentence_id (list[tuple]): index of start and end char in original text.
        prob (list[float]): probabilities of predicted spans.
    """
    prompt_end_token_id = offset_mapping[1:].index([0, 0])
    bias = offset_mapping[prompt_end_token_id][1] + 1
    for idx in range(1, prompt_end_token_id + 1):
        offset_mapping[idx][0] -= bias
        offset_mapping[idx][1] -= bias

    sentence_id = []
    prob = []
    for start, end in span_set:
        prob.append(start[1] * end[1])
        start_id = offset_mapping[start[0]][0]
        end_id = offset_mapping[end[0]][1]
        sentence_id.append((start_id, end_id))
    return sentence_id, prob


def dbc2sbc(s):
    rs = ""
    for char in s:
        code = ord(char)
        if code == 0x3000:
            code = 0x0020
        else:
            code -= 0xFEE0
        if not (0x0021 <= code and code <= 0x7E):
            rs += char
            continue
        rs += chr(code)
    return rs


class WordTagRelationExtractor(object):
    """Implement of information extractor."""

    _chain_items = {"和", "与", "兼", "及", "以及", "还有", "并"}
    _all_items = None
    _jux_buf = []

    def __init__(self, schema):
        self._schema = schema

    @property
    def schema(self):
        return self._schema

    @classmethod
    def from_dict(cls, config_dict):
        """Make an instance from a configuration dictionary.

        Args:
            config_dict (Dict[str, Any]): configuration dict.
        """
        res = {}

        for i, trip_config in enumerate(config_dict):
            head_role_type = trip_config["head_role"]
            if head_role_type not in res:
                res[head_role_type] = {"trigger": {}, "g_t_map": {}, "rel_group": {}, "trig_word": {}}
            group_name = trip_config["group"]
            if "rel_group" in trip_config:
                res[head_role_type]["rel_group"][group_name] = trip_config["rel_group"]
            if group_name not in res[head_role_type]["trig_word"]:
                res[head_role_type]["trig_word"][group_name] = set()
            for trig_word in trip_config["trig_word"]:
                res[head_role_type]["trigger"][trig_word] = {
                    "trigger_type": trip_config["trig_type"],
                    "group_name": group_name,
                    "rev_flag": trip_config["reverse"],
                }
                res[head_role_type]["trig_word"][group_name].add(trig_word)
            res[head_role_type]["g_t_map"][group_name] = trip_config["tail_role"]

        return cls(res)

    @classmethod
    def from_json(cls, json_str):
        """Implement an instance from JSON str."""
        config_dict = json.loads(json_str)
        return cls.from_dict(config_dict)

    @classmethod
    def from_pkl(cls, pkl_path):
        """Implement an instance from a serialized pickle package."""
        with open(pkl_path, "rb") as fp:
            schema = pickle.load(fp)
        return cls(schema)

    @classmethod
    def from_config(cls, config_path):
        """Implement an instance from a configuration file."""
        with open(config_path, encoding="utf-8") as fp:
            config_json = json.load(fp)
        return cls.from_dict(config_json)

    def add_schema_from_dict(self, config_dict):
        """Add the schema from the dict."""
        for i, trip_config in enumerate(config_dict):
            head_role_type = trip_config["head_role"]
            if head_role_type not in self._schema:
                self._schema[head_role_type] = {"trigger": {}, "g_t_map": {}, "rel_group": {}, "trig_word": {}}
            group_name = trip_config["group"]
            if "rel_group" in self._schema:
                self._schema[head_role_type]["rel_group"][group_name] = trip_config["rel_group"]
            if group_name not in self._schema[head_role_type]["trig_word"]:
                self._schema[head_role_type]["trig_word"][group_name] = set()
            for trig_word in trip_config["trig_word"]:
                self._schema[head_role_type]["trigger"][trig_word] = {
                    "trigger_type": trip_config["trig_type"],
                    "group_name": group_name,
                    "rev_flag": trip_config["reverse"],
                }
                self._schema[head_role_type]["trig_word"][group_name].add(trig_word)
            self._schema[head_role_type]["g_t_map"][group_name] = trip_config["tail_role"]

    def _judge_jux(self, wordtag_item):
        """Judge whether `wordtag_item` is a relevance component between two juxtaposed items.

        Args:
            wordtag_item (dict): input item.

        Returns:
            bool: [description]
        """
        if wordtag_item["item"] in {"、", " ", "《", "》", "/"}:
            return True
        if wordtag_item["item"] in self._chain_items and wordtag_item["wordtag_label"] == "连词":
            return True
        return False

    def _search_jux(self, cur_item, cur_pos=0, jux_type=None, jux_word=None, status_flag=None, search_list=None):
        """Find juxtaposed items with `cur_item` at `cur_pos` in `self._all_items`.

        Args:
            cur_item (Dict[str, Any]): the item current viewing.
            cur_pos (int, optional): current position of viewing item. Defaults to 0.
            jux_type (Set[str], optional):  wordtag labels that can be considered as juxtaposed item. Defaults to None.
            jux_word (Set[str], optional):  words that can be considered as juxtaposed item. Defaults to None.
            status_flag (bool, optional): if True, on the juxtaposed item, or on chain item. Defaults to None.

        Returns:
            int: end position of juxtable items.
        """
        if search_list is None:
            search_list = self._all_items

        if jux_type is None and jux_word is None:
            raise ValueError("`jux_type` and `jux_word` are both None.")

        if status_flag is True:
            self._jux_buf.append(cur_item)

        if cur_pos >= len(search_list) - 1:
            return cur_pos

        next_item = search_list[cur_pos + 1]

        if self._judge_jux(next_item) is True:
            return self._search_jux(
                cur_item=next_item,
                cur_pos=cur_pos + 1,
                jux_type=jux_type,
                jux_word=jux_word,
                status_flag=False,
                search_list=search_list,
            )

        next_flag = True
        if jux_type is not None:
            next_flag = next_flag and self._match_item(next_item, jux_type)
        if jux_word is not None:
            next_flag = next_flag and (next_item["item"] in jux_word)
        if next_flag is True:
            return self._search_jux(
                cur_item=next_item, cur_pos=cur_pos + 1, jux_type=jux_type, jux_word=jux_word, status_flag=True
            )
        if next_flag is not True:
            while self._judge_jux(search_list[cur_pos]) is True:
                cur_pos -= 1
        return cur_pos

    @staticmethod
    def _match_item(item, type_can):
        match_key = item["wordtag_label"].split("_")[0]
        return match_key in type_can or item["wordtag_label"] in type_can

    def _trig_handler(self, cur_item, head_conf):
        """Whether current item is a trigger, if True, return corresponding flag and configuration.

        Args:
            cur_item (Dict[str, Any]): current viewing ite,
            st_conf (Dict[str, Any]): config

        Returns:
            Tuple[str, Union[None, dict]]: [description]
        """
        trigger_conf = head_conf["trigger"]
        if cur_item["item"] in trigger_conf:
            # find a trigger, then judge whether it is a tail-trigger or a rel trigger.
            if trigger_conf[cur_item["item"]]["trigger_type"] == "role":
                # find a tail-trigger, then judge wordtag label.
                group_name = trigger_conf[cur_item["item"]]["group_name"]
                for tail_conf in head_conf["g_t_map"][group_name]:
                    if self._match_item(cur_item, tail_conf["main"]) is True:
                        return "trig_t", tail_conf
                else:
                    return "un_trig", None
            else:
                return "trig_g", None
        else:
            return "un_trig", None

    def _find_tail(self, search_range, sg_conf, head_hype):
        """Find tail role in `search_range`

        Args:
            search_range (List[int]): index range of `self._all_items`, items to be checked.
            sg_conf (Dict[str, Any]): configuration of group.
            head_type (str): wordtag label of head role item.
        """
        for i in search_range:
            item = self._all_items[i]
            if item["item"] in {"，", "？", "、", "。", "；"}:
                return -2, None
            for j, tail_conf in enumerate(sg_conf):
                flag = self._match_item(item, tail_conf["main"])
                if flag is True:
                    return i, tail_conf
                if item["wordtag_label"].startswith(head_hype):
                    return -1, None

        return -2, None

    def _find_supp(self, search_range, search_type):
        res = []
        for i in search_range:
            item = self._all_items[i]
            if item["item"] == "，":
                break
            if any(item["wordtag_label"].startswith(sup_t) for sup_t in search_type):
                res.append(item)
        return res if len(res) > 0 else None

    def _make_output(self, head_item, tail_item, group, source, support=None, trig_word=None, **kwargs):
        """Make formatted outputs of mined results.

        Args:
            head_item (Dict[str, Any]): [description]
            head_index (int): [description]
            tail_item (List[Dict[str, Any]]): [description]
            tail_indices (List[int]): [description]
            group (str): [description]
            source (str): [description]
            support (List[Dict[str, Any]], optional): [description]. Defaults to None.
            support_indices (List[int], optional): [description]. Defaults to None.
            trig_word (List[str], optional): [description]. Defaults to None.
            trig_indices (List[int], optional): [description]. Defaults to None.
        """
        res = {
            "HEAD_ROLE": {
                "item": head_item["item"],
                "type": head_item["wordtag_label"],
                "offset": head_item["offset"],
            },
            "TAIL_ROLE": [
                {"item": ti["item"], "offset": ti["offset"], "type": ti["wordtag_label"]} for ti in tail_item
            ],
            "GROUP": group,
            "SRC": source,
        }
        if support is not None:
            res["SUPPORT"] = [
                {
                    "item": si["item"],
                    "offset": si["offset"],
                    "type": si["wordtag_label"],
                }
                for si in support
            ]
        if trig_word is not None:
            res["TRIG"] = [
                {
                    "item": ti["item"],
                    "offset": ti["offset"],
                }
                for ti in trig_word
            ]
        return res

    def _reverse(self, res, group_name=None):
        ret = []
        for rev_head in res["TAIL_ROLE"]:
            rev_tmp = {
                "HEAD_ROLE": rev_head,
                "TAIL_ROLE": [res["HEAD_ROLE"]],
                "GROUP": group_name if group_name is not None else res["GROUP"],
            }
            if "SUPPORT" in res:
                rev_tmp["SUPPORT"] = res["SUPPORT"]
            if "TRIG" in res:
                rev_tmp["TRIG"] = res["TRIG"]
            rev_tmp["SRC"] = "REVERSE" if group_name is not None else res["SRC"]
            ret.append(rev_tmp)
        return ret

    def extract_spo(self, all_items):
        """Pipeline of mining procedure.

        Args:
            all_items ([type]): [description]
        """
        self._all_items = all_items

        res_cand = []

        # Match head role, and consider it as central, search others.
        for i, head_cand in enumerate(self._all_items):
            last_end = i
            try:
                datetime.strptime(head_cand["item"], "%Y年%m月%d日")
                head_cand["wordtag_label"] = "时间类_具体时间"
            except ValueError:
                pass

            if head_cand["wordtag_label"] in self._schema:
                head_conf = self._schema[head_cand["wordtag_label"]]
                head_type = head_cand["wordtag_label"]
            else:
                match_key = head_cand["wordtag_label"].split("_")[0]
                if match_key in self._schema:
                    head_conf = self._schema[match_key]
                    head_type = match_key
                else:
                    continue

            trig_status = "un_trig"

            # Consider `head_cand` as a start item, find trigger words behind.
            # We suppose that minning strategy is directed, so only search items behinds head.
            # If need, we can reverse constructed triples.
            j = i + 1
            while j < len(self._all_items):
                cur_item = all_items[j]
                cur_pos = j
                j += 1

                trig_status, trig_conf = self._trig_handler(cur_item, self._schema[head_type])

                # Find a tail role, generate corresponding triple.
                if trig_status == "trig_t":
                    trig_status = "un_trig"
                    tail_flag = True
                    for k in range(i + 1, j):
                        if self._all_items[k]["wordtag_label"] == head_cand["wordtag_label"]:
                            tail_flag = False
                            break
                    if tail_flag is False:
                        continue

                    group_name = head_conf["trigger"][cur_item["item"]]["group_name"]
                    del self._jux_buf[:]
                    idx = self._search_jux(
                        cur_item=cur_item, cur_pos=cur_pos, jux_type=trig_conf["main"], status_flag=True
                    )
                    supports = self._find_supp(search_range=range(j - 1, i, -1), search_type=trig_conf["support"])

                    tmp = self._make_output(
                        head_item=head_cand,
                        tail_item=self._jux_buf[:],
                        group=group_name,
                        support=supports,
                        source="TAIL",
                    )

                    # Reverse triple if group has relative.
                    if (
                        group_name in head_conf.get("rel_group", {})
                        or head_conf["trigger"][cur_item["item"]]["rev_flag"] is True
                    ):
                        rev_tmp = self._reverse(tmp, head_conf.get("rel_group", {}).get(group_name, None))
                        res_cand.extend(rev_tmp[:])
                    if head_conf["trigger"][cur_item["item"]]["rev_flag"] is False:
                        res_cand.append(tmp.copy())

                    j = idx + 1
                    last_end = idx
                    continue

                # Find a group trigger word, look for tail role items of current head role and group argument.
                # Searching range is items behind group trigger and items between head rold and group trigger word.
                if trig_status == "trig_g":
                    trig_status = "un_trig"
                    group_name = head_conf["trigger"][cur_item["item"]]["group_name"]

                    del self._jux_buf[:]
                    g_start_idx = j - 1
                    g_idx = self._search_jux(
                        cur_item=cur_item,
                        cur_pos=cur_pos,
                        jux_word=head_conf["trig_word"][group_name],
                        status_flag=True,
                    )

                    g_trig_words = self._jux_buf[:]
                    j = g_idx + 1

                    # Search right.
                    if j < len(self._all_items) - 1:
                        tail_idx, tail_conf = self._find_tail(
                            range(g_idx + 1, len(self._all_items)), head_conf["g_t_map"][group_name], head_type
                        )

                        if tail_idx > 0:
                            # Find a tail.
                            tail_item = self._all_items[tail_idx]
                            del self._jux_buf[:]
                            idx = self._search_jux(
                                cur_item=tail_item, cur_pos=tail_idx, status_flag=True, jux_type=tail_conf["main"]
                            )
                            tail_cand = self._jux_buf[:]
                            supports = self._find_supp(range(tail_idx - 1, i, -1), tail_conf["support"])

                            tmp = self._make_output(
                                head_item=head_cand,
                                tail_item=tail_cand,
                                group=group_name,
                                source="HGT",
                                support=supports,
                                trig_word=g_trig_words,
                            )

                            if (
                                group_name in head_conf.get("rel_group", {})
                                or head_conf["trigger"][cur_item["item"]]["rev_flag"] is True
                            ):
                                rev_tmp = self._reverse(tmp, head_conf.get("rel_group", {}).get(group_name, None))
                                res_cand.extend(rev_tmp[:])
                            if head_conf["trigger"][cur_item["item"]]["rev_flag"] is False:
                                res_cand.append(tmp.copy())

                            j = idx + 1
                            last_end = idx
                            continue

                    # Search left
                    if g_idx - i > len(g_trig_words):
                        tail_idx, tail_conf = self._find_tail(
                            range(g_start_idx, last_end, -1), head_conf["g_t_map"][group_name], head_type
                        )
                        tail_item = self._all_items[tail_idx]
                        if tail_idx > 0:
                            del self._jux_buf[:]
                            _ = self._search_jux(
                                cur_item=tail_item,
                                cur_pos=0,
                                jux_type=tail_conf["main"],
                                status_flag=True,
                                search_list=self._all_items[i + 1 : tail_idx][::-1],
                            )
                            tail_cand = self._jux_buf[:]
                            supports = self._find_supp(range(g_idx - 1, last_end, -1), tail_conf["support"])
                            last_end = g_idx

                            tmp = self._make_output(
                                head_item=head_cand,
                                tail_item=tail_cand,
                                group=group_name,
                                trig_word=g_trig_words,
                                source="HTG",
                                support=supports,
                            )

                            if (
                                group_name in head_conf.get("rel_group", {})
                                or head_conf["trigger"][cur_item["item"]]["rev_flag"] is True
                            ):
                                rev_tmp = self._reverse(tmp, head_conf.get("rel_group", {}).get(group_name, None))
                                res_cand.extend(rev_tmp[:])
                            if head_conf["trigger"][cur_item["item"]]["rev_flag"] is False:
                                res_cand.append(tmp.copy())
                            continue
        return res_cand


@dataclass
class DataCollatorGP:
    tokenizer: PretrainedTokenizerBase
    padding: Union[bool, str, PaddingStrategy] = True
    max_length: Optional[int] = None
    label_maps: Optional[dict] = None
    task_type: Optional[str] = None

    def __call__(self, features: List[Dict[str, Union[List[int], paddle.Tensor]]]) -> Dict[str, paddle.Tensor]:
        new_features = [{k: v for k, v in f.items() if k not in ["offset_mapping", "text"]} for f in features]

        batch = self.tokenizer.pad(
            new_features,
            padding=self.padding,
        )

        batch = [paddle.to_tensor(batch[k]) for k in batch.keys()]
        batch.append([feature["offset_mapping"] for feature in features])
        batch.append([feature["text"] for feature in features])
        return batch


@dataclass
class DataCollatorForErnieCtm:
    tokenizer: PretrainedTokenizerBase
    padding: Union[bool, str, PaddingStrategy] = True
    model: Optional[str] = "wordtag"

    def __call__(self, features: List[Dict[str, Union[List[int], paddle.Tensor]]]) -> Dict[str, paddle.Tensor]:
        no_pad = "seq_len" if self.model == "wordtag" else "label_indices"
        new_features = [{k: v for k, v in f.items() if k != no_pad} for f in features]
        batch = self.tokenizer.pad(
            new_features,
            padding=self.padding,
        )

        batch = [paddle.to_tensor(batch[k]) for k in batch.keys()]
        batch.append(paddle.to_tensor([f[no_pad] for f in features]))
        return batch


def gp_decode(batch_outputs, offset_mappings, texts, label_maps, task_type="relation_extraction"):
    if task_type == "entity_extraction":
        batch_ent_results = []
        for entity_output, offset_mapping, text in zip(batch_outputs[0].numpy(), offset_mappings, texts):
            entity_output[:, [0, -1]] -= np.inf
            entity_output[:, :, [0, -1]] -= np.inf
            entity_probs = F.softmax(paddle.to_tensor(entity_output), axis=1).numpy()
            ent_list = []
            for l, start, end in zip(*np.where(entity_output > 0.0)):
                ent_prob = entity_probs[l, start, end]
                start, end = (offset_mapping[start][0], offset_mapping[end][-1])
                ent = {
                    "text": text[start:end],
                    "type": label_maps["id2entity"][str(l)],
                    "start_index": start,
                    "probability": ent_prob,
                }
                ent_list.append(ent)
            batch_ent_results.append(ent_list)
        return batch_ent_results
    else:
        batch_ent_results = []
        batch_rel_results = []
        for entity_output, head_output, tail_output, offset_mapping, text in zip(
            batch_outputs[0].numpy(),
            batch_outputs[1].numpy(),
            batch_outputs[2].numpy(),
            offset_mappings,
            texts,
        ):
            entity_output[:, [0, -1]] -= np.inf
            entity_output[:, :, [0, -1]] -= np.inf
            entity_probs = F.softmax(paddle.to_tensor(entity_output), axis=1).numpy()
            head_probs = F.softmax(paddle.to_tensor(head_output), axis=1).numpy()
            tail_probs = F.softmax(paddle.to_tensor(tail_output), axis=1).numpy()

            ents = set()
            ent_list = []
            for l, start, end in zip(*np.where(entity_output > 0.0)):
                ent_prob = entity_probs[l, start, end]
                ents.add((start, end))
                start, end = (offset_mapping[start][0], offset_mapping[end][-1])
                ent = {
                    "text": text[start:end],
                    "type": label_maps["id2entity"][str(l)],
                    "start_index": start,
                    "probability": ent_prob,
                }
                ent_list.append(ent)
            batch_ent_results.append(ent_list)

            rel_list = []
            for sh, st in ents:
                for oh, ot in ents:
                    p1s = np.where(head_output[:, sh, oh] > 0.0)[0]
                    p2s = np.where(tail_output[:, st, ot] > 0.0)[0]
                    ps = set(p1s) & set(p2s)
                    for p in ps:
                        rel_prob = head_probs[p, sh, oh] * tail_probs[p, st, ot]
                        if task_type == "relation_extraction":
                            rel = {
                                "subject": text[offset_mapping[sh][0] : offset_mapping[st][1]],
                                "predicate": label_maps["id2relation"][str(p)],
                                "object": text[offset_mapping[oh][0] : offset_mapping[ot][1]],
                                "subject_start_index": offset_mapping[sh][0],
                                "object_start_index": offset_mapping[oh][0],
                                "probability": rel_prob,
                            }
                        else:
                            rel = {
                                "aspect": text[offset_mapping[sh][0] : offset_mapping[st][1]],
                                "sentiment": label_maps["id2relation"][str(p)],
                                "opinion": text[offset_mapping[oh][0] : offset_mapping[ot][1]],
                                "aspect_start_index": offset_mapping[sh][0],
                                "opinion_start_index": offset_mapping[oh][0],
                                "probability": rel_prob,
                            }
                        rel_list.append(rel)
            batch_rel_results.append(rel_list)
        return (batch_ent_results, batch_rel_results)


DocSpan = namedtuple("DocSpan", ["start", "length"])

Example = namedtuple(
    "Example",
    [
        "keys",
        "key_labels",
        "doc_tokens",
        "text",
        "qas_id",
        "model_type",
        "seq_labels",
        "ori_boxes",
        "boxes",
        "segment_ids",
        "symbol_ids",
        "im_base64",
        "image_rois",
    ],
)

Feature = namedtuple(
    "Feature",
    [
        "unique_id",
        "example_index",
        "qas_id",
        "doc_span_index",
        "tokens",
        "token_to_orig_map",
        "token_is_max_context",
        "token_ids",
        "position_ids",
        "text_type_ids",
        "text_symbol_ids",
        "overlaps",
        "key_labels",
        "seq_labels",
        "se_seq_labels",
        "bio_seq_labels",
        "bioes_seq_labels",
        "keys",
        "model_type",
        "doc_tokens",
        "doc_labels",
        "text",
        "boxes",
        "segment_ids",
        "im_base64",
        "image_rois",
    ],
)


class Compose(object):
    """compose"""

    def __init__(self, transforms, ctx=None):
        """init"""
        self.transforms = transforms
        self.ctx = ctx

    def __call__(self, data):
        """call"""
        ctx = self.ctx if self.ctx else {}
        for f in self.transforms:
            try:
                data = f(data, ctx)
            except Exception as e:
                stack_info = traceback.format_exc()
                logger.warning("fail to map op [{}] with error: {} and stack:\n{}".format(f, e, str(stack_info)))
                raise e
        return data


def batch_arrange(batch_samples, fields):
    def _segm(samples):
        """"""
        assert "gt_poly" in samples
        segms = samples["gt_poly"]
        if "is_crowd" in samples:
            is_crowd = samples["is_crowd"]
            if len(segms) != 0:
                assert len(segms) == is_crowd.shape[0]

        gt_masks = []
        valid = True
        for i in range(len(segms)):
            segm = segms[i]
            gt_segm = []
            if "is_crowd" in samples and is_crowd[i]:
                gt_segm.append([[0, 0]])
            else:
                for poly in segm:
                    if len(poly) == 0:
                        valid = False
                        break
                    gt_segm.append(np.array(poly).reshape(-1, 2))
            if (not valid) or len(gt_segm) == 0:
                break
            gt_masks.append(gt_segm)
        return gt_masks

    def im_shape(samples, dim=3):
        # hard code
        assert "h" in samples
        assert "w" in samples
        if dim == 3:  # RCNN, ..
            return np.array((samples["h"], samples["w"], 1), dtype=np.float32)
        else:  # YOLOv3, ..
            return np.array((samples["h"], samples["w"]), dtype=np.int32)

    arrange_batch = []
    for samples in batch_samples:
        one_ins = ()
        for i, field in enumerate(fields):
            if field == "gt_mask":
                one_ins += (_segm(samples),)
            elif field == "im_shape":
                one_ins += (im_shape(samples),)
            elif field == "im_size":
                one_ins += (im_shape(samples, 2),)
            else:
                if field == "is_difficult":
                    field = "difficult"
                assert field in samples, "{} not in samples".format(field)
                one_ins += (samples[field],)
        arrange_batch.append(one_ins)
    return arrange_batch


class ProcessReader(object):
    """
    Args:
        dataset (DataSet): DataSet object
        sample_transforms (list of BaseOperator): a list of sample transforms
            operators.
        batch_transforms (list of BaseOperator): a list of batch transforms
            operators.
        batch_size (int): batch size.
        shuffle (bool): whether shuffle dataset or not. Default False.
        drop_last (bool): whether drop last batch or not. Default False.
        drop_empty (bool): whether drop sample when it's gt is empty or not.
            Default True.
        mixup_epoch (int): mixup epoc number. Default is -1, meaning
            not use mixup.
        cutmix_epoch (int): cutmix epoc number. Default is -1, meaning
            not use cutmix.
        class_aware_sampling (bool): whether use class-aware sampling or not.
            Default False.
        worker_num (int): number of working threads/processes.
            Default -1, meaning not use multi-threads/multi-processes.
        use_process (bool): whether use multi-processes or not.
            It only works when worker_num > 1. Default False.
        bufsize (int): buffer size for multi-threads/multi-processes,
            please note, one instance in buffer is one batch data.
        memsize (str): size of shared memory used in result queue when
            use_process is true. Default 3G.
        inputs_def (dict): network input definition use to get input fields,
            which is used to determine the order of returned data.
        devices_num (int): number of devices.
        num_trainers (int): number of trainers. Default 1.
    """

    def __init__(
        self,
        dataset=None,
        sample_transforms=None,
        batch_transforms=None,
        batch_size=None,
        shuffle=False,
        drop_last=False,
        drop_empty=True,
        mixup_epoch=-1,
        cutmix_epoch=-1,
        class_aware_sampling=False,
        use_process=False,
        use_fine_grained_loss=False,
        num_classes=80,
        bufsize=-1,
        memsize="3G",
        inputs_def=None,
        devices_num=1,
        num_trainers=1,
    ):
        """"""
        self._fields = copy.deepcopy(inputs_def["fields"]) if inputs_def else None

        # transform
        self._sample_transforms = Compose(sample_transforms, {"fields": self._fields})
        self._batch_transforms = None

        if batch_transforms:
            batch_transforms = [bt for bt in batch_transforms]
            self._batch_transforms = Compose(batch_transforms, {"fields": self._fields})

        self._batch_size = batch_size
        self._shuffle = shuffle
        self._drop_last = drop_last
        self._drop_empty = drop_empty

        # sampling
        self._mixup_epoch = mixup_epoch // num_trainers
        self._cutmix_epoch = cutmix_epoch // num_trainers
        self._class_aware_sampling = class_aware_sampling

        self._indexes = None
        self._pos = -1
        self._epoch = -1
        self._curr_iter = 0

    def process(self, dataset):
        """process"""
        batch = self._load_batch(dataset)
        res = self.worker(self._drop_empty, batch)
        return res

    def _load_batch(self, dataset):
        batch = []
        for data in dataset:
            sample = copy.deepcopy(data)
            batch.append(sample)
        return batch

    def worker(self, drop_empty=True, batch_samples=None):
        """
        sample transform and batch transform.
        """
        batch = []
        for sample in batch_samples:
            sample = self._sample_transforms(sample)
            batch.append(sample)
        if len(batch) > 0 and self._batch_transforms:
            batch = self._batch_transforms(batch)
        if len(batch) > 0 and self._fields:
            batch = batch_arrange(batch, self._fields)
        return batch


def pad_batch_data(
    insts,
    pad_idx=0,
    max_seq_len=None,
    return_pos=False,
    return_input_mask=False,
    return_max_len=False,
    return_num_token=False,
    return_seq_lens=False,
    pad_2d_pos_ids=False,
    pad_segment_id=False,
    select=False,
    extract=False,
):
    """
    Pad the instances to the max sequence length in batch, and generate the
    corresponding position data and attention bias.
    """
    return_list = []
    max_len = max(len(inst) for inst in insts) if max_seq_len is None else max_seq_len
    # Any token included in dict can be used to pad, since the paddings' loss
    # will be masked out by weights and make no effect on parameter gradients.
    if pad_2d_pos_ids:
        boxes = [x + [[0, 0, 0, 0]] * (max_len - len(x)) for x in insts]
        boxes = np.array(boxes, dtype="int64")
        return boxes

    inst_data = np.array([inst + list([pad_idx] * (max_len - len(inst))) for inst in insts])
    return_list += [inst_data.astype("int64").reshape([-1, max_len, 1])]

    # position data
    if return_pos:
        inst_pos = np.array([list(range(0, len(inst))) + [pad_idx] * (max_len - len(inst)) for inst in insts])

        return_list += [inst_pos.astype("int64").reshape([-1, max_len, 1])]

    if return_input_mask:
        # This is used to avoid attention on paddings.
        input_mask_data = np.array([[1] * len(inst) + [0] * (max_len - len(inst)) for inst in insts])
        input_mask_data = np.expand_dims(input_mask_data, axis=-1)
        return_list += [input_mask_data.astype("float32")]

    if return_max_len:
        return_list += [max_len]

    if return_num_token:
        num_token = 0
        for inst in insts:
            num_token += len(inst)
        return_list += [num_token]

    if return_seq_lens:
        seq_lens = np.array([len(inst) for inst in insts])
        return_list += [seq_lens.astype("int64").reshape([-1, 1])]

    return return_list if len(return_list) > 1 else return_list[0]


class ImageReader(object):
    def __init__(
        self,
        super_rel_pos,
        tokenizer,
        max_key_len=16,
        max_seq_len=512,
        image_size=1024,
        block_w=7,
        block_h=7,
        im_npos=224,
    ):
        self.tokenizer = tokenizer
        self.vocab = self.tokenizer.get_vocab()

        self.pad_id = self.vocab["[PAD]"]
        self.cls_id = self.vocab["[CLS]"]
        self.sep_id = self.vocab["[SEP]"]
        self.mask_id = self.vocab["[MASK]"]
        self.pad = "[PAD]"
        self.cls = "[CLS]"
        self.sep = "[SEP]"
        self.mask = "[MASK]"

        self.super_rel_pos = super_rel_pos
        self.max_key_len = max_key_len
        self.max_seq_len = max_seq_len
        self.doc_stride = 128
        self.unique_id = 10000000

        self.examples = {}
        self.features = {}

        self.image_size = image_size
        self.block_w = block_w
        self.block_h = block_h
        self.im_npos = im_npos
        self.image_rois = []
        cut_width, cut_height = int(self.image_size / self.block_w), int(self.image_size / self.block_h)
        for idh in range(self.block_h):
            for idw in range(self.block_w):
                self.image_rois.append([idw * cut_width, idh * cut_height, cut_width, cut_height])

        sample_trans = [
            DecodeImage(),
            ResizeImage(target_size=self.im_npos, interp=1),
            NormalizeImage(
                is_channel_first=False,
                mean=[103.530, 116.280, 123.675],
                std=[57.375, 57.120, 58.395],
            ),
            Permute(to_bgr=False),
        ]

        batch_trans = [PadBatch(pad_to_stride=32, use_padded_im_info=True)]

        inputs_def = {
            "fields": ["image", "im_info", "im_id", "gt_bbox"],
        }
        self.data_loader = ProcessReader(
            sample_transforms=sample_trans,
            batch_transforms=batch_trans,
            shuffle=False,
            drop_empty=True,
            inputs_def=inputs_def,
        )

    def ppocr2example(self, ocr_res, img_path, queries):
        examples = []
        segments = []
        for rst in ocr_res:
            left = min(rst[0][0][0], rst[0][3][0])
            top = min(rst[0][0][-1], rst[0][1][-1])
            width = max(rst[0][1][0], rst[0][2][0]) - min(rst[0][0][0], rst[0][3][0])
            height = max(rst[0][2][-1], rst[0][3][-1]) - min(rst[0][0][-1], rst[0][1][-1])
            segments.append({"bbox": Bbox(*[left, top, width, height]), "text": rst[-1][0]})
        segments.sort(key=cmp_to_key(two_dimension_sort_layout))
        # 2. im_base64
        img_base64 = img2base64(img_path)
        # 3. doc_tokens, doc_boxes, segment_ids
        doc_tokens = []
        doc_boxes = []
        ori_boxes = []
        doc_segment_ids = []

        im_w_box = max([seg["bbox"].left + seg["bbox"].width for seg in segments]) + 20
        im_h_box = max([seg["bbox"].top + seg["bbox"].height for seg in segments]) + 20
        img = Image.open(img_path)
        im_w, im_h = img.size
        im_w, im_h = max(im_w, im_w_box), max(im_h, im_h_box)

        scale_x = self.image_size / im_w
        scale_y = self.image_size / im_h
        for segment_id, segment in enumerate(segments):
            bbox = segment["bbox"]  # x, y, w, h
            x1, y1, w, h = bbox.left, bbox.top, bbox.width, bbox.height
            sc_w = int(min(w * scale_x, self.image_size - 1))
            sc_h = int(min(h * scale_y, self.image_size - 1))
            sc_y1 = int(max(0, min(y1 * scale_y, self.image_size - h - 1)))
            sc_x1 = int(max(0, min(x1 * scale_x, self.image_size - w - 1)))
            if w < 0:
                raise ValueError("Incorrect bbox, please check the input word boxes.")
            ori_bbox = Bbox(*[x1, y1, w, h])
            sc_bbox = Bbox(*[sc_x1, sc_y1, sc_w, sc_h])
            text = segment["text"]
            char_num = []
            eng_word = ""
            for char in text:
                if not check(char) and not eng_word:
                    doc_tokens.append([char])
                    doc_segment_ids.append([segment_id])
                    char_num.append(2)
                elif not check(char) and eng_word:
                    doc_tokens.append([eng_word])
                    doc_segment_ids.append([segment_id])
                    char_num.append(len(eng_word))
                    eng_word = ""
                    doc_tokens.append([char])
                    doc_segment_ids.append([segment_id])
                    char_num.append(2)
                else:
                    eng_word += char
            if eng_word:
                doc_tokens.append([eng_word])
                doc_segment_ids.append([segment_id])
                char_num.append(len(eng_word))
            ori_char_width = round(ori_bbox.width / sum(char_num), 1)
            sc_char_width = round(sc_bbox.width / sum(char_num), 1)
            for chr_idx in range(len(char_num)):
                if chr_idx == 0:
                    doc_boxes.append(
                        [Bbox(*[sc_bbox.left, sc_bbox.top, (sc_char_width * char_num[chr_idx]), sc_bbox.height])]
                    )
                    ori_boxes.append(
                        [Bbox(*[ori_bbox.left, ori_bbox.top, (ori_char_width * char_num[chr_idx]), ori_bbox.height])]
                    )
                else:
                    doc_boxes.append(
                        [
                            Bbox(
                                *[
                                    sc_bbox.left + (sc_char_width * sum(char_num[:chr_idx])),
                                    sc_bbox.top,
                                    (sc_char_width * char_num[chr_idx]),
                                    sc_bbox.height,
                                ]
                            )
                        ]
                    )
                    ori_boxes.append(
                        [
                            Bbox(
                                *[
                                    ori_bbox.left + (ori_char_width * sum(char_num[:chr_idx])),
                                    ori_bbox.top,
                                    (ori_char_width * char_num[chr_idx]),
                                    ori_bbox.height,
                                ]
                            )
                        ]
                    )

        qas_id = 0
        for query in queries:
            example = Example(
                keys=[query],
                key_labels=[0],
                doc_tokens=doc_tokens,
                seq_labels=[0 for one in doc_tokens],
                text="",
                qas_id="0_" + str(qas_id),
                model_type=None,
                ori_boxes=ori_boxes,
                boxes=doc_boxes,
                segment_ids=doc_segment_ids,
                symbol_ids=None,
                image_rois=self.image_rois,
                im_base64=img_base64,
            )

            examples.append(example)
            qas_id += 1
        return examples

    def box2example(self, ocr_res, img_path, queries):
        """
        ocr_res = [[word_str, [x1, y1, x2, y2]], [word_str, [x1, y1, x2, y2]], ...]
        """
        examples = []
        doc_boxes = []
        ori_boxes = []
        boxes = [x[1] for x in ocr_res]
        im_w_box = max([b[2] for b in boxes]) + 20
        im_h_box = max([b[3] for b in boxes]) + 20
        img = Image.open(img_path)
        im_w, im_h = img.size
        im_w, im_h = max(im_w, im_w_box), max(im_h, im_h_box)

        scale_x = self.image_size / im_w
        scale_y = self.image_size / im_h
        for box in boxes:
            x1, y1, x2, y2 = box
            if x2 <= x1 or y2 <= y1:
                raise ValueError("Invalid bbox format")
            w = max(x1, x2) - min(x1, x2)
            h = max(y1, y2) - min(y1, y2)
            ori_boxes.append([Bbox(*[x1, y1, w, h])])
            w = int(min(w * scale_x, self.image_size - 1))
            h = int(min(h * scale_y, self.image_size - 1))
            x1 = int(max(0, min(x1 * scale_x, self.image_size - w - 1)))
            y1 = int(max(0, min(y1 * scale_y, self.image_size - h - 1)))
            if w < 0:
                raise ValueError("Invalid bbox format")
            doc_boxes.append([Bbox(*[x1, y1, w, h])])

        img_base64 = img2base64(img_path)

        doc_tokens = [[x[0]] for x in ocr_res]
        doc_segment_ids = [[0]] * len(doc_tokens)

        qas_id = 0
        for query in queries:
            example = Example(
                keys=[query],
                key_labels=[0],
                doc_tokens=doc_tokens,
                seq_labels=[0 for one in doc_tokens],
                text="",
                qas_id=str(qas_id),
                model_type=None,
                ori_boxes=ori_boxes,
                boxes=doc_boxes,
                segment_ids=doc_segment_ids,
                symbol_ids=None,
                image_rois=self.image_rois,
                im_base64=img_base64,
            )

            if not (len(example.doc_tokens) == len(example.boxes) == len(example.segment_ids)):
                raise ValueError(
                    "Incorrect word_boxes, the format should be `List[str, Tuple[float, float, float, float]]`"
                )

            examples.append(example)
            qas_id += 1

        return examples

    def example2feature(self, example, tokenizer, max_line_id=128):
        features = []
        all_doc_tokens = []
        tok_to_orig_index = []
        boxes = []
        segment_ids = []
        all_doc_labels = []

        query_tokens = tokenizer.tokenize("&" + str(example.keys[0]))[1:][: self.max_key_len]

        for i, (token_list, box_list, seg_list, l) in enumerate(
            zip(example.doc_tokens, example.boxes, example.segment_ids, example.seq_labels)
        ):
            assert len(token_list) == len(box_list) == len(seg_list)
            for idt, (token, box, seg) in enumerate(zip(token_list, box_list, seg_list)):
                sub_tokens = tokenizer.tokenize("&" + token)[1:]
                for ii, sub_token in enumerate(sub_tokens):
                    width_split = box.width / len(sub_tokens)
                    boxes.append([box.left + ii * width_split, box.top, width_split, box.height])
                    segment_ids.append(seg)
                    tok_to_orig_index.append(i)
                    all_doc_tokens.append(sub_token)
                    all_doc_labels.extend([0])

        max_tokens_for_doc = self.max_seq_len - len(query_tokens) - 4
        doc_spans = []
        start_offset = 0
        while start_offset < len(all_doc_tokens):
            length = len(all_doc_tokens) - start_offset
            if length > max_tokens_for_doc:
                length = max_tokens_for_doc
            doc_spans.append(DocSpan(start=start_offset, length=length))
            if start_offset + length == len(all_doc_tokens):
                break
            start_offset += min(length, self.doc_stride)

        for (doc_span_index, doc_span) in enumerate(doc_spans):
            tokens = []
            labels = []
            feature_segment_ids = []
            feature_boxes = []
            token_to_orig_map = {}
            token_is_max_context = {}
            text_type_ids = []
            tokens.append(self.cls)
            feature_boxes.append([0, 0, 0, 0])
            labels.append(0)
            text_type_ids.append(0)
            feature_segment_ids.append(max_line_id - 1)

            for i in range(doc_span.length):
                split_token_index = doc_span.start + i
                token_to_orig_map[len(tokens)] = tok_to_orig_index[split_token_index]
                is_max_context = self._check_is_max_context(doc_spans, doc_span_index, split_token_index)
                token_is_max_context[len(tokens)] = is_max_context
                tokens.append(all_doc_tokens[split_token_index])

                feature_boxes.append(boxes[split_token_index])
                feature_segment_ids.append(segment_ids[split_token_index])
                text_type_ids.append(0)
                labels.append(all_doc_labels[split_token_index])

            tokens.append(self.sep)
            feature_boxes.append([0, 0, 0, 0])
            text_type_ids.append(0)
            feature_segment_ids.append(max_line_id - 1)
            labels.append(0)
            for token in query_tokens:
                tokens.append(token)
                feature_boxes.append([0, 0, 0, 0])
                feature_segment_ids.append(max_line_id - 1)
                text_type_ids.append(1)
                labels.append(0)

            tokens = tokens + [self.sep]
            feature_boxes.extend([[0, 0, 0, 0]])
            feature_segment_ids = feature_segment_ids + [max_line_id - 1]
            text_type_ids = text_type_ids + [1]
            labels.append(0)

            position_ids = list(range(len(tokens)))
            token_ids = tokenizer.convert_tokens_to_ids(tokens)
            feature_segment_ids = [x % max_line_id for x in feature_segment_ids]

            feature = Feature(
                unique_id=self.unique_id,
                example_index=0,
                qas_id=example.qas_id,
                doc_span_index=doc_span_index,
                tokens=tokens,
                token_to_orig_map=token_to_orig_map,
                token_is_max_context=token_is_max_context,
                token_ids=token_ids,
                position_ids=position_ids,
                text_type_ids=text_type_ids,
                text_symbol_ids=None,
                overlaps=None,
                keys=example.keys,
                seq_labels=labels,
                se_seq_labels=None,
                bio_seq_labels=None,
                bioes_seq_labels=None,
                key_labels=example.key_labels,
                model_type=example.model_type,
                doc_tokens=example.doc_tokens,
                doc_labels=example.seq_labels,
                text=example.text,
                boxes=feature_boxes,
                segment_ids=feature_segment_ids,
                im_base64=example.im_base64,
                image_rois=example.image_rois,
            )
            features.append(feature)
            self.unique_id += 1
        return features

    def _pad_batch_records(self, batch_records, max_line_id=128, phase="infer"):
        """pad batch records"""
        return_list = []
        batch_token_ids = []
        batch_sent_ids = []
        batch_pos_ids = []
        batch_2d_pos_ids = []
        batch_segment_ids = []
        batch_labels = []
        batch_unique_id = []
        batch_image_base64 = []
        batch_image_rois = []

        for i in range(len(batch_records)):
            batch_token_ids.append(batch_records[i].token_ids)
            batch_sent_ids.append(batch_records[i].text_type_ids)
            batch_segment_ids.append(batch_records[i].segment_ids)
            batch_labels.append(batch_records[i].seq_labels)
            batch_unique_id.append(batch_records[i].unique_id)
            batch_pos_ids.append(batch_records[i].position_ids)
            batch_2d_pos_ids.append(batch_records[i].boxes)
            batch_image_base64.append(batch_records[i].im_base64)
            batch_image_rois.append(batch_records[i].image_rois)

        padded_token_ids, _ = pad_batch_data(batch_token_ids, pad_idx=self.pad_id, return_input_mask=True)
        padded_sent_ids = pad_batch_data(batch_sent_ids, pad_idx=self.pad_id)
        padded_pos_ids = pad_batch_data(batch_pos_ids, pad_idx=self.pad_id)
        new_padded_pos_ids = []
        for idp, pos_ids in enumerate(padded_pos_ids):
            new_padded_pos_ids.append(
                np.concatenate((pos_ids, np.array([[x] for x in range(self.block_w * self.block_h)])), axis=0)
            )
        padded_pos_ids = np.array(new_padded_pos_ids)
        padded_2d_pos_ids = pad_batch_data(batch_2d_pos_ids, pad_2d_pos_ids=True, select=False, extract=True)
        new_padded_2d_pos_ids = []
        for pos_ids_2d, batch_record in zip(padded_2d_pos_ids, batch_records):
            new_padded_2d_pos_ids.append(np.concatenate((pos_ids_2d, np.array(batch_record.image_rois)), axis=0))
        padded_2d_pos_ids = np.array(new_padded_2d_pos_ids)
        padded_segment_ids = pad_batch_data(batch_segment_ids, pad_idx=max_line_id - 1)

        input_mask_mat = self._build_input_mask(
            np.array([list(x) + [[-1] for _ in range(self.block_w * self.block_h)] for x in padded_token_ids])
        )
        super_rel_pos = self._build_rel_pos(
            np.array([list(x) + [[-1] for _ in range(self.block_w * self.block_h)] for x in padded_token_ids])
        )

        unique_id = np.array(batch_unique_id).astype("float32").reshape([-1, 1])

        bsz, seq_len, _ = padded_token_ids.shape
        task_ids = np.ones((bsz, seq_len, 1)).astype("int64")
        for b in range(bsz):
            if np.sum(padded_2d_pos_ids[b]) > 0:
                task_ids[b, :, :] = 0
            else:
                task_ids[b, :, :] = 1

        coco_data = self.generate_coco_data(
            [""] * len(batch_image_base64),
            batch_image_base64,
            [self.image_size] * len(batch_image_base64),
            [self.image_size] * len(batch_image_base64),
            batch_image_rois,
        )

        image_data = self.im_make_batch(
            self.data_loader.process(coco_data),
            self.block_w * self.block_h,
            len(batch_image_base64),
        )

        return_list = [
            padded_token_ids,
            padded_sent_ids,
            padded_pos_ids,
            padded_2d_pos_ids,
            padded_segment_ids,
            task_ids,
            input_mask_mat,
            super_rel_pos,
            unique_id,
            image_data,
        ]
        return return_list

    def data_generator(self, ocr_res, img_path, queries, batch_size, ocr_type="ppocr", phase="infer"):
        if ocr_type == "ppocr":
            self.examples[phase] = self.ppocr2example(ocr_res, img_path, queries)
        elif ocr_type == "word_boxes":
            self.examples[phase] = self.box2example(ocr_res, img_path, queries)
        self.features[phase] = sum([self.example2feature(e, self.tokenizer) for e in self.examples[phase]], [])
        for batch_data in self._prepare_batch_data(self.features[phase], batch_size, phase=phase):
            yield self._pad_batch_records(batch_data)

    def _prepare_batch_data(self, features, batch_size, phase=None):
        """generate batch records"""
        batch_records = []
        for feature in features:
            to_append = len(batch_records) < batch_size
            if to_append:
                batch_records.append(feature)
            else:
                yield batch_records
                batch_records = [feature]

        if phase == "infer" and batch_records:
            yield batch_records

    def _build_input_mask(self, padded_token_ids):
        """build_input_mask"""
        bsz, seq_len, _ = padded_token_ids.shape
        return np.ones((bsz, seq_len, seq_len)).astype("float32")

    def _build_rel_pos(self, padded_token_ids):
        """build relative position"""
        bsz, seq_len, _ = padded_token_ids.shape
        rel_pos = np.zeros((bsz, seq_len, seq_len)).astype("int64")
        return rel_pos

    def generate_coco_data(
        self,
        batch_image_path,
        batch_image_base64,
        batch_scaled_width,
        batch_scaled_height,
        batch_rois,
    ):
        """generator coco data"""

        def transform(dataset):
            roidbs = []
            for i in dataset:
                rvl_rec = {
                    "im_file": i[0],
                    "im_id": np.array([i[1]]),
                    "h": i[2],
                    "w": i[3],
                    "gt_bbox": i[4],
                    "cover_box": i[5],
                    "im_base64": i[6],
                }

                roidbs.append(rvl_rec)
            return roidbs

        result = []
        for image_path, im_base64, width, height, roi in zip(
            batch_image_path,
            batch_image_base64,
            batch_scaled_width,
            batch_scaled_height,
            batch_rois,
        ):
            result.append((image_path, 0, height, width, roi, None, im_base64))
        return transform(result)

    def im_make_batch(self, dataset, image_boxes_nums, bsize):
        """make image batch"""
        img_batch = np.array([i[0] for i in dataset], "float32")
        return img_batch

    def BIO2SPAN(self, BIO):
        start_label, end_label = [], []
        for seq in BIO:
            first_one = True
            start_pos = [1 if x == 2 else 0 for x in seq]
            if sum(start_pos) == 0 and sum(seq) != 0:
                start_pos = []
                for idp, p in enumerate(seq):
                    if p == 1 and first_one:
                        start_pos.append(1)
                        first_one = False
                    else:
                        start_pos.append(0)

            start_label.append(start_pos)

            end_tmp = []
            for index, s in enumerate(seq):
                if s == -100 or s == 0:
                    end_tmp.append(s)
                elif s == 2 and index + 1 < len(seq) and (seq[index + 1] == 0 or seq[index + 1] == 2):
                    end_tmp.append(1)
                elif s == 2 and index + 1 < len(seq) and seq[index + 1] != 0:
                    end_tmp.append(0)
                elif s == 2 and index + 1 == len(seq):
                    end_tmp.append(1)
                elif s == 1 and (index + 1 == len(seq) or seq[index + 1] != 1):
                    end_tmp.append(1)
                else:
                    end_tmp.append(0)
            end_label.append(end_tmp)

        return start_label, end_label

    def _check_is_max_context(self, doc_spans, cur_span_index, position):
        best_score = None
        best_span_index = None
        for (span_index, doc_span) in enumerate(doc_spans):
            end = doc_span.start + doc_span.length - 1
            if position < doc_span.start:
                continue
            if position > end:
                continue
            num_left_context = position - doc_span.start
            num_right_context = end - position
            score = min(num_left_context, num_right_context) + 0.01 * doc_span.length
            if best_score is None or score > best_score:
                best_score = score
                best_span_index = span_index
        return cur_span_index == best_span_index


def get_doc_pred(result, ans_pos, example, tokenizer, feature, do_lower_case, all_key_probs, example_index):
    def _compute_softmax(scores):
        """Compute softmax probability over raw logits."""
        if len(scores) == 0:
            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

    preds = []
    for start_index, end_index in ans_pos:
        # process data
        tok_tokens = feature.tokens[start_index : end_index + 1]
        tok_text = " ".join(tok_tokens)
        # De-tokenize WordPieces that have been split off.
        tok_text = tok_text.replace(" ##", "")
        tok_text = tok_text.replace("##", "")
        tok_text = tok_text.strip()
        tok_text = "".join(tok_text.split())

        orig_doc_start = feature.token_to_orig_map[start_index]
        orig_doc_end = feature.token_to_orig_map[end_index]
        orig_tokens = example.doc_tokens[orig_doc_start : orig_doc_end + 1]

        # Clean whitespace
        orig_text = "".join(["".join(x) for x in orig_tokens])
        final_text = get_final_text(tok_text, orig_text, tokenizer, do_lower_case)

        probs = []
        for idx, logit in enumerate(result.seq_logits[start_index : end_index + 1]):
            if idx == 0:
                # -1 is for B in  OIB or I in OI
                probs.append(_compute_softmax(logit)[-1])
            else:
                # 1 is for I in OIB or I in OI
                probs.append(_compute_softmax(logit)[1])
        avg_prob = sum(probs) / len(probs)
        preds.append({"value": final_text, "prob": round(avg_prob, 2), "start": orig_doc_start, "end": orig_doc_end})
    return preds


def get_final_text(pred_text, orig_text, tokenizer, do_lower_case):
    """Project the tokenized prediction back to the original text."""

    def _strip_spaces(text):
        ns_chars = []
        ns_to_s_map = 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)

    tok_text = " ".join(tokenizer.tokenize(orig_text))

    start_position = tok_text.find(pred_text)
    if start_position == -1:
        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):
        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 six.iteritems(tok_ns_to_s_map):
        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:
        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:
        return orig_text

    output_text = orig_text[orig_start_position : (orig_end_position + 1)]
    return output_text


def find_bio_pos(label):
    """find answer position from BIO label"""
    e = []
    cand_ans = []
    last_l = None
    for idx, l in enumerate(label):
        if l == "O":
            if e:
                cand_ans.append([e[0], e[-1]])
            e = []
        elif l.startswith("B"):
            if last_l == "O" or last_l is None:
                if len(e) != 0:
                    e = []
                e.append(idx)
            else:  # I B
                if e:
                    cand_ans.append([e[0], e[-1]])
                    e = []
                e.append(idx)
        elif l.startswith("I"):
            if len(e) == 0:
                continue
            else:
                e.append(idx)
        last_l = l
    if e:
        cand_ans.append([e[0], e[-1]])
    return cand_ans


def viterbi_decode(logits):
    np_logits = np.array(logits)  # shape: L * D
    length, dim = np_logits.shape
    f = np.zeros(np_logits.shape)
    path = [["" for i in range(dim)] for j in range(length)]
    label_scheme = "OIB"
    # oib label 0:O, 1:I, 2:B
    # illegal matrix: [O, I ,B, start, end] * [O, I, B, start, end]
    illegal = np.array([[0, -1, 0, -1, 0], [0, 0, 0, -1, 0], [0, 0, 0, 0, 0], [0, -1, 0, 0, 0], [-1, -1, -1, -1, -1]])
    illegal = illegal * 1000

    f[0, :] = np_logits[0, :] + illegal[3, :3]
    path[0] = [label_scheme[i] for i in range(dim)]

    for step in range(1, length):
        last_s = f[step - 1, :]
        for d in range(dim):
            cand_score = illegal[:3, d] + last_s + np_logits[step, d]
            f[step, d] = np.max(cand_score)
            path[step][d] = path[step - 1][np.argmax(cand_score)] + label_scheme[d]
    final_path = path[-1][np.argmax(f[-1, :])]
    return final_path


def find_answer_pos(logits, feature):
    start_index = -1
    end_index = -1
    ans = []
    cand_ans = []

    best_path = viterbi_decode(logits)
    cand_ans = find_bio_pos(best_path)

    for start_index, end_index in cand_ans:
        is_valid = True
        if start_index not in feature.token_to_orig_map:
            is_valid = False
        if end_index not in feature.token_to_orig_map:
            is_valid = False
        if not feature.token_is_max_context.get(start_index, False):
            is_valid = False
        if end_index < start_index:
            is_valid = False
        if is_valid:
            ans.append([start_index, end_index])

    return ans


def calEuclidean(x_list, y_list):
    """
    Calculate euclidean distance
    """
    if x_list is None or y_list is None:
        return None
    else:
        dist = np.sqrt(np.square(x_list[0] - y_list[0]) + np.square(x_list[1] - y_list[1]))
        return dist


def longestCommonSequence(question_tokens, context_tokens):
    """
    Longest common sequence
    """
    max_index = -1
    max_len = 0
    m, n = len(question_tokens), len(context_tokens)
    dp = [[0] * (n + 1) for _ in range(m + 1)]
    for i in range(1, m + 1):
        for j in range(1, n + 1):
            if question_tokens[i - 1].lower() == context_tokens[j - 1][0].lower():
                dp[i][j] = 1 + dp[i - 1][j - 1]
                if dp[i][j] > max_len:
                    max_len = dp[i][j]
                    max_index = j - 1
    return max_index, max_len


def sort_res(prompt, ans_list, context, boxes, lang="en"):
    if len(ans_list) == 1:
        return ans_list
    else:
        ans_val = []
        for ans in ans_list:
            ans_val.append(ans["value"])
        if len(set(ans_val)) == len(ans_val):
            sorted_ans_list = sorted(ans_list, key=lambda x: x["prob"], reverse=True)
            return sorted_ans_list
        else:
            if lang == "en":
                clean_prompt = [word for word in prompt.split(" ")]
            else:
                clean_prompt = [word for word in prompt]

            max_index, max_len = longestCommonSequence(clean_prompt, context)
            if max_index == -1:
                sorted_ans_list = sorted(ans_list, key=lambda x: x["prob"], reverse=True)
                return sorted_ans_list
            else:
                prompt_center = []
                for idx in range(max_index - max_len + 1, max_index + 1):
                    box = boxes[idx][0]
                    x = box.left + box.width / 2
                    y = box.top + box.height / 2
                    prompt_center.append([x, y])

                ans_center = []
                ans_prob = []
                for ans in ans_list:
                    ans_prob.append(ans["prob"])
                    cent_list = []
                    for idx in range(ans["start"], ans["end"] + 1):
                        box = boxes[idx][0]
                        x = box.left + box.width / 2
                        y = box.top + box.height / 2
                        cent_list.append([x, y])
                    ans_center.append(cent_list)

                ans_odist = []
                for ans_c in ans_center:
                    odist = 0
                    for a_c in ans_c:
                        for p_c in prompt_center:
                            odist += calEuclidean(a_c, p_c)
                    odist /= len(ans_c)
                    ans_odist.append(odist * (-1))

                ans_score = np.sum([ans_prob, ans_odist], axis=0).tolist()
                sorted_ans_list = sorted(ans_list, key=lambda x: ans_score[ans_list.index(x)], reverse=True)
                return sorted_ans_list