wehub-resource-sync
123 lines
5.2 KiB
Python
123 lines
5.2 KiB
Python
import os
|
|
import json
|
|
import argparse
|
|
import sys
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, 'src/')
|
|
|
|
from tqdm import tqdm
|
|
from typing import Dict, Any
|
|
|
|
from logger_config import logger
|
|
from data_utils import load_query_answers, load_corpus, save_to_readable_format
|
|
|
|
parser = argparse.ArgumentParser(description='data preprocessing for NQ & TriviaQA in DPR paper')
|
|
parser.add_argument('--out-dir', default='./data/dpr/', type=str, metavar='N',
|
|
help='output directory')
|
|
parser.add_argument('--task', default='nq', type=str, metavar='N',
|
|
help='task name, nq or tq')
|
|
parser.add_argument('--train-pred-path', default='amlt/0621_cont_100k_psg16_ft/nq/nq_train.dpr.json',
|
|
type=str, metavar='N', help='path to train predictions to construct negatives')
|
|
parser.add_argument('--dev-pred-path', default='amlt/0621_cont_100k_psg16_ft/nq/nq_dev.dpr.json',
|
|
type=str, metavar='N', help='path to dev predictions to construct negatives')
|
|
parser.add_argument('--num-negatives', default=100, type=int, metavar='N',
|
|
help='number of negative passages')
|
|
parser.add_argument('--depth', default=100, type=int, metavar='N',
|
|
help='depth to choose negative passages from')
|
|
|
|
args = parser.parse_args()
|
|
os.makedirs(args.out_dir, exist_ok=True)
|
|
logger.info('Args: {}'.format(json.dumps(args.__dict__, ensure_ascii=False, indent=4)))
|
|
assert args.task in ['nq', 'tq']
|
|
|
|
|
|
def _load_qid_to_positive(path: str) -> Dict[str, str]:
|
|
if args.task != 'nq':
|
|
logger.warning('Only NQ has manually labeled positives')
|
|
return {}
|
|
|
|
examples = json.load(open(path, 'r', encoding='utf-8'))
|
|
qid_to_pos_id = {}
|
|
for ex in examples:
|
|
positive_ctxs = ex['positive_ctxs']
|
|
if len(positive_ctxs) > 0:
|
|
qid_to_pos_id[ex['question']] = str(int(positive_ctxs[0]['passage_id']) - 1)
|
|
|
|
logger.info('Get {} manually labeled positives from {}'.format(len(qid_to_pos_id), path))
|
|
return qid_to_pos_id
|
|
|
|
|
|
def _write_prepared_data_to_disk(out_path: str,
|
|
split: str,
|
|
queries: Dict[str, Dict[str, Any]],
|
|
preds_path: str):
|
|
qid_to_pos_id = _load_qid_to_positive(path='{}/biencoder-nq-{}.json'.format(args.out_dir, split))
|
|
cnt_filtered = 0
|
|
|
|
preds = json.load(open(preds_path, 'r', encoding='utf-8'))
|
|
|
|
with open(out_path, 'w', encoding='utf-8') as writer:
|
|
for query_id in tqdm(queries, mininterval=1, desc='prepare {} data'.format(split)):
|
|
cur_pred: dict = preds[query_id] if query_id in preds else preds[query_id.strip()]
|
|
|
|
positive_ids, negative_ids = [], []
|
|
manual_positive_id = qid_to_pos_id.get(query_id, None)
|
|
if manual_positive_id:
|
|
positive_ids.append(manual_positive_id)
|
|
|
|
for ctx in cur_pred['contexts'][:args.depth]:
|
|
doc_id = str(ctx['docid'])
|
|
if doc_id == manual_positive_id:
|
|
continue
|
|
elif ctx['has_answer']:
|
|
positive_ids.append(doc_id)
|
|
else:
|
|
negative_ids.append(doc_id)
|
|
|
|
if not positive_ids or not negative_ids:
|
|
cnt_filtered += 1
|
|
continue
|
|
|
|
np.random.shuffle(negative_ids)
|
|
negative_ids = negative_ids[:args.num_negatives]
|
|
|
|
doc_id_to_score = {str(ctx['docid']): float(ctx['score']) for ctx in cur_pred['contexts']}
|
|
doc_id_to_score[manual_positive_id] = 1000.
|
|
example = {
|
|
'query_id': query_id,
|
|
'query': queries[query_id]['query'],
|
|
'answers': queries[query_id]['answers'],
|
|
'positives': {'doc_id': positive_ids,
|
|
'score': [doc_id_to_score.get(doc_id, -1.) for doc_id in positive_ids]
|
|
},
|
|
'negatives': {'doc_id': negative_ids,
|
|
'score': [doc_id_to_score.get(doc_id, -1.) for doc_id in negative_ids]
|
|
},
|
|
}
|
|
|
|
writer.write(json.dumps(example, ensure_ascii=False, separators=(',', ':')))
|
|
writer.write('\n')
|
|
|
|
if cnt_filtered > 0:
|
|
logger.info('{} questions are filtered out'.format(cnt_filtered))
|
|
|
|
logger.info('Done write {} data to {}'.format(split, out_path))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
for split in ['dev', 'train']:
|
|
_write_prepared_data_to_disk(
|
|
out_path=os.path.join(args.out_dir, '{}_hard_{}.jsonl'.format(args.task, split)),
|
|
split=split,
|
|
queries=load_query_answers(path=os.path.join(args.out_dir, '{}_{}_queries.tsv'.format(args.task, split))),
|
|
preds_path=(args.train_pred_path if split == 'train' else args.dev_pred_path)
|
|
)
|
|
|
|
corpus = load_corpus(path=os.path.join(args.out_dir, 'passages.jsonl.gz'))
|
|
for split in ['dev', 'train']:
|
|
save_to_readable_format(in_path=os.path.join(args.out_dir, '{}_hard_{}.jsonl'.format(args.task, split)),
|
|
corpus=corpus)
|
|
|
|
logger.info('Done')
|