pregenerate_training_data.py

# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors, The HuggingFace Inc. team and Huawei Noah's Ark Lab.
# Copyright (c) 2018, NVIDIA CORPORATION.  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 json
import collections
import logging
import os
import shelve
from argparse import ArgumentParser
import pickle
from tqdm import tqdm, trange
from tempfile import TemporaryDirectory
from multiprocessing import Pool

import numpy as np
from random import random, randrange, randint, shuffle, choice
import re
import string
from collections import Counter

from transformer.tokenization import BertTokenizer
# docs = None # global parameter

vocab_list = None
# This is used for running on Huawei Cloud.
oncloud = True
try:
    import moxing as mox
    cache_dir = "/cache"
    os.system("pip install pathlib")

except:
    oncloud = False # say random a filename instead
    cache_dir = "./"
from pathlib import Path

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt='%m/%d/%Y %H:%M:%S',
                    level=logging.INFO)
logger = logging.getLogger(__name__)


class DocumentDatabase:
    def __init__(self, reduce_memory=False):
        if reduce_memory:
            self.temp_dir = TemporaryDirectory()
            self.working_dir = Path(self.temp_dir.name)
            self.document_shelf_filepath = self.working_dir / 'shelf.db'
            self.document_shelf = shelve.open(os.path.join(cache_dir,'shelf.db'),
                                              flag='n', protocol=-1)
            self.documents = None
        else:
            self.documents = []
            self.document_shelf = None
            self.document_shelf_filepath = None
            self.temp_dir = None
        self.doc_lengths = []
        self.doc_cumsum = None
        self.cumsum_max = None
        self.reduce_memory = reduce_memory

    def add_document(self, document):
        if not document:
            return
        if self.reduce_memory:
            current_idx = len(self.doc_lengths)
            self.document_shelf[str(current_idx)] = document
        else:
            self.documents.append(document)
        self.doc_lengths.append(len(document))

    def _precalculate_doc_weights(self):
        self.doc_cumsum = np.cumsum(self.doc_lengths)
        self.cumsum_max = self.doc_cumsum[-1]

    def sample_doc(self, current_idx, sentence_weighted=True):
        # Uses the current iteration counter to ensure we don't sample the same doc twice
        if sentence_weighted:
            # With sentence weighting, we sample docs proportionally to their sentence length
            if self.doc_cumsum is None or len(self.doc_cumsum) != len(self.doc_lengths):
                self._precalculate_doc_weights()
            rand_start = self.doc_cumsum[current_idx]
            rand_end = rand_start + self.cumsum_max - self.doc_lengths[current_idx]
            sentence_index = randrange(rand_start, rand_end) % self.cumsum_max
            sampled_doc_index = np.searchsorted(self.doc_cumsum, sentence_index, side='right')
        else:
            # If we don't use sentence weighting, then every doc has an equal chance to be chosen
            sampled_doc_index = (current_idx + randrange(1, len(self.doc_lengths))) % len(self.doc_lengths)
        assert sampled_doc_index != current_idx
        if self.reduce_memory:
            return self.document_shelf[str(sampled_doc_index)]
        else:
            return self.documents[sampled_doc_index]

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

    def __getitem__(self, item):
        if self.reduce_memory:
            return self.document_shelf[str(item)]
        else:
            return self.documents[item]

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, traceback):
        if self.document_shelf is not None:
            self.document_shelf.close()
        if self.temp_dir is not None:
            self.temp_dir.cleanup()


def truncate_seq_pair(tokens_a, tokens_b, max_num_tokens):
    """Truncates a pair of sequences to a maximum sequence length. Lifted from Google's BERT repo."""
    while True:
        total_length = len(tokens_a) + len(tokens_b)
        if total_length <= max_num_tokens:
            break

        trunc_tokens = tokens_a if len(tokens_a) > len(tokens_b) else tokens_b
        assert len(trunc_tokens) >= 1

        # We want to sometimes truncate from the front and sometimes from the
        # back to add more randomness and avoid biases.
        if random() < 0.5:
            del trunc_tokens[0]
        else:
            trunc_tokens.pop()


MaskedLmInstance = collections.namedtuple("MaskedLmInstance",
                                          ["index", "label"])


