unilm/layoutreader/s2s_ft/utils.py

from __future__ import absolute_import, division, print_function

import logging
import os
import json
import random
import glob
import re

import torch
import tqdm
import torch.utils.data

logger = logging.getLogger(__name__)


class Seq2seqDatasetForBert(torch.utils.data.Dataset):
    def __init__(
            self, features, max_source_len, max_target_len,
            vocab_size, cls_id, sep_id, pad_id, mask_id,
            random_prob, keep_prob, offset, num_training_instances,
            span_len=1, span_prob=1.0):
        self.features = features
        self.max_source_len = max_source_len
        self.max_target_len = max_target_len
        self.offset = offset
        if offset > 0:
            logger.info("  ****  Set offset %d in Seq2seqDatasetForBert ****  ", offset)
        self.cls_id = cls_id
        self.sep_id = sep_id
        self.pad_id = pad_id
        self.random_prob = random_prob
        self.keep_prob = keep_prob
        self.mask_id = mask_id
        self.vocab_size = vocab_size
        self.num_training_instances = num_training_instances
        self.span_len = span_len
        self.span_prob = span_prob

    def __len__(self):
        return int(self.num_training_instances)

    def __trunk(self, ids, max_len):
        if len(ids) > max_len - 1:
            ids = ids[:max_len - 1]
        ids = ids + [self.sep_id]
        return ids

    def __pad(self, ids, max_len):
        if len(ids) < max_len:
            return ids + [self.pad_id] * (max_len - len(ids))
        else:
            assert len(ids) == max_len
            return ids

    def __getitem__(self, idx):
        idx = (self.offset + idx) % len(self.features)
        feature = self.features[idx]
        source_ids = self.__trunk([self.cls_id] + feature["source_ids"], self.max_source_len)
        target_ids = self.__trunk(feature["target_ids"], self.max_target_len)
        pseudo_ids = []
        for tk_id in target_ids:
            p = random.random()
            if p < self.keep_prob:
                pseudo_ids.append(tk_id)
            elif p < self.keep_prob + self.random_prob:
                pseudo_ids.append(random.randint(0, self.vocab_size - 1))
            else:
                pseudo_ids.append(self.mask_id)

        num_source_tokens = len(source_ids)
        num_target_tokens = len(target_ids)

        source_ids = self.__pad(source_ids, self.max_source_len)
        target_ids = self.__pad(target_ids, self.max_target_len)
        pseudo_ids = self.__pad(pseudo_ids, self.max_target_len)

        if self.span_len > 1:
            span_ids = []
            span_id = 1
            while len(span_ids) < num_target_tokens:
                p = random.random()
                if p < self.span_prob:
                    span_len = random.randint(2, self.span_len)
                    span_len = min(span_len, num_target_tokens - len(span_ids))
                else:
                    span_len = 1
                span_ids.extend([span_id] * span_len)
                span_id += 1
            span_ids = self.__pad(span_ids, self.max_target_len)
            return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, span_ids
        else:
            return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens


