preproc.py

import random
from tqdm import tqdm
import spacy
import ujson as json
from collections import Counter
import numpy as np
from codecs import open

'''
The content of this file is mostly copied from https://github.com/HKUST-KnowComp/R-Net/blob/master/prepro.py
'''

nlp = spacy.blank("en")


def word_tokenize(sent):
    doc = nlp(sent)
    return [token.text for token in doc]


def convert_idx(text, tokens):
    current = 0
    spans = []
    for token in tokens:
        current = text.find(token, current)
        if current < 0:
            print("Token {} cannot be found".format(token))
            raise Exception()
        spans.append((current, current + len(token)))
        current += len(token)
    return spans


def process_file(filename, data_type, word_counter, char_counter):
    print("Generating {} examples...".format(data_type))
    examples = []
    eval_examples = {}
    total = 0
    with open(filename, "r") as fh:
        source = json.load(fh)
        for article in tqdm(source["data"]):
            for para in article["paragraphs"]:
                context = para["context"].replace(
                    "''", '" ').replace("``", '" ')
                context_tokens = word_tokenize(context)
                context_chars = [list(token) for token in context_tokens]
                spans = convert_idx(context, context_tokens)
                for token in context_tokens:
                    word_counter[token] += len(para["qas"])
                    for char in token:
                        char_counter[char] += len(para["qas"])
                for qa in para["qas"]:
                    ques = qa["question"].replace(
                        "''", '" ').replace("``", '" ')
                    ques_tokens = word_tokenize(ques)
                    ques_chars = [list(token) for token in ques_tokens]
                    for token in ques_tokens:
                        word_counter[token] += 1
                        for char in token:
                            char_counter[char] += 1
                    y1s, y2s = [], []
                    answer_texts = []
                    for answer in qa["answers"]:
                        answer_text = answer["text"]
                        answer_start = answer['answer_start']
                        answer_end = answer_start + len(answer_text)
                        answer_texts.append(answer_text)
                        answer_span = []
                        for idx, span in enumerate(spans):
                            if not (answer_end <= span[0] or answer_start >= span[1]):
                                answer_span.append(idx)
                        y1, y2 = answer_span[0], answer_span[-1]
                        y1s.append(y1)
                        y2s.append(y2)
                    example = {"context_tokens": context_tokens, "context_chars": context_chars,
                               "ques_tokens": ques_tokens,
                               "ques_chars": ques_chars, "y1s": y1s[-1], "y2s": y2s[-1], "id": total}
                    examples.append(example)
                    eval_examples[str(total)] = {
                        "context": context, "spans": spans, "answers": answer_texts, "uuid": qa["id"]}
                    total += 1

        print("{} questions in total".format(len(examples)))

    return examples, eval_examples

def get_embedding(counter, data_type, limit=-1, emb_file=None, vec_size=None):
    print("Generating {} embedding...".format(data_type))
    embedding_dict = {}
    filtered_elements = [k for k, v in counter.items() if v < limit]
    print('{} words have been filtered'.format(len(filtered_elements)))
    if emb_file is not None:
        assert vec_size is not None
        with open(emb_file, "r", encoding="utf-8") as fh:
            for line in tqdm(fh):
                array = line.split()
                word = "".join(array[0:-vec_size])
                vector = list(map(float, array[-vec_size:]))
                if word in counter and counter[word] > limit:
                    embedding_dict[word] = vector
        print("{} / {} tokens have corresponding {} embedding vector".format(
            len(embedding_dict), len(counter)-len(filtered_elements), data_type))
    else:
        assert vec_size is not None
        for token in filtered_elements:
            embedding_dict[token] = [np.random.normal(
                scale=0.1) for _ in range(vec_size)]
        print("{} tokens have corresponding embedding vector".format(
            len(filtered_elements)))

    NULL = "--NULL--"
    OOV = "--OOV--"
    token2idx_dict = {token: idx for idx,
                                     token in enumerate(embedding_dict.keys(), 2)}
    token2idx_dict[NULL] = 0
    token2idx_dict[OOV] = 1
    embedding_dict[NULL] = [0. for _ in range(vec_size)]
    embedding_dict[OOV] = [0. for _ in range(vec_size)]
    idx2emb_dict = {idx: embedding_dict[token]
                    for token, idx in token2idx_dict.items()}
    emb_mat = [idx2emb_dict[idx] for idx in range(len(idx2emb_dict))]
    return emb_mat, token2idx_dict