def create_masked_lm_predictions(tokens, masked_lm_prob, max_predictions_per_seq, whole_word_mask, vocab_list):
    """Creates the predictions for the masked LM objective. This is mostly copied from the Google BERT repo, but
    with several refactors to clean it up and remove a lot of unnecessary variables."""
    cand_indices = []
    for (i, token) in enumerate(tokens):
        if token == "[CLS]" or token == "[SEP]":
            continue
        # Whole Word Masking means that if we mask all of the wordpieces
        # corresponding to an original word. When a word has been split into
        # WordPieces, the first token does not have any marker and any subsequence
        # tokens are prefixed with ##. So whenever we see the ## token, we
        # append it to the previous set of word indexes.
        #
        # Note that Whole Word Masking does *not* change the training code
        # at all -- we still predict each WordPiece independently, softmaxed
        # over the entire vocabulary.
        if (whole_word_mask and len(cand_indices) >= 1 and token.startswith("##")):
            cand_indices[-1].append(i)
        else:
            cand_indices.append([i])

    num_to_mask = min(max_predictions_per_seq,
                      max(1, int(round(len(tokens) * masked_lm_prob))))
    shuffle(cand_indices)
    masked_lms = []
    covered_indexes = set()
    for index_set in cand_indices:
        if len(masked_lms) >= num_to_mask:
            break
        # If adding a whole-word mask would exceed the maximum number of
        # predictions, then just skip this candidate.
        if len(masked_lms) + len(index_set) > num_to_mask:
            continue
        is_any_index_covered = False
        for index in index_set:
            if index in covered_indexes:
                is_any_index_covered = True
                break
        if is_any_index_covered:
            continue
        for index in index_set:
            covered_indexes.add(index)

            # 80% of the time, replace with [MASK]
            if random() < 0.8:
                masked_token = "[MASK]"
            else:
                # 10% of the time, keep original
                if random() < 0.5:
                    masked_token = tokens[index]
                # 10% of the time, replace with random word
                else:
                    masked_token = choice(vocab_list)
            masked_lms.append(MaskedLmInstance(index=index, label=tokens[index]))
            tokens[index] = masked_token

    assert len(masked_lms) <= num_to_mask
    masked_lms = sorted(masked_lms, key=lambda x: x.index)
    mask_indices = [p.index for p in masked_lms]
    masked_token_labels = [p.label for p in masked_lms]

    return tokens, mask_indices, masked_token_labels