# DONE: finish this!!! the 2D input id settings.
class Seq2seqDatasetForLayoutlm(torch.utils.data.Dataset):
    def __init__(
            self, features, max_source_len, max_target_len,
            vocab_size, cls_id, sep_id, pad_id, mask_id,
            random_prob, keep_prob, offset, num_training_instances, layout_flag=True,
            span_len=1, span_prob=1.0):

        self.layout_flag = layout_flag

        self.features = features
        self.max_source_len = max_source_len
        self.max_target_len = max_target_len
        self.offset = offset
        if offset > 0:
            logger.info("  ****  Set offset %d in Seq2seqDatasetForBert ****  ", offset)
        self.cls_id = cls_id
        self.sep_id = sep_id
        self.pad_id = pad_id
        self.random_prob = random_prob
        self.keep_prob = keep_prob
        self.mask_id = mask_id
        self.vocab_size = vocab_size
        self.num_training_instances = num_training_instances
        self.span_len = span_len
        self.span_prob = span_prob

        self.index_sp_id = 0

    def __len__(self):
        return int(self.num_training_instances)

    def __clip_index(self, ids):
        replace_value = 0
        for i in range(len(ids)):
            if ids[i] > self.max_source_len - 1:
                ids[i] = replace_value
        return ids

    def __trunk(self, ids, max_len, simple=False, value=None):
        trunk_value = value if value is not None else self.sep_id
        if len(ids) > max_len - 1:
            ids = ids[:max_len - 1]
        if simple:
            ids = ids + [trunk_value]
        else:
            ids = ids + [[trunk_value, 1000, 1000, 1000, 1000]]
        return ids

    def __pad(self, ids, max_len, simple=False, value=None):
        pad_value = value if value is not None else self.pad_id
        if len(ids) < max_len:
            if simple:
                return ids + [pad_value] * (max_len - len(ids))
            else:
                return ids + [[pad_value, 0, 0, 0, 0]] * (max_len - len(ids))
        else:
            assert len(ids) == max_len
            return ids

    def __getitem__(self, idx):
        if self.layout_flag:
            return self.__getitem_layout__(idx)
        else:
            return self.__getitem_bert__(idx)

    def __getitem_bert__(self, idx):
        idx = (self.offset + idx) % len(self.features)
        feature = self.features[idx]
        source_ids = self.__trunk([self.cls_id] + feature["source_ids"], self.max_source_len, simple=True)
        target_ids = self.__trunk(feature["target_ids"], self.max_target_len, simple=True)
        target_index = self.__trunk(feature['target_index'], self.max_target_len, simple=True, value=self.index_sp_id)

        pseudo_ids = []
        for tk_id in target_ids:
            p = random.random()
            if p < self.keep_prob:
                pseudo_ids.append(tk_id)
            elif p < self.keep_prob + self.random_prob:
                pseudo_ids.append(random.randint(0, self.vocab_size - 1))
            else:
                pseudo_ids.append(self.mask_id)

        num_source_tokens = len(source_ids)
        num_target_tokens = len(target_ids)

        source_ids = self.__pad(source_ids, self.max_source_len, simple=True)
        target_ids = self.__pad(target_ids, self.max_target_len, simple=True)
        pseudo_ids = self.__pad(pseudo_ids, self.max_target_len, simple=True)
        target_index = self.__pad(target_index, self.max_target_len, simple=True, value=self.index_sp_id)
        target_index = self.__clip_index(target_index)

        if self.span_len > 1:
            span_ids = []
            span_id = 1
            while len(span_ids) < num_target_tokens:
                p = random.random()
                if p < self.span_prob:
                    span_len = random.randint(2, self.span_len)
                    span_len = min(span_len, num_target_tokens - len(span_ids))
                else:
                    span_len = 1
                span_ids.extend([span_id] * span_len)
                span_id += 1
            span_ids = self.__pad(span_ids, self.max_target_len)
            return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, span_ids, target_index
        else:
            return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, target_index

    def __getitem_layout__(self, idx):
        # TODO: how to initialize the random and masked tokens' pos emb
        #   Simple Solution: only mask the text
        idx = (self.offset + idx) % len(self.features)
        feature = self.features[idx]
        source_ids = self.__trunk([[self.cls_id, 0, 0, 0, 0]] + feature["source_ids"], self.max_source_len)
        target_ids = self.__trunk(feature["target_ids"], self.max_target_len)
        target_index = self.__trunk(feature['target_index'], self.max_target_len, simple=True, value=self.index_sp_id)

        pseudo_ids = []
        for tk_id in target_ids:
            p = random.random()
            if p < self.keep_prob:
                pseudo_ids.append(tk_id)
            elif p < self.keep_prob + self.random_prob:
                pseudo_ids.append([random.randint(0, self.vocab_size - 1)] + [0, 0, 0, 0])  # tk_id[1:])
            else:
                pseudo_ids.append([self.mask_id] + [0, 0, 0, 0])  # tk_id[1:])

        num_source_tokens = len(source_ids)
        num_target_tokens = len(target_ids)

        source_ids = self.__pad(source_ids, self.max_source_len)
        target_ids = self.__pad(target_ids, self.max_target_len)
        pseudo_ids = self.__pad(pseudo_ids, self.max_target_len)
        target_index = self.__pad(target_index, self.max_target_len, simple=True, value=self.index_sp_id)
        target_index = self.__clip_index(target_index)

        if self.span_len > 1:
            span_ids = []
            span_id = 1
            while len(span_ids) < num_target_tokens:
                p = random.random()
                if p < self.span_prob:
                    span_len = random.randint(2, self.span_len)
                    span_len = min(span_len, num_target_tokens - len(span_ids))
                else:
                    span_len = 1
                span_ids.extend([span_id] * span_len)
                span_id += 1
            span_ids = self.__pad(span_ids, self.max_target_len)
            return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, span_ids, target_index
        else:
            return source_ids, target_ids, pseudo_ids, num_source_tokens, num_target_tokens, target_index


