utils_unilm.py

from random import randint, shuffle, choice
from random import random as rand
import math
import numpy as np
import torch
import torch.utils.data


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:
            try:
                batch_tensors.append(torch.tensor(x, dtype=torch.long))
            except:
                batch_tensors.append(None)
    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.vocab_words = None
        self.call_count = 0
        self.offline_mode = False
        self.skipgram_size_geo_list = None
        self.span_same_mask = False

    def init_skipgram_size_geo_list(self, p):
        if p > 0:
            g_list = []
            t = p
            for _ in range(self.skipgram_size):
                g_list.append(t)
                t *= (1 - p)
            s = sum(g_list)
            self.skipgram_size_geo_list = [x / s for x in g_list]

    def __call__(self, instance):
        raise NotImplementedError

    # pre_whole_word: tokenize to words before masking
    # post whole word (--mask_whole_word): expand to words after masking
    def get_masked_pos(self,
                       tokens,
                       n_pred,
                       add_skipgram=False,
                       mask_segment=None,
                       protect_range=None):
        if self.pre_whole_word:
            pre_word_split = _get_word_split_index(tokens, 0, len(tokens))
        else:
            pre_word_split = list(range(0, len(tokens) + 1))

        span_list = list(zip(pre_word_split[:-1], pre_word_split[1:]))

        # candidate positions of masked tokens
        cand_pos = []
        special_pos = set()
        if mask_segment:
            for i, sp in enumerate(span_list):
                sp_st, sp_end = sp
                if (sp_end - sp_st == 1) and tokens[sp_st].endswith('SEP]'):
                    segment_index = i
                    break
        for i, sp in enumerate(span_list):
            sp_st, sp_end = sp
            if (sp_end - sp_st == 1) and (tokens[sp_st].endswith('CLS]')
                                          or tokens[sp_st].endswith('SEP]')):
                special_pos.add(i)
            else:
                if mask_segment:
                    if ((i < segment_index) and
                        ('a' in mask_segment)) or ((i > segment_index) and
                                                   ('b' in mask_segment)):
                        cand_pos.append(i)
                else:
                    cand_pos.append(i)
        shuffle(cand_pos)

        masked_pos = set()
        for i_span in cand_pos:
            if len(masked_pos) >= n_pred:
                break
            cand_st, cand_end = span_list[i_span]
            if len(masked_pos) + cand_end - cand_st > n_pred:
                continue
            if any(p in masked_pos for p in range(cand_st, cand_end)):
                continue

            n_span = 1
            rand_skipgram_size = 0
            # ngram
            if self.skipgram_size_geo_list:
                # sampling ngram size from geometric distribution
                rand_skipgram_size = np.random.choice(
                    len(self.skipgram_size_geo_list),
                    1,
                    p=self.skipgram_size_geo_list)[0] + 1
            else:
                if add_skipgram and (self.skipgram_prb > 0) and (
                        self.skipgram_size >= 2) and (rand() <
                                                      self.skipgram_prb):
                    rand_skipgram_size = min(randint(2, self.skipgram_size),
                                             len(span_list) - i_span)
            for n in range(2, rand_skipgram_size + 1):
                tail_st, tail_end = span_list[i_span + n - 1]
                if (tail_end - tail_st == 1) and (tail_st in special_pos):
                    break
                if len(masked_pos) + tail_end - cand_st > n_pred:
                    break
                n_span = n
            st_span, end_span = i_span, i_span + n_span

            if self.mask_whole_word:
                # pre_whole_word==False: position index of span_list is the same as tokens
                st_span, end_span = _expand_whole_word(tokens, st_span,
                                                       end_span)

            skip_pos = None

            for sp in range(st_span, end_span):
                for mp in range(span_list[sp][0], span_list[sp][1]):
                    if not (skip_pos and (mp in skip_pos)) and (
                            mp not in special_pos
                    ) and not (protect_range and
                               (protect_range[0] <= mp < protect_range[1])):
                        masked_pos.add(mp)

        if len(masked_pos) < n_pred:
            shuffle(cand_pos)
            for pos in cand_pos:
                if len(masked_pos) >= n_pred:
                    break
                if pos not in masked_pos:
                    masked_pos.add(pos)
        masked_pos = list(masked_pos)
        if len(masked_pos) > n_pred:
            # shuffle(masked_pos)
            masked_pos = masked_pos[:n_pred]
        return masked_pos

    def replace_masked_tokens(self, tokens, masked_pos):
        if self.span_same_mask:
            masked_pos = sorted(list(masked_pos))
        prev_pos, prev_rand = None, None
        for pos in masked_pos:
            if self.span_same_mask and (pos - 1 == prev_pos):
                t_rand = prev_rand
            else:
                t_rand = rand()
            if t_rand < 0.8:  # 80%
                tokens[pos] = '[MASK]'
            elif t_rand < 0.9:  # 10%
                tokens[pos] = get_random_word(self.vocab_words)
            prev_pos, prev_rand = pos, t_rand