def create_instances_from_document(
        doc_database, doc_idx, max_seq_length, short_seq_prob,
        masked_lm_prob, max_predictions_per_seq, whole_word_mask, vocab_list, bi_text=True):
    """This code is mostly a duplicate of the equivalent function from Google BERT's repo.
    However, we make some changes and improvements. Sampling is improved and no longer requires a loop in this function.
    Also, documents are sampled proportionally to the number of sentences they contain, which means each sentence
    (rather than each document) has an equal chance of being sampled as a false example for the NextSentence task."""
    document = doc_database[doc_idx]
    # Account for [CLS], [SEP], [SEP]
    max_num_tokens = max_seq_length - 3

    # We *usually* want to fill up the entire sequence since we are padding
    # to `max_seq_length` anyways, so short sequences are generally wasted
    # computation. However, we *sometimes*
    # (i.e., short_seq_prob == 0.1 == 10% of the time) want to use shorter
    # sequences to minimize the mismatch between pre-training and fine-tuning.
    # The `target_seq_length` is just a rough target however, whereas
    # `max_seq_length` is a hard limit.
    target_seq_length = max_num_tokens
    if random() < short_seq_prob:
        target_seq_length = randint(2, max_num_tokens)

    # We DON'T just concatenate all of the tokens from a document into a long
    # sequence and choose an arbitrary split point because this would make the
    # next sentence prediction task too easy. Instead, we split the input into
    # segments "A" and "B" based on the actual "sentences" provided by the user
    # input.
    instances = []
    current_chunk = []
    current_length = 0
    i = 0
    while i < len(document):
        segment = document[i]
        current_chunk.append(segment)
        current_length += len(segment)
        if i == len(document) - 1 or current_length >= target_seq_length:
            if current_chunk:
                # `a_end` is how many segments from `current_chunk` go into the `A`
                # (first) sentence.
                a_end = 1
                if len(current_chunk) >= 2:
                    a_end = randrange(1, len(current_chunk))

                tokens_a = []
                for j in range(a_end):
                    tokens_a.extend(current_chunk[j])

                tokens_b = []

                # Random next
                if bi_text and (len(current_chunk) == 1 or random() < 0.5):
                    is_random_next = True
                    target_b_length = target_seq_length - len(tokens_a)

                    # Sample a random document, with longer docs being sampled more frequently
                    random_document = doc_database.sample_doc(current_idx=doc_idx, sentence_weighted=True)

                    random_start = randrange(0, len(random_document))
                    for j in range(random_start, len(random_document)):
                        tokens_b.extend(random_document[j])
                        if len(tokens_b) >= target_b_length:
                            break
                    # We didn't actually use these segments so we "put them back" so
                    # they don't go to waste.
                    num_unused_segments = len(current_chunk) - a_end
                    i -= num_unused_segments
                # Actual next
                else:
                    is_random_next = False
                    for j in range(a_end, len(current_chunk)):
                        tokens_b.extend(current_chunk[j])

                if not tokens_a or len(tokens_a) == 0:
                    tokens_a = ["."]

                if not tokens_b or len(tokens_b) == 0:
                    tokens_b = ["."]

                assert len(tokens_a) >= 1
                assert len(tokens_b) >= 1

                truncate_seq_pair(tokens_a, tokens_b, max_num_tokens)

                tokens = ["[CLS]"] + tokens_a + ["[SEP]"] + tokens_b + ["[SEP]"]
                # The segment IDs are 0 for the [CLS] token, the A tokens and the first [SEP]
                # They are 1 for the B tokens and the final [SEP]
                segment_ids = [0 for _ in range(len(tokens_a) + 2)] + [1 for _ in range(len(tokens_b) + 1)]

                tokens, masked_lm_positions, masked_lm_labels = create_masked_lm_predictions(
                    tokens, masked_lm_prob, max_predictions_per_seq, whole_word_mask, vocab_list)

                instance = {
                    "tokens": tokens,
                    "segment_ids": segment_ids,
                    "is_random_next": is_random_next,
                    "masked_lm_positions": masked_lm_positions,
                    "masked_lm_labels": masked_lm_labels}

                instances.append(instance)
            current_chunk = []
            current_length = 0
        i += 1

    return instances


def create_training_file(_docs, vocab_list,  args, epoch_num, bi_text=True):
    # if pkl_filename is not None and os.path.exists(pkl_filename):
    #     print("loading docs from {}".format(pkl_filename))
    #     docs = pickle.load(open(pkl_filename,"rb"))
    epoch_filename = args.output_dir / "epoch_{}.json".format(epoch_num)
    print("generating {}".format(epoch_filename))
    # print("have {} docs in total".format(len(docs)))
    print("have {} words in total".format(len(vocab_list)))

    num_instances = 0
    with epoch_filename.open('w',encoding="utf-8") as epoch_file:


            docs = DocumentDatabase(reduce_memory=False)
            for i,doc in enumerate(_docs):
                docs.add_document(doc)
            print("processed {} docs".format(len(docs)))
            for doc_idx in trange(len(docs), desc="Document"):
                doc_instances = create_instances_from_document(
                    docs, doc_idx, max_seq_length=args.max_seq_len, short_seq_prob=args.short_seq_prob,
                    masked_lm_prob=args.masked_lm_prob, max_predictions_per_seq=args.max_predictions_per_seq,
                    whole_word_mask=args.do_whole_word_mask, vocab_list=vocab_list, bi_text=bi_text)
                doc_instances = [json.dumps(instance) for instance in doc_instances]
                for instance in doc_instances:
                    epoch_file.write(instance + '\n')
                    num_instances += 1
    metrics_filename = args.output_dir / "epoch_{}_metrics.json".format(epoch_num)
    with metrics_filename.open('w',encoding="utf-8") as metrics_file:
        metrics = {
            "num_training_examples": num_instances,
            "max_seq_len": args.max_seq_len
        }
        metrics_file.write(json.dumps(metrics))

    return epoch_filename, metrics_filename