def batch_list_to_batch_tensors(batch):
    batch_tensors = []
    for x in zip(*batch):
        if isinstance(x[0], torch.Tensor):
            batch_tensors.append(torch.stack(x))
        else:
            batch_tensors.append(torch.tensor(x, dtype=torch.long))
    return batch_tensors


def get_max_epoch_model(output_dir):
    fn_model_list = glob.glob(os.path.join(output_dir, "model.*.bin"))
    fn_optim_list = glob.glob(os.path.join(output_dir, "optim.*.bin"))
    if (not fn_model_list) or (not fn_optim_list):
        return None
    os.path.basename(output_dir)
    both_set = set([int(os.path.basename(fn).split('.')[1]) for fn in fn_model_list]
                   ) & set([int(os.path.basename(fn).split('.')[1]) for fn in fn_optim_list])
    if both_set:
        return max(both_set)
    else:
        return None


def load_and_cache_examples(
        example_file, tokenizer, local_rank, cached_features_file, shuffle=True):
    # Make sure only the first process in distributed training process the dataset, and the others will use the cache
    if local_rank not in [-1, 0]:
        torch.distributed.barrier()

    if cached_features_file is not None and os.path.exists(cached_features_file):
        logger.info("Loading features from cached file %s", cached_features_file)
        features = torch.load(cached_features_file)
    else:
        logger.info("Creating features from dataset file at %s", example_file)

        examples = []
        with open(example_file, mode="r", encoding="utf-8") as reader:
            for i, line in enumerate(reader):
                if i == 100:
                    break
                examples.append(json.loads(line))
        features = []

        for example in tqdm.tqdm(examples):
            if isinstance(example["src"], list):
                source_tokens = example["src"]
                target_tokens = example["tgt"]
            else:
                source_tokens = tokenizer.tokenize(example["src"])
                target_tokens = tokenizer.tokenize(example["tgt"])
            features.append({
                "source_ids": tokenizer.convert_tokens_to_ids(source_tokens),
                "target_ids": tokenizer.convert_tokens_to_ids(target_tokens),
            })

        if shuffle:
            random.shuffle(features)

        if local_rank in [-1, 0] and cached_features_file is not None:
            logger.info("Saving features into cached file %s", cached_features_file)
            torch.save(features, cached_features_file)

    # Make sure only the first process in distributed training process the dataset, and the others will use the cache
    if local_rank == 0:
        torch.distributed.barrier()

    return features


def load_and_cache_line_order_examples(
        example_path, tokenizer, local_rank, cached_features_file, max_src_length=1024,
        layout_flag=True, shuffle=True,
        src_shuffle_rate=0,
        file_info_flag=False,
        ):
    # Make sure only the first process in distributed training process the dataset, and the others will use the cache
    if local_rank not in [-1, 0]:
        torch.distributed.barrier()

    if cached_features_file is not None and os.path.exists(cached_features_file) and False:
        logger.info("Loading features from cached file %s", cached_features_file)
        features = torch.load(cached_features_file)
    else:
        logger.info("Creating features from dataset at %s", example_path)

        examples = []

        with open(example_path, 'r') as layout_reader:
            logger.info(f'Start loading {example_path}')
            for i, line in enumerate(layout_reader):
                examples.append(json.loads(line))

        features = []

        for layout in tqdm.tqdm(examples):
            bleu = layout['bleu']

            if random.random() < src_shuffle_rate:
                # print('Random!!!')
                # DONE: the random src! here has bug! index also need shuffle
                src_layout = layout['src']
                tgt_index = layout['tgt_index']

                source_length = len(src_layout)
                shuffle_index = list(range(source_length))
                random.shuffle(shuffle_index)

                shuffle_layout = ['' for _ in range(source_length)]
                for i, j in enumerate(shuffle_index):
                    # NOTE: map i-th token to j-th token
                    shuffle_layout[j] = src_layout[i]

                shuffle_target_index = [shuffle_index[i] for i in tgt_index]

                layout['tgt_index'] = shuffle_target_index
                layout['src'] = shuffle_layout

            mask = tokenizer.mask_token_id
            src_ids = [tokenizer.convert_tokens_to_ids([str(tmp_i)])[:1] + src_layout for tmp_i, src_layout in enumerate(layout['src'])]
            tgt_ids = [tokenizer.convert_tokens_to_ids([str(tmp_i)])[:1] + tgt_layout for tmp_i, tgt_layout in enumerate(layout['tgt'])]
            tgt_index = layout['tgt_index']

            feature = {
                "source_ids": src_ids,
                "target_ids": tgt_ids,
                "target_index": tgt_index,
                'bleu': bleu
            }

            if file_info_flag:
                file_info = {'original_filename': layout['filename'], 'filename': layout['filename'],
                             'page_idx': 0}
                feature['file_info'] = file_info

            features.append(feature)

        if shuffle:
            random.shuffle(features)

        if local_rank in [-1, 0] and cached_features_file is not None:
            logger.info("Saving features into cached file %s", cached_features_file)
            torch.save(features, cached_features_file)

    # Make sure only the first process in distributed training process the dataset, and the others will use the cache
    if local_rank == 0:
        torch.distributed.barrier()

    return features


