lmqg/qg_squad

1""" Script to process raw SQuAD file for Question Generation format2You need to run `python -m spacy download en_core_web_sm`.3Split when uploading to dataset hub by4```5gsplit -l 3300 -d --additional-suffix=.jsonl train.jsonl train6gsplit -l 3300 -d --additional-suffix=.jsonl test.jsonl test7gsplit -l 3300 -d --additional-suffix=.jsonl dev.jsonl dev8```9"""10import json11import os12import re13from glob import glob14from tqdm import tqdm15from typing import List, Dict16 17import spacy18 19SPLITTER = spacy.load('en_core_web_sm')20HIGHLIGHT_TOKEN = '<hl>'21 22 23def get_sentence(document: str):24    return [str(s) for s in SPLITTER(document).sents]25 26 27def jsonline_reader(filename: str):28    with open(filename, 'r') as f_reader:29        examples = [json.loads(i) for i in f_reader.read().split('\n') if len(i) > 0]30    return examples31 32 33def process_single_data(data: Dict):34    """ Convert single raw json data into QG format """35    example = {'question': data["question"], 'paragraph': data["context"], 'answer': data["answer"]}36 37    # get sentence38    position = example['paragraph'].find(example['answer'])39    assert position != -140    before_tmp = get_sentence(example['paragraph'][:position])41    if len(before_tmp) == 0:42        before = ''43        before_sentence = ''44    else:45        if before_tmp[-1].endswith('.'):46            before = ' '.join(before_tmp)47            before_sentence = ''48        else:49            before = ' '.join(before_tmp[:-1])50            before_sentence = before_tmp[-1]51            before_sentence = before_sentence if before_sentence.endswith(' ') else '{} '.format(before_sentence)52    after_tmp = get_sentence(example['paragraph'][position + len(example['answer']):])53    if len(after_tmp) == 0:54        after = ''55        after_sentence = ''56    else:57        after = ' '.join(after_tmp[1:])58        after_sentence = after_tmp[0]59        after_sentence = after_sentence if after_sentence.startswith(' ') else ' {}'.format(after_sentence)60    example['sentence'] = '{}{}{}'.format(before_sentence, example['answer'], after_sentence)61 62    # get paragraph_sentence63    before = '' if before == '' else '{} '.format(before)64    after = '' if after == '' else ' {}'.format(after)65    source_text = '{0}{1} {2} {1}{3}'.format(before, HIGHLIGHT_TOKEN, example['sentence'], after)66    example['paragraph_sentence'] = re.sub(r'\s+', ' ', source_text)67 68    # get paragraph_answer69    source_text = '{0}{1} {2} {1}{3}'.format(70        example['paragraph'][:position], HIGHLIGHT_TOKEN, example['answer'],71        example['paragraph'][position + len(example['answer']):])72    example['paragraph_answer'] = re.sub(r'\s+', ' ', source_text)73 74    # get sentence_answer75    if len(before_tmp) == 0 or before_tmp[-1].endswith('.'):76        before = ''77    else:78        before = before_tmp[-1] if before_tmp[-1].endswith(' ') else '{} '.format(before_tmp[-1])79    if len(after_tmp) == 0:80        after = ''81    else:82        after = after_tmp[0] if after_tmp[0].startswith(' ') else ' {}'.format(after_tmp[0])83    source_text = '{0}{1} {2} {1}{3}'.format(before, HIGHLIGHT_TOKEN, example['answer'], after)84    example['sentence_answer'] = re.sub(r'\s+', ' ', source_text)85 86    return example87 88 89if __name__ == '__main__':90    output = './data/processed'91    os.makedirs(output, exist_ok=True)92    path = {'train': 'data/raw/train*.jsonl', 'dev': 'data/raw/dev.jsonl', 'test': 'data/raw/test.jsonl'}93    for k, v in path.items():94        json_data = []95        for _file in sorted(glob(v)):96            json_data += jsonline_reader(_file)97        with open('{}/{}.jsonl'.format(output, k), 'w') as f:98            for single_data in tqdm(json_data):99                single_data = process_single_data(single_data)100                f.write(json.dumps(single_data) + '\n')101