# Input file format :
# 1. One sentence per line. These should ideally be actual sentences,
#    not entire paragraphs or arbitrary spans of text. (Because we use
#    the sentence boundaries for the "next sentence prediction" task).
# 2. Blank lines between documents. Document boundaries are needed
#    so that the "next sentence prediction" task doesn't span between documents.


def truncate_tokens_pair(tokens_a, tokens_b, max_len):
    if len(tokens_a) + len(tokens_b) > max_len - 3:
        while len(tokens_a) + len(tokens_b) > max_len - 3:
            if len(tokens_a) > len(tokens_b):
                tokens_a = tokens_a[:-1]
            else:
                tokens_b = tokens_b[:-1]
    return tokens_a, tokens_b


def truncate_tokens_signle(tokens_a, max_len):
    if len(tokens_a) > max_len - 2:
        tokens_a = tokens_a[:max_len - 2]
    return tokens_a


class Seq2SeqDataset(torch.utils.data.Dataset):
    """ Load sentence pair (sequential or random order) from corpus """
    def __init__(self,
                 file,
                 batch_size,
                 tokenizer,
                 max_len,
                 short_sampling_prob=0.1,
                 sent_reverse_order=False,
                 bi_uni_pipeline=[]):
        super().__init__()
        self.tokenizer = tokenizer  # tokenize function
        self.max_len = max_len  # maximum length of tokens
        self.short_sampling_prob = short_sampling_prob
        self.bi_uni_pipeline = bi_uni_pipeline
        assert len(self.bi_uni_pipeline) == 4
        self.batch_size = batch_size
        self.sent_reverse_order = sent_reverse_order

        # read the file into memory
        self.ex_list = []
        self.idx = 0
        # with open(file, "r", encoding='utf-8') as f:
        #     for i, line in enumerate(f):
        #         sample = eval(line.strip())
        #         src_tk = tokenizer.tokenize(sample["src_text"])
        #         tgt_tk = tokenizer.tokenize(sample["tgt_text"])
        #         assert len(src_tk) > 0
        #         assert len(tgt_tk) > 0
        #         self.ex_list.append((src_tk, tgt_tk))
        self.read_data(file, tokenizer)
        # file_data = open(file, "r", encoding='utf-8')
        #
        # threads = min(8, cpu_count())
        # with Pool(threads) as p:
        #     annotate_ = partial(
        #         self.read_data,
        #         tokenizer=self.tokenizer)
        #     self.ex_list = list(
        #         tqdm(
        #             p.imap(annotate_, file_data.readlines(), chunksize=32),
        #             total=len(file_data.readlines()),
        #             desc="convert squad examples to features",
        #         )
        #     )
        # fin = open("look_new.json", "w",encoding="utf-8")
        # for jj, m in enumerate(self.ex_list):
        #     fin.write(str(jj)+"\t"+str(m)+"\n")
        # print('Load {0} documents'.format(len(self.ex_list)))
        # exit()

    def gene_query(self, title):
        a = rand()
        query = None
        if a < 0.4 or len(title) >= 45:
            title_len = len(title)
            query_len = randint(int(title_len * .2), int(title_len * .4))
            start = randint(0, title_len - query_len)
            query = title[start:start + query_len]
        elif 0.4 <= a < 0.6:
            query = title[:4] + title[-4:]
        elif 0.6 <= a < 0.8:
            query = title[:8]
        else:
            query = title[-8:]
        return query

    def read_data(self, file, tokenizer):
        with open(file, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            for line in lines:
                line = line.strip().split("\t")
                self.ex_list.append((line[1], self.gene_query(line[1])))
                # src_tk = tokenizer.tokenize(sample["src_text"])
                # tgt_tk = tokenizer.tokenize(sample["tgt_text"])
                # src_tk = sample["src_text"]
                # tgt_tk = sample["tgt_text"]
                # return (src_tk, tgt_tk)

    def __len__(self):
        return len(self.ex_list)

    def __getitem__(self, idx):
        instance = self.ex_list[idx]
        # proc = choice(self.bi_uni_pipeline)
        new_instance = ()
        proc = None
        if self.idx <= 1:
            proc = self.bi_uni_pipeline[0]
        elif 2 <= self.idx <= 3:
            proc = self.bi_uni_pipeline[1]
        elif self.idx == 4:
            proc = self.bi_uni_pipeline[2]
        elif self.idx == 5:
            self.idx = -1
            proc = self.bi_uni_pipeline[3]
        self.idx += 1
        new_instance += proc(instance)
        return new_instance

    def __iter__(self):  # iterator to load data
        for __ in range(math.ceil(len(self.ex_list) / float(self.batch_size))):
            batch = []
            for __ in range(self.batch_size):
                idx = randint(0, len(self.ex_list) - 1)
                batch.append(self.__getitem__(idx))
            # To Tensor
            yield batch_list_to_batch_tensors(batch)


class Preprocess4Seq2seq(Pipeline):
    """ Pre-processing steps for pretraining transformer """
    def __init__(self,
                 max_pred,
                 mask_prob,
                 vocab_words,
                 indexer,
                 max_len=512,
                 skipgram_prb=0,
                 skipgram_size=0,
                 mask_whole_word=False,
                 mask_source_words=True,
                 tokenizer=None):
        super().__init__()
        self.max_len = max_len
        self.max_pred = max_pred  # max tokens of prediction
        self.mask_prob = mask_prob  # masking probability
        self.vocab_words = vocab_words  # vocabulary (sub)words
        self.indexer = indexer  # function from token to token index
        self._tril_matrix = torch.tril(
            torch.ones((max_len, max_len), dtype=torch.long))
        self.skipgram_prb = skipgram_prb
        self.skipgram_size = skipgram_size
        self.mask_whole_word = mask_whole_word
        self.mask_source_words = mask_source_words
        self.tokenizer = tokenizer

    def __call__(self, instance):
        next_sentence_label = None
        tokens_a, tokens_b = instance[:2]
        tokens_a = self.tokenizer.tokenize(tokens_a)
        tokens_b = self.tokenizer.tokenize(tokens_b)
        # -3  for special tokens [CLS], [SEP], [SEP]
        tokens_a, tokens_b = truncate_tokens_pair(tokens_a, tokens_b,
                                                  self.max_len)
        # Add Special Tokens
        tokens = ['[CLS]'] + tokens_a + ['[SEP]'] + tokens_b + ['[SEP]']
        segment_ids = [4] * (len(tokens_a) + 2) + [5] * (len(tokens_b) + 1)
        # For masked Language Models
        # the number of prediction is sometimes less than max_pred when sequence is short
        effective_length = len(tokens_b)
        if self.mask_source_words:
            effective_length += len(tokens_a)
        n_pred = min(self.max_pred,
                     max(1, int(round(effective_length * self.mask_prob))))
        # candidate positions of masked tokens
        cand_pos = []
        special_pos = set()
        for i, tk in enumerate(tokens):
            # only mask tokens_b (target sequence)
            # we will mask [SEP] as an ending symbol
            if (i >= len(tokens_a) + 2) and (tk != '[CLS]'):
                cand_pos.append(i)
            elif self.mask_source_words and (i < len(tokens_a) + 2) and (
                    tk != '[CLS]') and (not tk.startswith('[SEP')):
                cand_pos.append(i)
            else:
                special_pos.add(i)
        shuffle(cand_pos)

        masked_pos = set()
        max_cand_pos = max(cand_pos)
        for pos in cand_pos:
            if len(masked_pos) >= n_pred:
                break
            if pos in masked_pos:
                continue

            def _expand_whole_word(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

            if (self.skipgram_prb > 0) and (self.skipgram_size >= 2) and (
                    rand() < self.skipgram_prb):
                # ngram
                cur_skipgram_size = randint(2, self.skipgram_size)
                if self.mask_whole_word:
                    st_pos, end_pos = _expand_whole_word(
                        pos, pos + cur_skipgram_size)
                else:
                    st_pos, end_pos = pos, pos + cur_skipgram_size
            else:
                # directly mask
                if self.mask_whole_word:
                    st_pos, end_pos = _expand_whole_word(pos, pos + 1)
                else:
                    st_pos, end_pos = pos, pos + 1

            for mp in range(st_pos, end_pos):
                if (0 < mp <= max_cand_pos) and (mp not in special_pos):
                    masked_pos.add(mp)
                else:
                    break

        masked_pos = list(masked_pos)
        if len(masked_pos) > n_pred:
            shuffle(masked_pos)
            masked_pos = masked_pos[:n_pred]

        masked_tokens = [tokens[pos] for pos in masked_pos]
        for pos in masked_pos:
            if rand() < 0.8:  # 80%
                tokens[pos] = '[MASK]'
            elif rand() < 0.5:  # 10%
                tokens[pos] = get_random_word(self.vocab_words)
        # when n_pred < max_pred, we only calculate loss within n_pred
        masked_weights = [1] * len(masked_tokens)

        # Token Indexing
        masked_ids = self.indexer(masked_tokens)
        # Token Indexing
        input_ids = self.indexer(tokens)

        # Zero Padding
        n_pad = self.max_len - len(input_ids)
        input_ids.extend([0] * n_pad)
        segment_ids.extend([0] * n_pad)

        input_mask = torch.zeros(self.max_len, self.max_len, dtype=torch.long)
        input_mask[:, :len(tokens_a) + 2].fill_(1)
        second_st, second_end = len(tokens_a) + 2, len(tokens_a) + len(
            tokens_b) + 3
        input_mask[second_st:second_end, second_st:second_end].copy_(
            self._tril_matrix[:second_end - second_st, :second_end -
                              second_st])

        # Zero Padding for masked target
        if self.max_pred > n_pred:
            n_pad = self.max_pred - n_pred
            if masked_ids is not None:
                masked_ids.extend([0] * n_pad)
            if masked_pos is not None:
                masked_pos.extend([0] * n_pad)
            if masked_weights is not None:
                masked_weights.extend([0] * n_pad)

        # print("seq2seq", masked_ids)
        return (input_ids, segment_ids, input_mask, masked_ids, masked_pos,
                masked_weights, next_sentence_label)


class Preprocess4BiLM(Pipeline):
    """ Pre-processing steps for pretraining transformer """
    def __init__(self,
                 max_pred,
                 mask_prob,
                 vocab_words,
                 indexer,
                 max_len=512,
                 skipgram_prb=0,
                 skipgram_size=0,
                 mask_whole_word=False,
                 mask_source_words=True,
                 tokenizer=None):
        super().__init__()
        self.max_len = max_len
        self.max_pred = max_pred  # max tokens of prediction
        self.mask_prob = mask_prob  # masking probability
        self.vocab_words = vocab_words  # vocabulary (sub)words
        self.indexer = indexer  # function from token to token index
        self._tril_matrix = torch.tril(
            torch.ones((max_len, max_len), dtype=torch.long))
        self.skipgram_prb = skipgram_prb
        self.skipgram_size = skipgram_size
        self.mask_whole_word = mask_whole_word
        self.mask_source_words = mask_source_words
        self.tokenizer = tokenizer

    def __call__(self, instance):
        tokens_a, tokens_b = instance[:2]
        if rand() <= 0.5:
            next_sentence_label = 1.0
        else:
            tokens_a, tokens_b = tokens_b, tokens_a
            next_sentence_label = 0.0

        tokens_a = self.tokenizer.tokenize(tokens_a)
        tokens_b = self.tokenizer.tokenize(tokens_b)
        # -3  for special tokens [CLS], [SEP], [SEP]
        tokens_a, tokens_b = truncate_tokens_pair(tokens_a, tokens_b,
                                                  self.max_len)
        # Add Special Tokens
        tokens = ['[CLS]'] + tokens_a + ['[SEP]'] + tokens_b + ['[SEP]']
        segment_ids = [0] * (len(tokens_a) + 2) + [1] * (len(tokens_b) + 1)

        # For masked Language Models
        # the number of prediction is sometimes less than max_pred when sequence is short
        effective_length = len(tokens_b)
        if self.mask_source_words:
            effective_length += len(tokens_a)
        n_pred = min(self.max_pred,
                     max(1, int(round(effective_length * self.mask_prob))))
        # candidate positions of masked tokens
        cand_pos = []
        special_pos = set()
        for i, tk in enumerate(tokens):
            # only mask tokens_b (target sequence)
            # we will mask [SEP] as an ending symbol
            if (i >= len(tokens_a) + 2) and (tk != '[CLS]'):
                cand_pos.append(i)
            elif self.mask_source_words and (i < len(tokens_a) + 2) and (
                    tk != '[CLS]') and (not tk.startswith('[SEP')):
                cand_pos.append(i)
            else:
                special_pos.add(i)
        shuffle(cand_pos)

        masked_pos = set()
        max_cand_pos = max(cand_pos)
        for pos in cand_pos:
            if len(masked_pos) >= n_pred:
                break
            if pos in masked_pos:
                continue

            def _expand_whole_word(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

            if (self.skipgram_prb > 0) and (self.skipgram_size >= 2) and (
                    rand() < self.skipgram_prb):
                # ngram
                cur_skipgram_size = randint(2, self.skipgram_size)
                if self.mask_whole_word:
                    st_pos, end_pos = _expand_whole_word(
                        pos, pos + cur_skipgram_size)
                else:
                    st_pos, end_pos = pos, pos + cur_skipgram_size
            else:
                # directly mask
                if self.mask_whole_word:
                    st_pos, end_pos = _expand_whole_word(pos, pos + 1)
                else:
                    st_pos, end_pos = pos, pos + 1

            for mp in range(st_pos, end_pos):
                if (0 < mp <= max_cand_pos) and (mp not in special_pos):
                    masked_pos.add(mp)
                else:
                    break

        masked_pos = list(masked_pos)
        if len(masked_pos) > n_pred:
            shuffle(masked_pos)
            masked_pos = masked_pos[:n_pred]

        masked_tokens = [tokens[pos] for pos in masked_pos]
        for pos in masked_pos:
            if rand() < 0.8:  # 80%
                tokens[pos] = '[MASK]'
            elif rand() < 0.5:  # 10%
                tokens[pos] = get_random_word(self.vocab_words)
        # when n_pred < max_pred, we only calculate loss within n_pred
        masked_weights = [1] * len(masked_tokens)

        # Token Indexing
        masked_ids = self.indexer(masked_tokens)
        # Token Indexing
        input_ids = self.indexer(tokens)

        # Zero Padding
        n_pad = self.max_len - len(input_ids)
        input_ids.extend([0] * n_pad)
        segment_ids.extend([0] * n_pad)

        input_mask = torch.ones(self.max_len, self.max_len, dtype=torch.long)
        # input_mask[:, :len(tokens_a)+2].fill_(1)
        # second_st, second_end = len(
        #     tokens_a)+2, len(tokens_a)+len(tokens_b)+3
        # input_mask[second_st:second_end, second_st:second_end].copy_(
        #     self._tril_matrix[:second_end-second_st, :second_end-second_st])

        # Zero Padding for masked target
        if self.max_pred > n_pred:
            n_pad = self.max_pred - n_pred
            if masked_ids is not None:
                masked_ids.extend([0] * n_pad)
            if masked_pos is not None:
                masked_pos.extend([0] * n_pad)
            if masked_weights is not None:
                masked_weights.extend([0] * n_pad)

        # print("bilm", masked_ids)
        return (input_ids, segment_ids, input_mask, masked_ids, masked_pos,
                masked_weights, next_sentence_label)


class Preprocess4RightLM(Pipeline):
    """ Pre-processing steps for pretraining transformer """
    def __init__(self,
                 max_pred,
                 mask_prob,
                 vocab_words,
                 indexer,
                 max_len=512,
                 skipgram_prb=0,
                 skipgram_size=0,
                 mask_whole_word=False,
                 mask_source_words=True,
                 tokenizer=None):
        super().__init__()
        self.max_len = max_len
        self.max_pred = max_pred  # max tokens of prediction
        self.mask_prob = mask_prob  # masking probability
        self.vocab_words = vocab_words  # vocabulary (sub)words
        self.indexer = indexer  # function from token to token index
        self._tril_matrix = torch.tril(
            torch.ones((max_len, max_len), dtype=torch.long))
        self.skipgram_prb = skipgram_prb
        self.skipgram_size = skipgram_size
        self.mask_whole_word = mask_whole_word
        self.mask_source_words = mask_source_words
        self.tokenizer = tokenizer

    def __call__(self, instance):
        next_sentence_label = None
        tokens_a, _ = instance[:2]
        tokens_a = self.tokenizer.tokenize(tokens_a)
        tokens_a = truncate_tokens_signle(tokens_a, self.max_len)
        # Add Special Tokens
        tokens = ['[CLS]'] + tokens_a + ['[SEP]']
        segment_ids = [2] * (len(tokens_a) + 2)

        # For masked Language Models
        # the number of prediction is sometimes less than max_pred when sequence is short
        effective_length = 0
        if self.mask_source_words:
            effective_length += len(tokens_a)
        n_pred = min(self.max_pred,
                     max(1, int(round(effective_length * self.mask_prob))))
        # candidate positions of masked tokens
        cand_pos = []
        special_pos = set()
        for i, tk in enumerate(tokens):
            # only mask tokens_b (target sequence)
            # we will mask [SEP] as an ending symbol
            # if (i >= len(tokens_a)+2) and (tk != '[CLS]'):
            #     cand_pos.append(i)
            if (tk != '[CLS]') and (tk != '[SEP]'):
                cand_pos.append(i)
            else:
                special_pos.add(i)

        shuffle(cand_pos)

        masked_pos = set()

        try:
            max_cand_pos = max(cand_pos)
        except:
            max_cand_pos = 0

        for pos in cand_pos:
            if len(masked_pos) >= n_pred:
                break
            if pos in masked_pos:
                continue

            def _expand_whole_word(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

            if (self.skipgram_prb > 0) and (self.skipgram_size >= 2) and (
                    rand() < self.skipgram_prb):
                # ngram
                cur_skipgram_size = randint(2, self.skipgram_size)
                if self.mask_whole_word:
                    st_pos, end_pos = _expand_whole_word(
                        pos, pos + cur_skipgram_size)
                else:
                    st_pos, end_pos = pos, pos + cur_skipgram_size
            else:
                # directly mask
                if self.mask_whole_word:
                    st_pos, end_pos = _expand_whole_word(pos, pos + 1)
                else:
                    st_pos, end_pos = pos, pos + 1

            for mp in range(st_pos, end_pos):
                if (0 < mp <= max_cand_pos) and (mp not in special_pos):
                    masked_pos.add(mp)
                else:
                    break

        masked_pos = list(masked_pos)
        if len(masked_pos) > n_pred:
            shuffle(masked_pos)
            masked_pos = masked_pos[:n_pred]

        masked_tokens = [tokens[pos] for pos in masked_pos]
        for pos in masked_pos:
            if rand() < 0.8:  # 80%
                tokens[pos] = '[MASK]'
            elif rand() < 0.5:  # 10%
                tokens[pos] = get_random_word(self.vocab_words)
        # when n_pred < max_pred, we only calculate loss within n_pred
        masked_weights = [1] * len(masked_tokens)

        # Token Indexing
        masked_ids = self.indexer(masked_tokens)
        # Token Indexing
        input_ids = self.indexer(tokens)

        # Zero Padding
        n_pad = self.max_len - len(input_ids)
        input_ids.extend([0] * n_pad)
        segment_ids.extend([0] * n_pad)

        input_mask = torch.ones(self.max_len, self.max_len, dtype=torch.long)
        # input_mask[:, :len(tokens_a)+2].fill_(1)
        second_st, second_end = 0, len(tokens_a) + 2
        input_mask[second_st:second_end, second_st:second_end].copy_(
            self._tril_matrix[:second_end - second_st, :second_end -
                              second_st])

        # Zero Padding for masked target
        if self.max_pred > n_pred:
            n_pad = self.max_pred - n_pred
            if masked_ids is not None:
                masked_ids.extend([0] * n_pad)
            if masked_pos is not None:
                masked_pos.extend([0] * n_pad)
            if masked_weights is not None:
                masked_weights.extend([0] * n_pad)

        # print("rm", masked_ids)
        return (input_ids, segment_ids, input_mask, masked_ids, masked_pos,
                masked_weights, next_sentence_label)


class Preprocess4LeftLM(Pipeline):
    """ Pre-processing steps for pretraining transformer """
    def __init__(self,
                 max_pred,
                 mask_prob,
                 vocab_words,
                 indexer,
                 max_len=512,
                 skipgram_prb=0,
                 skipgram_size=0,
                 mask_whole_word=False,
                 mask_source_words=True,
                 tokenizer=None):
        super().__init__()
        self.max_len = max_len
        self.max_pred = max_pred  # max tokens of prediction
        self.mask_prob = mask_prob  # masking probability
        self.vocab_words = vocab_words  # vocabulary (sub)words
        self.indexer = indexer  # function from token to token index
        self._tril_matrix = torch.triu(
            torch.ones((max_len, max_len), dtype=torch.long))
        self.skipgram_prb = skipgram_prb
        self.skipgram_size = skipgram_size
        self.mask_whole_word = mask_whole_word
        self.mask_source_words = mask_source_words
        self.tokenizer = tokenizer

    def __call__(self, instance):
        next_sentence_label = None
        tokens_a, _ = instance[:2]

        tokens_a = self.tokenizer.tokenize(tokens_a)
        tokens_a = truncate_tokens_signle(tokens_a, self.max_len)
        # Add Special Tokens
        tokens = ['[CLS]'] + tokens_a + ['[SEP]']

        segment_ids = [3] * (len(tokens_a) + 2)

        # For masked Language Models
        # the number of prediction is sometimes less than max_pred when sequence is short
        effective_length = 0
        if self.mask_source_words:
            effective_length += len(tokens_a)
        n_pred = min(self.max_pred,
                     max(1, int(round(effective_length * self.mask_prob))))
        # candidate positions of masked tokens
        cand_pos = []
        special_pos = set()
        for i, tk in enumerate(tokens):
            # only mask tokens_b (target sequence)
            # we will mask [SEP] as an ending symbol
            # if (i >= len(tokens_a)+2) and (tk != '[CLS]'):
            #     cand_pos.append(i)
            if (tk != '[CLS]') and (tk != '[SEP]'):
                cand_pos.append(i)
            else:
                special_pos.add(i)

        shuffle(cand_pos)

        masked_pos = set()

        try:
            max_cand_pos = max(cand_pos)
        except:
            max_cand_pos = 0

        for pos in cand_pos:
            if len(masked_pos) >= n_pred:
                break
            if pos in masked_pos:
                continue

            def _expand_whole_word(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

            if (self.skipgram_prb > 0) and (self.skipgram_size >= 2) and (
                    rand() < self.skipgram_prb):
                # ngram
                cur_skipgram_size = randint(2, self.skipgram_size)
                if self.mask_whole_word:
                    st_pos, end_pos = _expand_whole_word(
                        pos, pos + cur_skipgram_size)
                else:
                    st_pos, end_pos = pos, pos + cur_skipgram_size
            else:
                # directly mask
                if self.mask_whole_word:
                    st_pos, end_pos = _expand_whole_word(pos, pos + 1)
                else:
                    st_pos, end_pos = pos, pos + 1

            for mp in range(st_pos, end_pos):
                if (0 < mp <= max_cand_pos) and (mp not in special_pos):
                    masked_pos.add(mp)
                else:
                    break

        masked_pos = list(masked_pos)
        if len(masked_pos) > n_pred:
            shuffle(masked_pos)
            masked_pos = masked_pos[:n_pred]

        masked_tokens = [tokens[pos] for pos in masked_pos]
        for pos in masked_pos:
            if rand() < 0.8:  # 80%
                tokens[pos] = '[MASK]'
            elif rand() < 0.5:  # 10%
                tokens[pos] = get_random_word(self.vocab_words)
        # when n_pred < max_pred, we only calculate loss within n_pred
        masked_weights = [1] * len(masked_tokens)

        # Token Indexing
        masked_ids = self.indexer(masked_tokens)
        # Token Indexing
        input_ids = self.indexer(tokens)

        # Zero Padding
        n_pad = self.max_len - len(input_ids)
        input_ids.extend([0] * n_pad)
        segment_ids.extend([0] * n_pad)

        input_mask = torch.ones(self.max_len, self.max_len, dtype=torch.long)
        # input_mask[:, :len(tokens_a)+2].fill_(1)
        second_st, second_end = 0, len(tokens_a) + 2
        input_mask[second_st:second_end, second_st:second_end].copy_(
            self._tril_matrix[:second_end - second_st, :second_end -
                              second_st])

        # Zero Padding for masked target
        if self.max_pred > n_pred:
            n_pad = self.max_pred - n_pred
            if masked_ids is not None:
                masked_ids.extend([0] * n_pad)
            if masked_pos is not None:
                masked_pos.extend([0] * n_pad)
            if masked_weights is not None:
                masked_weights.extend([0] * n_pad)

        # print("lm", masked_ids)
        return (input_ids, segment_ids, input_mask, masked_ids, masked_pos,
                masked_weights, next_sentence_label)


class Preprocess4Seq2seqDecode(Pipeline):
    """ Pre-processing steps for pretraining transformer """
    def __init__(self, vocab_words, indexer, max_len=512, max_tgt_length=128):
        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.max_tgt_length = max_tgt_length

    def __call__(self, instance):
        tokens_a, max_a_len = instance

        # Add Special Tokens
        padded_tokens_a = ['[CLS]'] + tokens_a + ['[SEP]']
        assert len(padded_tokens_a) <= max_a_len + 2
        if max_a_len + 2 > len(padded_tokens_a):
            padded_tokens_a += ['[PAD]'] * \
                (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 = [4] * (len(padded_tokens_a)) + [5] * (
            max_len_in_batch - len(padded_tokens_a))

        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
        input_ids = self.indexer(tokens)

        # Zero Padding
        input_mask = torch.zeros(max_len_in_batch,
                                 max_len_in_batch,
                                 dtype=torch.long)
        input_mask[:, :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)