def merge_documents(path):
    # dir,filename = os.path.split(path)
    final_filename = os.path.join(path,"corpus.txt")
    if os.path.exists(final_filename):
        return  final_filename
    with open(final_filename,"w",encoding="utf-8") as f:
        for subpath in os.listdir(path):
            subpath = os.path.join(path,subpath)
            if os.path.isdir(subpath):
                for document in os.listdir(subpath):
                    document = os.path.join(subpath,document)
                    print(document)
                    with open(document,encoding="utf-8") as ff:
                        for line in ff:
                            f.write(line)
    return final_filename


def get_filename(path):
    documents = []
    for subpath in os.listdir(path):
        subpath = os.path.join(path, subpath)
        if os.path.isdir(subpath):
            for document in os.listdir(subpath):
                document = os.path.join(subpath, document)
                documents.append(document)
    return documents

def is_valid(line):
    l = len(line)
    if l > 1000000 or l < 50:
        return False
    count = Counter(line)
    alpha_cnt = sum(count[ch] for ch in string.ascii_letters)
    if alpha_cnt < 50 or alpha_cnt / l < 0.7:
        return False
    if count['/'] / l > 0.05:  # filter hyperlinks
        return False
    if count['\\'] / l > 0.05:  # filter latex math equations
        return False
    if count['|'] / l > 0.05 or line[0] == '|':  # filter remaining tables
        return False
    return True

def pre_cleanup(line):
    line = line.replace('\t', ' ')  # replace tab with spaces
    line = ' '.join(line.strip().split())  # remove redundant spaces
    line = re.sub(r'\.{4,}', '...', line)  # remove extra dots
    line = line.replace('<<', '«').replace('>>', '»')  # group << together
    line = re.sub(' (,:\.\)\]»)', r'\1', line)  # remove space before >>
    line = re.sub('(\[\(«) ', r'\1', line)  # remove space after <<
    line = line.replace(',,', ',').replace(',.', '.')  # remove redundant punctuations
    line = re.sub(r' \*([^\s])', r' \1', line)  # remove redundant asterisks
    return ' '.join(line.strip().split())  # remove redundant spaces

def post_cleanup(line):
    line = re.sub(r'\\', ' ', line)  # remove all backslashes
    return ' '.join(line.strip().split())  # remove redundant space


def read_documents(filename, tokenizer):

    print("process {} with tokenizer {}".format(filename,id(tokenizer)))
    docs=[]
    with open(filename, encoding="utf-8") as f:
        doc = []
        for line in tqdm(f, desc="Loading Dataset", unit=" lines"):
            if is_valid(line):
                line = post_cleanup(line)
                if line == "":
                    docs.append(doc)
                    doc = []
                    if len(docs) % 100 == 0:
                        logger.info('loaded {} docs!'.format(len(docs)))
                else:
                    tokens = tokenizer.tokenize(line)
                    doc.append(tokens)
        if doc:
            docs.append(doc)  # If the last doc didn't end on a newline, make sure it still gets added

    print("finished processing {} with tokenizer {}".format(filename, id(tokenizer)))

    return docs
    

def split_file(filename="final_doc.txt", maxline = 10):
    file_count = 0
    line_count = 0
    with open(filename) as f, open(filename+"_split_{}".format(file_count),"w",encoding="utf-8") as newf:
        for line in f:
            line_count += 1
            if line_count > maxline : #and line == ""
                file_count = file_count + 1
                newf.close()
                newf = open(filename + "_split_{}".format(file_count),"w",encoding="utf-8")
                line_count = 0
            else:
                newf.write(line)
    return [filename+"_split_{}".format(i) for i in range(file_count+1)]



