chore: import upstream snapshot with attribution

layoutreader/s2s_ft/s2s_loader.py

@@ -0,0 +1,165 @@
+import numpy as np
+
+from random import randint, shuffle, choice
+from random import random as rand
+import math
+import logging
+import torch
+import torch.utils.data
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_random_word(vocab_words):
+    i = randint(0, len(vocab_words)-1)
+    return vocab_words[i]
+
+
+def batch_list_to_batch_tensors(batch):
+    batch_tensors = []
+    for x in zip(*batch):
+        if x[0] is None:
+            batch_tensors.append(None)
+        elif isinstance(x[0], torch.Tensor):
+            batch_tensors.append(torch.stack(x))
+        else:
+            batch_tensors.append(torch.tensor(x, dtype=torch.long))
+    return batch_tensors
+
+
+def _get_word_split_index(tokens, st, end):
+    split_idx = []
+    i = st
+    while i < end:
+        if (not tokens[i].startswith('##')) or (i == st):
+            split_idx.append(i)
+        i += 1
+    split_idx.append(end)
+    return split_idx
+
+
+def _expand_whole_word(tokens, st, end):
+    new_st, new_end = st, end
+    while (new_st >= 0) and tokens[new_st].startswith('##'):
+        new_st -= 1
+    while (new_end < len(tokens)) and tokens[new_end].startswith('##'):
+        new_end += 1
+    return new_st, new_end
+
+
+class Pipeline():
+    """ Pre-process Pipeline Class : callable """
+
+    def __init__(self):
+        super().__init__()
+        self.skipgram_prb = None
+        self.skipgram_size = None
+        self.pre_whole_word = None
+        self.mask_whole_word = None
+        self.word_subsample_prb = None
+        self.sp_prob = None
+        self.pieces_dir = None
+        self.vocab_words = None
+        self.pieces_threshold = 10
+        self.call_count = 0
+        self.offline_mode = False
+        self.skipgram_size_geo_list = None
+        self.span_same_mask = False
+
+    def __call__(self, instance):
+        raise NotImplementedError
+
+
+class Preprocess4Seq2seqDecoder(Pipeline):
+    """ Pre-processing steps for pretraining transformer """
+
+    def __init__(self, vocab_words, indexer, max_len=512, max_tgt_length=128, 
+                 mode="s2s", pos_shift=False, source_type_id=0, target_type_id=1, 
+                 cls_token='[CLS]', sep_token='[SEP]', pad_token='[PAD]', layout_flag=False):
+        super().__init__()
+        self.max_len = max_len
+        self.vocab_words = vocab_words  # vocabulary (sub)words
+        self.indexer = indexer  # function from token to token index
+        self.max_len = max_len
+        self._tril_matrix = torch.tril(torch.ones((max_len, max_len), dtype=torch.long))
+        self.task_idx = 3   # relax projection layer for different tasks
+        assert mode in ("s2s", "l2r")
+        self.mode = mode
+        self.max_tgt_length = max_tgt_length
+        self.pos_shift = pos_shift
+
+        self.layout_flag = layout_flag
+
+        if layout_flag:
+            self.cls_token = [cls_token, 0, 0, 0, 0]
+            self.sep_token = [sep_token, 1000, 1000, 1000, 1000]
+            self.pad_token = [pad_token, 0, 0, 0, 0]
+        else:
+            self.cls_token = cls_token
+            self.sep_token = sep_token
+            self.pad_token = pad_token
+
+        self.source_type_id = source_type_id
+        self.target_type_id = target_type_id
+
+        self.cc = 0
+
+    def __call__(self, instance):
+        tokens_a, max_a_len = instance
+
+        # NOTE: must pad to the max src length
+        max_a_len = 511
+
+        padded_tokens_a = [self.cls_token] + tokens_a + [self.sep_token]
+        assert len(padded_tokens_a) <= max_a_len + 2
+        if max_a_len + 2 > len(padded_tokens_a):
+            padded_tokens_a += [self.pad_token] * \
+                (max_a_len + 2 - len(padded_tokens_a))
+        assert len(padded_tokens_a) == max_a_len + 2
+        max_len_in_batch = min(self.max_tgt_length + max_a_len + 2, self.max_len)
+        tokens = padded_tokens_a
+        segment_ids = [self.source_type_id] * (len(padded_tokens_a)) \
+                + [self.target_type_id] * (max_len_in_batch - len(padded_tokens_a))
+
+        mask_qkv = None
+
+        position_ids = []
+        for i in range(len(tokens_a) + 2):
+            position_ids.append(i)
+        for i in range(len(tokens_a) + 2, max_a_len + 2):
+            position_ids.append(0)
+        for i in range(max_a_len + 2, max_len_in_batch):
+            position_ids.append(i - (max_a_len + 2) + len(tokens_a) + 2)
+
+        # Token Indexing
+        if not self.layout_flag:
+            input_ids = self.indexer(tokens)
+        else:
+            raw_text = [x[0] for x in tokens]
+            raw_text_ids = self.indexer(raw_text)
+            input_ids = [[i] + x[1:] for i, x in zip(raw_text_ids, tokens)]
+
+        self.cc += 1
+        if self.cc < 5:
+            if not self.layout_flag:
+                logger.info("Input src = %s" % " ".join(self.vocab_words[tk_id] for tk_id in input_ids))
+            else:
+                logger.info("Input src = %s" % " ".join(self.vocab_words[tk_id[0]] for tk_id in input_ids))
+
+        # Zero Padding
+        input_mask = torch.zeros(
+            max_len_in_batch, max_len_in_batch, dtype=torch.long)
+        if self.mode == "s2s":
+            input_mask[:, :len(tokens_a)+2].fill_(1)
+        else:
+            st, end = 0, len(tokens_a) + 2
+            input_mask[st:end, st:end].copy_(
+                self._tril_matrix[:end, :end])
+            input_mask[end:, :len(tokens_a)+2].fill_(1)
+        second_st, second_end = len(padded_tokens_a), max_len_in_batch
+
+        input_mask[second_st:second_end, second_st:second_end].copy_(
+            self._tril_matrix[:second_end-second_st, :second_end-second_st])
+
+        return input_ids, segment_ids, position_ids, input_mask, mask_qkv, self.task_idx