# def build_features(config, examples, data_type, out_file, word2idx_dict, char2idx_dict, is_test=False):
#     para_limit = config.test_para_limit if is_test else config.para_limit
#     ques_limit = config.test_ques_limit if is_test else config.ques_limit
#     ans_limit = 100 if is_test else config.ans_limit
#     word_len = config.word_len
#
#     def filter_func(example, is_test=False):
#         return len(example["context_tokens"]) > para_limit or \
#                len(example["ques_tokens"]) > ques_limit or \
#                (example["y2s"][0] - example["y1s"][0]) > ans_limit
#
#     print("Processing {} examples...".format(data_type))
#     total = 0
#     total_ = 0
#     meta = {}
#     N = len(examples)
#     context_idxs = np.zeros([N, para_limit], dtype=np.int32)
#     context_char_idxs = np.zeros([N, para_limit, word_len], dtype=np.int32)
#     ques_idxs = np.zeros([N, ques_limit], dtype=np.int32)
#     ques_char_idxs = np.zeros([N, ques_limit, word_len], dtype=np.int32)
#     y1s = np.zeros([N], dtype=np.int32)
#     y2s = np.zeros([N], dtype=np.int32)
#     ids = np.zeros([N], dtype=np.int64)
#     for n, example in tqdm(enumerate(examples)):
#         total_ += 1
#
#         if filter_func(example, is_test):
#             continue
#
#         total += 1
#
#         def _get_word(word):
#             for each in (word, word.lower(), word.capitalize(), word.upper()):
#                 if each in word2idx_dict:
#                     return word2idx_dict[each]
#             return 1
#
#         def _get_char(char):
#             if char in char2idx_dict:
#                 return char2idx_dict[char]
#             return 1
#
#         for i, token in enumerate(example["context_tokens"]):
#             context_idxs[n][i] = _get_word(token)
#
#         for i, token in enumerate(example["ques_tokens"]):
#             ques_idxs[n][i] = _get_word(token)
#
#         for i, token in enumerate(example["context_chars"]):
#             for j, char in enumerate(token):
#                 if j == word_len:
#                     break
#                 context_char_idxs[n, i, j] = _get_char(char)
#
#         for i, token in enumerate(example["ques_chars"]):
#             for j, char in enumerate(token):
#                 if j == word_len:
#                     break
#                 ques_char_idxs[n, i, j] = _get_char(char)
#
#         start, end = example["y1s"][-1], example["y2s"][-1]
#         y1s[n], y2s[n] = start, end
#         ids[n] = example["id"]
#
#     np.savez(out_file, context_idxs=context_idxs, context_char_idxs=context_char_idxs, ques_idxs=ques_idxs,
#              ques_char_idxs=ques_char_idxs, y1s=y1s, y2s=y2s, ids=ids)
#     print("Built {} / {} instances of features in total".format(total, total_))
#     meta["total"] = total
#     return meta