def parse_parameters():
    parser = ArgumentParser()
    # parser.add_argument('--train_corpus', type=Path, required=False, default="./bert_en_clean/")
    parser.add_argument('--train_corpus', type=Path, required=False,
                        default="./data/corpora")
    #C:/Users/w50011414/Desktop/generate_data/data
    #D:/codes/filesystem/danteng/bert/bert-data/

    parser.add_argument("--output_dir", type=Path, required=False, default="data/train_data")
    parser.add_argument("--bert_model", type=str, required=False, default="./models/bert-base-uncased")
    parser.add_argument("--do_lower_case", action="store_true", default=True)
    parser.add_argument("--do_whole_word_mask", action="store_true",
                        help="Whether to use whole word masking rather than per-WordPiece masking.")
    parser.add_argument("--reduce_memory", action="store_true", default=False,
                        help="Reduce memory usage for large datasets by keeping data on disc rather than in memory")

    parser.add_argument("--num_workers", type=int, default=2,
                        help="The number of workers to use to write the files")
    parser.add_argument("--epochs_to_generate", type=int, default=5,
                        help="Number of epochs of data to pregenerate")
    parser.add_argument("--max_seq_len", type=int, default=128)
    parser.add_argument("--short_seq_prob", type=float, default=0.1,
                        help="Probability of making a short sentence as a training example")
    parser.add_argument("--masked_lm_prob", type=float, default=0.15,
                        help="Probability of masking each token for the LM task")  # no [mask] symbol in corpus
    parser.add_argument("--max_predictions_per_seq", type=int, default=20,
                        help="Maximum number of tokens to mask in each sequence")
    # parser.add_argument('--data_url', type=str, default="")
    parser.add_argument('--one_seq', action='store_true')
    parser.add_argument("--multi_read", type=int, default=1,
                        help="Maximum number of tokens to mask in each sequence")

    # add 1. for huawei yun.
    parser.add_argument("--data_url", type=str, default="", help="s3 url")
    parser.add_argument("--train_url", type=str, default="", help="s3 url")
    parser.add_argument("--init_method", default='', type=str)
    parser.add_argument("--remote_path", default='', type=str) #benyou/en-bert-pytorch/max-seq-512-WIKI-BOOK-10-epochs


    args = parser.parse_args()
    return args


def main():
    args = parse_parameters()

    tokenizer = BertTokenizer.from_pretrained(args.bert_model, do_lower_case=args.do_lower_case)
    vocab_list = list(tokenizer.vocab.keys())
    doc_num = 0
    print(len(vocab_list))
    final_filename = merge_documents(args.train_corpus)
    with DocumentDatabase(reduce_memory=args.reduce_memory) as docs:
        # with args.train_corpus.open() as f:
        with open(final_filename,encoding="utf-8") as f:
            doc = []
            for line in tqdm(f, desc="Loading Dataset", unit=" lines"):
                line = line.strip()
                if is_valid(line):
                    line = pre_cleanup(line)
                    line = post_cleanup(line)
                if line == "":
                    docs.add_document(doc)
                    doc = []
                    doc_num += 1
                    if doc_num % 100 == 0:
                        logger.info('loaded {} docs!'.format(doc_num))
                else:
                    tokens = tokenizer.tokenize(line)
                    doc.append(tokens)
            if doc:
                docs.add_document(doc)  # If the last doc didn't end on a newline, make sure it still gets added
        if len(docs) <= 1:
            exit("ERROR: No document breaks were found in the input file! These are necessary to allow the script to "
                 "ensure that random NextSentences are not sampled from the same document. Please add blank lines to "
                 "indicate breaks between documents in your input file. If your dataset does not contain multiple "
                 "documents, blank lines can be inserted at any natural boundary, such as the ends of chapters, "
                 "sections or paragraphs.")

        args.output_dir.mkdir(exist_ok=True)

        if args.num_workers > 1:
            writer_workers = Pool(min(args.num_workers, args.epochs_to_generate))
            arguments = [(docs, vocab_list, args, idx) for idx in range(args.epochs_to_generate)]
            writer_workers.starmap(create_training_file, arguments)

        else:
            for epoch in trange(args.epochs_to_generate, desc="Epoch"):
                bi_text = True if not args.one_seq else False
                epoch_file, metric_file = create_training_file(docs, vocab_list, args, epoch, bi_text=bi_text)

                if oncloud:
                    # add 3. for huawei yun.
                    logging.info(mox.file.list_directory(str(args.output_dir), recursive=True))
                    mox.file.copy_parallel(str(args.output_dir), args.train_url)

                    # logging.info(mox.file.list_directory(str(args.output_dir), recursive=True))
                    # logging.info(mox.file.list_directory('.', recursive=True))
                    # mox.file.copy_parallel(str(args.output_dir), args.data_url)
                    # mox.file.copy_parallel('.', args.data_url)

                    os.remove(str(epoch_file))
                    os.remove(str(metric_file))



if __name__ == '__main__':
    # global docs
    main()