def load_and_cache_layoutlm_examples(
        example_path, tokenizer, local_rank, cached_features_file, max_src_length=1024,
        layout_flag=True, shuffle=True,
        src_shuffle_rate=0,
        file_info_flag=False
        ):
    # Make sure only the first process in distributed training process the dataset, and the others will use the cache
    if local_rank not in [-1, 0]:
        torch.distributed.barrier()

    if cached_features_file is not None and os.path.exists(cached_features_file):
        logger.info("Loading features from cached file %s", cached_features_file)
        features = torch.load(cached_features_file)
    else:
        logger.info("Creating features from dataset at %s", example_path)

        examples = []

        if os.path.isdir(example_path):
            text_files = glob.glob(f'{example_path}/*text*.json')
            layout_files = [re.sub('text|txt', 'layout', x, 1) for x in text_files]
        else:
            text_files = [example_path]
            layout_files = [re.sub('text|txt', 'layout', example_path, 1)]
        for text_file, layout_file in zip(text_files, layout_files):
            with open(text_file, mode='r', encoding='utf-8') as text_reader, \
                    open(layout_file, mode='r', encoding='utf-8') as layout_reader:
                logger.info(f'Start loading {text_file}')
                for i, (text_line, layout_line) in enumerate(zip(text_reader, layout_reader)):
                    if (i + 1) % 10000 == 0:
                        logger.info(f'{i + 1} lines ...')
                    examples.append((json.loads(text_line), json.loads(layout_line)))

        features = []

        def tokenize_text_and_layout_src(_text, _layout, _layout_flag):
            ret = []
            index_split = {}
            words = _text.split()
            # note: (OLD) the index should start from 1: 0-the cls token in src
            # note: (NEW) we need to remove the src embedding's CLS SEP token so we can still start from 0
            # note: (NEWER) we need to at least one blank pos for ignore index in loss function (we use sep's index)
            # NOTE: (NEWER-ER) 1 for all padding tgt index
            new_token_index = 1  # first ordinary index
            for i, (word, box) in enumerate(zip(words, _layout)):

                if (not box[2] >= box[0]) or (not box[3] >= box[1]):
                    continue

                tokens = tokenizer.tokenize(word)
                tokens = tokenizer.convert_tokens_to_ids(tokens)
                new_token_ids = []
                for token in tokens:
                    if _layout_flag:
                        ret.append([token] + box)
                    else:
                        ret.append(token)
                    new_token_ids.append(new_token_index)
                    new_token_index += 1
                index_split[i] = new_token_ids

            return ret, index_split

        def tokenize_text_and_layout_tgt(_text, _layout, _index, _index_split, _layout_flag):
            ret = []
            ret_index = []
            words = _text.split()
            for word, box, i in zip(words, _layout, _index):

                if (not box[2] >= box[0]) or (not box[3] >= box[1]):
                    continue

                tokens = tokenizer.tokenize(word)
                tokens = tokenizer.convert_tokens_to_ids(tokens)
                for token, ii in zip(tokens, _index_split[i]):
                    if _layout_flag:
                        ret.append([token] + box)
                    else:
                        ret.append(token)
                    ii = min(ii, max_src_length - 1)
                    ret_index.append(ii)
            return ret, ret_index

        for text, layout in tqdm.tqdm(examples):
            if 'bleu' in text:
                bleu = text['bleu']
            else:
                bleu = 0

            if random.random() < src_shuffle_rate:
                # print('Random!!!')
                # DONE: the random src! here has bug! index also need shuffle
                src_text = text['src']
                src_layout = layout['src']
                tgt_index = text['tgt_index']

                src_text = src_text.split()
                source_length = len(src_text)
                shuffle_index = list(range(source_length))
                random.shuffle(shuffle_index)

                shuffle_text = ['' for _ in range(source_length)]
                shuffle_layout = ['' for _ in range(source_length)]
                for i, j in enumerate(shuffle_index):
                    # NOTE: map i-th token to j-th token
                    shuffle_text[j] = src_text[i]
                    shuffle_layout[j] = src_layout[i]

                shuffle_target_index = [shuffle_index[i] for i in tgt_index]

                text['src'] = ' '.join(shuffle_text)
                text['tgt_index'] = shuffle_target_index
                layout['src'] = shuffle_layout

            src_ids, src_index_split = tokenize_text_and_layout_src(text['src'], layout['src'],
                                                                    _layout_flag=layout_flag)
            tgt_ids, tgt_index = tokenize_text_and_layout_tgt(text['tgt'], layout['tgt'], text['tgt_index'],
                                                              src_index_split, _layout_flag=layout_flag)

            feature = {
                "source_ids": src_ids,
                "target_ids": tgt_ids,
                "target_index": tgt_index,
                'bleu': bleu
            }

            if file_info_flag:
                file_info = {'original_filename': text['original_filename'], 'filename': text['filename'], 'page_idx': text['page_idx']}
                feature['file_info'] = file_info

            features.append(feature)

        if shuffle:
            random.shuffle(features)

        if local_rank in [-1, 0] and cached_features_file is not None:
            if not os.path.exists(os.path.dirname(cached_features_file)):
                os.makedirs(os.path.dirname(cached_features_file))
            logger.info("Saving features into cached file %s", cached_features_file)
            torch.save(features, cached_features_file)

    # Make sure only the first process in distributed training process the dataset, and the others will use the cache
    if local_rank == 0:
        torch.distributed.barrier()

    return features