def build_features(config, examples, data_type, out_file,word2idx_dict, char2idx_dict, is_test=False):
    para_limit = config.test_para_limit if is_test else config.para_limit
    ques_limit = config.test_ques_limit if is_test else config.ques_limit
    ans_limit = 100 if is_test else config.ans_limit
    word_len = config.word_len

    def filter_func(example, is_test=False):
        return len(example["context_tokens"]) > para_limit or \
               len(example["ques_tokens"]) > ques_limit or \
               (example["y2s"] - example["y1s"]) > ans_limit

    print("Processing {} examples...".format(data_type))
    total = 0
    meta = {}
    N = len(examples)
    example_ids=[]
    for n, example in tqdm(enumerate(examples)):
        new_example = {"context_tokens": [], "context_chars": [],
            "ques_tokens": [],
            "ques_chars": [], "y1s": 0, "y2s": 0, "id": 0}

        if filter_func(example, is_test):
            continue

        total+=1

        def _get_word(word):
            for each in (word, word.lower(), word.capitalize(), word.upper()):
                if each in word2idx_dict:
                    return word2idx_dict[each]
            return 1

        def _get_char(char):
            if char in char2idx_dict:
                return char2idx_dict[char]
            return 1

        for token in example["context_tokens"]:
            new_example['context_tokens'].append(_get_word(token))

        for token in example["ques_tokens"]:
            new_example['ques_tokens'].append(_get_word(token))

        for token in example["context_chars"]:
            chars=[]
            for j, char in enumerate(token):
                if j == word_len:
                    break
                chars.append(_get_char(char))
            chars+=[0]*(word_len-len(token))
            new_example['context_chars'].append(chars)

        for token in example["ques_chars"]:
            chars=[]
            for j, char in enumerate(token):
                if j == word_len:
                    break
                chars.append(_get_char(char))
            chars+=[0]*(word_len-len(token))
            new_example['ques_chars'].append(chars)

        new_example['y1s']=example["y1s"]
        new_example['y2s']=example["y2s"]
        new_example['id']=example["id"]
        example_ids.append(new_example)

    save(out_file,example_ids,message=out_file)
    print("Built {} / {} instances of features in total".format(total, N))
    meta["total"] = total
    return meta

def save(filename, obj, message=None):
    if message is not None:
        print("Saving {}...".format(message))
        with open(filename, "w") as fh:
            json.dump(obj, fh)


def preproc(config):
    word_counter, char_counter = Counter(), Counter()
    train_examples, train_eval = process_file(
        config.train_file, "train", word_counter, char_counter)
    dev_examples, dev_eval = process_file(
        config.dev_file, "dev", word_counter, char_counter)
    test_examples, test_eval = process_file(
        config.test_file, "test", word_counter, char_counter)

    word_emb_file = config.fasttext_file if config.fasttext else config.glove_word_file
    char_emb_file = config.glove_char_file if config.pretrained_char else None
    char_emb_size = config.glove_char_size if config.pretrained_char else None
    char_emb_dim = config.glove_dim if config.pretrained_char else config.char_dim

    word_emb_mat, word2idx_dict = get_embedding(
        word_counter, "word", emb_file=word_emb_file, vec_size=config.glove_dim)
    char_emb_mat, char2idx_dict = get_embedding(
        char_counter, "char", emb_file=char_emb_file, vec_size=char_emb_dim)

    build_features(config, train_examples, "train",
                   config.train_token_file,word2idx_dict, char2idx_dict)
    dev_meta = build_features(config, dev_examples, "dev",
                              config.dev_token_file, word2idx_dict, char2idx_dict)
    test_meta = build_features(config, test_examples, "test",
                               config.test_token_file, word2idx_dict, char2idx_dict, is_test=True)

    save(config.word_emb_file, word_emb_mat, message="word embedding")
    # save(config.char_emb_file, char2idx_dict, message="char dict")
    save(config.train_eval_file, train_eval, message="train eval")
    save(config.dev_eval_file, dev_eval, message="dev eval")
    save(config.test_eval_file, test_eval, message="test eval")
    save(config.dev_meta, dev_meta, message="dev meta")
    save(config.test_meta, test_meta, message="test meta")
    save(config.word2idx_file, word2idx_dict, message="word dictionary")
    save(config.char2idx_file, char2idx_dict, message="char dictionary")