def convert_src_layout_inputs_to_tokens(inputs, converter, max_src_length, layout_flag=True):
    ret = []
    if not layout_flag:
        for line in inputs:
            ret.append(converter(line["source_ids"])[: max_src_length])
    else:
        for line in inputs:
            raw_text_ids = [x[0] for x in line['source_ids']]
            raw_text = converter(raw_text_ids)
            new_line = [[t] + x[1:] for t, x in zip(raw_text, line['source_ids'])][: max_src_length]
            ret.append(new_line)
    return ret


def convert_tgt_layout_inputs_to_tokens(inputs, converter, max_tgt_length, layout_flag=True):
    ret = []
    if not layout_flag:
        for line in inputs:
            ret.append(converter(line["target_ids"])[: max_tgt_length])
    else:
        for line in inputs:
            raw_text_ids = [x[0] for x in line['target_ids']]
            ret.append(converter(raw_text_ids)[: max_tgt_length])
    return ret


def get_tokens_from_src_and_index(src, index, modifier=None):
    result = []
    for i in index:
        i = modifier(i)
        i = min(i, len(src) - 1)
        if isinstance(src[i], list):
            result.append(src[i][0])
        else:
            result.append(src[i])
    return result


def get_layout_from_src_and_index(src, index, modifier=None):
    result = []
    s = set()
    for i in index:
        i = modifier(i)
        i = min(i, len(src) - 1)
        layout = src[i][1:]
        if repr(layout) not in s:
            result.append(layout)
            s.add(repr(layout))
    return result


def get_everything_from_src_and_index(src, index, modifier=None):
    result = []
    for i in index:
        i = modifier(i)
        i = min(i, len(src) - 1)
        result.append(src[i])
    return result