def _base_qrels(self, subset):
rnd, _ = subset.split('-', 1)
path = os.path.join(util.path_dataset(self), f'{rnd}.qrels')
if os.path.exists(path):
return trec.read_qrels_dict(path)
self.logger.info(f'missing qrels for {rnd} -- returning empty qrels')
return {}
def init(self, force=False):
base_path = util.path_dataset(self)
idxs = [self.index, self.index_stem, self.doc_store]
self._init_indices_parallel(idxs, self._init_iter_collection(), force)
qrels_file = os.path.join(base_path, 'qrels.robust2004.txt')
if (force or not os.path.exists(qrels_file)) and self._confirm_dua():
util.download(**_FILES['qrels'], file_name=qrels_file)
for fold in FOLDS:
fold_qrels_file = os.path.join(base_path, f'{fold}.qrels')
if (force or not os.path.exists(fold_qrels_file)):
all_qrels = trec.read_qrels_dict(qrels_file)
fold_qrels = {
qid: dids
for qid, dids in all_qrels.items() if qid in FOLDS[fold]
}
trec.write_qrels_dict(fold_qrels_file, fold_qrels)
query_file = os.path.join(base_path, 'topics.txt')
if (force or not os.path.exists(query_file)) and self._confirm_dua():
query_file_stream = util.download_stream(**_FILES['queries'],
encoding='utf8')
with util.finialized_file(query_file, 'wt') as f:
plaintext.write_tsv(f,
trec.parse_query_format(query_file_stream))
def qrels_path(self, fold_qrels_file):
if not fold_qrels_file.is_file():
with self.assessments.path.open("r") as fp:
all_qrels = trec.read_qrels_dict(fp)
fold_qrels = {
qid: dids
for qid, dids in all_qrels.items() if qid in self.qids
}
trec.write_qrels_dict(fold_qrels_file, fold_qrels)
return fold_qrels_file
def calc_metrics(self, qrelsf, runf, metrics, verbose=False):
rel_args = {}
output_map = {}
for metric in metrics:
for exp in TE_METRICS:
match = re.match(exp, str(metric))
if match:
params = match.groupdict()
rel, par, gain = params.get('rel'), params.get(
'par'), params.get('gain')
rel_args.setdefault((rel, gain), {})
in_name, out_name = TE_METRICS[exp]
if in_name not in rel_args[rel, gain]:
rel_args[rel, gain][in_name] = in_name
if par is not None:
rel_args[rel, gain][in_name] += par
elif par is not None:
rel_args[rel, gain][in_name] += f',{par}'
output_map[metric] = (rel, gain, out_name.format(**params))
break
full_trec_eval_result = {}
for rel_level, gains in rel_args:
with ExitStack() as stack:
if verbose:
logger.debug(
f'running trec_eval (rel_level={rel_level} gains={gains})'
)
if gains:
gain_map = [g.split('=') for g in gains.split(':')]
gain_map = {int(k): int(v) for k, v in gain_map}
tmpf = tempfile.NamedTemporaryFile()
stack.enter_context(tmpf)
qrels = trec.read_qrels_dict(qrelsf)
for qid in qrels:
for did in qrels[qid]:
qrels[qid][did] = gain_map.get(
qrels[qid][did], qrels[qid][did])
trec.write_run_dict(tmpf, qrels)
qrelsf_run = tmpf.name
else:
qrelsf_run = qrelsf
trec_eval_result = self._treceval(
qrelsf_run, runf, rel_args[rel_level, gains].values(),
rel_level)
for k, v in trec_eval_result.items():
full_trec_eval_result[rel_level, gains, k] = v
results = {}
for onir_metric, te_metric in output_map.items():
results[str(onir_metric)] = full_trec_eval_result[te_metric]
return results
def init(self, force=False):
idxs = [self.index_stem, self.doc_store]
self._init_indices_parallel(idxs, self._init_iter_collection(), force)
base_path = util.path_dataset(self)
needs_queries = []
if force or not os.path.exists(
os.path.join(base_path, 'train.queries.tsv')):
needs_queries.append(lambda it: plaintext.write_tsv(
os.path.join(base_path, 'train.queries.tsv'),
((qid, txt) for file, qid, txt in it
if file == 'queries.train.tsv' and qid not in MINI_DEV)))
if force or not os.path.exists(
os.path.join(base_path, 'minidev.queries.tsv')):
needs_queries.append(lambda it: plaintext.write_tsv(
os.path.join(base_path, 'minidev.queries.tsv'),
((qid, txt) for file, qid, txt in it
if file == 'queries.train.tsv' and qid in MINI_DEV)))
if force or not os.path.exists(
os.path.join(base_path, 'dev.queries.tsv')):
needs_queries.append(lambda it: plaintext.write_tsv(
os.path.join(base_path, 'dev.queries.tsv'),
((qid, txt) for file, qid, txt in it
if file == 'queries.dev.tsv')))
if force or not os.path.exists(
os.path.join(base_path, 'eval.queries.tsv')):
needs_queries.append(lambda it: plaintext.write_tsv(
os.path.join(base_path, 'eval.queries.tsv'),
((qid, txt) for file, qid, txt in it
if file == 'queries.eval.tsv')))
if needs_queries and self._confirm_dua():
with util.download_tmp(_SOURCES['queries']) as f, \
tarfile.open(fileobj=f) as tarf, \
contextlib.ExitStack() as ctxt:
def _extr_subf(subf):
for qid, txt in plaintext.read_tsv(
io.TextIOWrapper(tarf.extractfile(subf))):
yield subf, qid, txt
query_iter = [
_extr_subf('queries.train.tsv'),
_extr_subf('queries.dev.tsv'),
_extr_subf('queries.eval.tsv')
]
query_iter = tqdm(itertools.chain(*query_iter), desc='queries')
query_iters = util.blocking_tee(query_iter, len(needs_queries))
for fn, it in zip(needs_queries, query_iters):
ctxt.enter_context(
util.CtxtThread(functools.partial(fn, it)))
file = os.path.join(base_path, 'train.qrels')
if (force or not os.path.exists(file)) and self._confirm_dua():
stream = util.download_stream(_SOURCES['train-qrels'], 'utf8')
with util.finialized_file(file, 'wt') as out:
for qid, _, did, score in plaintext.read_tsv(stream):
if qid not in MINI_DEV:
trec.write_qrels(out, [(qid, did, score)])
file = os.path.join(base_path, 'minidev.qrels')
if (force or not os.path.exists(file)) and self._confirm_dua():
stream = util.download_stream(_SOURCES['train-qrels'], 'utf8')
with util.finialized_file(file, 'wt') as out:
for qid, _, did, score in plaintext.read_tsv(stream):
if qid in MINI_DEV:
trec.write_qrels(out, [(qid, did, score)])
file = os.path.join(base_path, 'dev.qrels')
if (force or not os.path.exists(file)) and self._confirm_dua():
stream = util.download_stream(_SOURCES['dev-qrels'], 'utf8')
with util.finialized_file(file, 'wt') as out:
for qid, _, did, score in plaintext.read_tsv(stream):
trec.write_qrels(out, [(qid, did, score)])
file = os.path.join(base_path, 'train.mspairs.gz')
if not os.path.exists(file) and os.path.exists(
os.path.join(base_path, 'qidpidtriples.train.full')):
# legacy
os.rename(os.path.join(base_path, 'qidpidtriples.train.full'),
file)
if (force or not os.path.exists(file)) and self._confirm_dua():
util.download(_SOURCES['qidpidtriples.train.full'], file)
if not self.config['init_skip_msrun']:
for file_name, subf in [('dev.msrun', 'top1000.dev'),
('eval.msrun', 'top1000.eval'),
('train.msrun', 'top1000.train.txt')]:
file = os.path.join(base_path, file_name)
if (force or not os.path.exists(file)) and self._confirm_dua():
run = {}
with util.download_tmp(_SOURCES[file_name]) as f, \
tarfile.open(fileobj=f) as tarf:
for qid, did, _, _ in tqdm(
plaintext.read_tsv(
io.TextIOWrapper(tarf.extractfile(subf)))):
if qid not in run:
run[qid] = {}
run[qid][did] = 0.
if file_name == 'train.msrun':
minidev = {
qid: dids
for qid, dids in run.items() if qid in MINI_DEV
}
with self.logger.duration('writing minidev.msrun'):
trec.write_run_dict(
os.path.join(base_path, 'minidev.msrun'),
minidev)
run = {
qid: dids
for qid, dids in run.items() if qid not in MINI_DEV
}
with self.logger.duration(f'writing {file_name}'):
trec.write_run_dict(file, run)
query_path = os.path.join(base_path, 'trec2019.queries.tsv')
if (force or not os.path.exists(query_path)) and self._confirm_dua():
stream = util.download_stream(_SOURCES['trec2019.queries'], 'utf8')
plaintext.write_tsv(query_path, plaintext.read_tsv(stream))
msrun_path = os.path.join(base_path, 'trec2019.msrun')
if (force or not os.path.exists(msrun_path)) and self._confirm_dua():
run = {}
with util.download_stream(_SOURCES['trec2019.msrun'],
'utf8') as stream:
for qid, did, _, _ in plaintext.read_tsv(stream):
if qid not in run:
run[qid] = {}
run[qid][did] = 0.
with util.finialized_file(msrun_path, 'wt') as f:
trec.write_run_dict(f, run)
qrels_path = os.path.join(base_path, 'trec2019.qrels')
if not os.path.exists(qrels_path) and self._confirm_dua():
util.download(_SOURCES['trec2019.qrels'], qrels_path)
qrels_path = os.path.join(base_path, 'judgedtrec2019.qrels')
if not os.path.exists(qrels_path):
os.symlink('trec2019.qrels', qrels_path)
query_path = os.path.join(base_path, 'judgedtrec2019.queries.tsv')
judged_qids = util.Lazy(
lambda: trec.read_qrels_dict(qrels_path).keys())
if (force or not os.path.exists(query_path)):
with util.finialized_file(query_path, 'wt') as f:
for qid, qtext in plaintext.read_tsv(
os.path.join(base_path, 'trec2019.queries.tsv')):
if qid in judged_qids():
plaintext.write_tsv(f, [(qid, qtext)])
msrun_path = os.path.join(base_path, 'judgedtrec2019.msrun')
if (force or not os.path.exists(msrun_path)) and self._confirm_dua():
with util.finialized_file(msrun_path, 'wt') as f:
for qid, dids in trec.read_run_dict(
os.path.join(base_path, 'trec2019.msrun')).items():
if qid in judged_qids():
trec.write_run_dict(f, {qid: dids})
# A subset of dev that only contains queries that have relevance judgments
judgeddev_path = os.path.join(base_path, 'judgeddev')
judged_qids = util.Lazy(lambda: trec.read_qrels_dict(
os.path.join(base_path, 'dev.qrels')).keys())
if not os.path.exists(f'{judgeddev_path}.qrels'):
os.symlink('dev.qrels', f'{judgeddev_path}.qrels')
if not os.path.exists(f'{judgeddev_path}.queries.tsv'):
with util.finialized_file(f'{judgeddev_path}.queries.tsv',
'wt') as f:
for qid, qtext in plaintext.read_tsv(
os.path.join(base_path, 'dev.queries.tsv')):
if qid in judged_qids():
plaintext.write_tsv(f, [(qid, qtext)])
if self.config['init_skip_msrun']:
if not os.path.exists(f'{judgeddev_path}.msrun'):
with util.finialized_file(f'{judgeddev_path}.msrun',
'wt') as f:
for qid, dids in trec.read_run_dict(
os.path.join(base_path, 'dev.msrun')).items():
if qid in judged_qids():
trec.write_run_dict(f, {qid: dids})
if not self.config['init_skip_train10']:
file = os.path.join(base_path, 'train10.queries.tsv')
if not os.path.exists(file):
with util.finialized_file(file, 'wt') as fout:
for qid, qtext in self.logger.pbar(
plaintext.read_tsv(
os.path.join(base_path, 'train.queries.tsv')),
desc='filtering queries for train10'):
if int(qid) % 10 == 0:
plaintext.write_tsv(fout, [(qid, qtext)])
file = os.path.join(base_path, 'train10.qrels')
if not os.path.exists(file):
with util.finialized_file(file, 'wt') as fout, open(
os.path.join(base_path, 'train.qrels'), 'rt') as fin:
for line in self.logger.pbar(
fin, desc='filtering qrels for train10'):
qid = line.split()[0]
if int(qid) % 10 == 0:
fout.write(line)
if not self.config['init_skip_msrun']:
file = os.path.join(base_path, 'train10.msrun')
if not os.path.exists(file):
with util.finialized_file(file, 'wt') as fout, open(
os.path.join(base_path, 'train.msrun'),
'rt') as fin:
for line in self.logger.pbar(
fin, desc='filtering msrun for train10'):
qid = line.split()[0]
if int(qid) % 10 == 0:
fout.write(line)
file = os.path.join(base_path, 'train10.mspairs.gz')
if not os.path.exists(file):
with gzip.open(file, 'wt') as fout, gzip.open(
os.path.join(base_path, 'train.mspairs.gz'),
'rt') as fin:
for qid, did1, did2 in self.logger.pbar(
plaintext.read_tsv(fin),
desc='filtering mspairs for train10'):
if int(qid) % 10 == 0:
plaintext.write_tsv(fout, [(qid, did1, did2)])
if not self.config['init_skip_train_med']:
med_qids = util.Lazy(
lambda: {
qid.strip()
for qid in util.download_stream(
'https://raw.githubusercontent.com/Georgetown-IR-Lab/covid-neural-ir/master/med-msmarco-train.txt',
'utf8',
expected_md5="dc5199de7d4a872c361f89f08b1163ef")
})
file = os.path.join(base_path, 'train_med.queries.tsv')
if not os.path.exists(file):
with util.finialized_file(file, 'wt') as fout:
for qid, qtext in self.logger.pbar(
plaintext.read_tsv(
os.path.join(base_path, 'train.queries.tsv')),
desc='filtering queries for train_med'):
if qid in med_qids():
plaintext.write_tsv(fout, [(qid, qtext)])
file = os.path.join(base_path, 'train_med.qrels')
if not os.path.exists(file):
with util.finialized_file(file, 'wt') as fout, open(
os.path.join(base_path, 'train.qrels'), 'rt') as fin:
for line in self.logger.pbar(
fin, desc='filtering qrels for train_med'):
qid = line.split()[0]
if qid in med_qids():
fout.write(line)
if not self.config['init_skip_msrun']:
file = os.path.join(base_path, 'train_med.msrun')
if not os.path.exists(file):
with util.finialized_file(file, 'wt') as fout, open(
os.path.join(base_path, 'train.msrun'),
'rt') as fin:
for line in self.logger.pbar(
fin, desc='filtering msrun for train_med'):
qid = line.split()[0]
if qid in med_qids():
fout.write(line)
file = os.path.join(base_path, 'train_med.mspairs.gz')
if not os.path.exists(file):
with gzip.open(file, 'wt') as fout, gzip.open(
os.path.join(base_path, 'train.mspairs.gz'),
'rt') as fin:
for qid, did1, did2 in self.logger.pbar(
plaintext.read_tsv(fin),
desc='filtering mspairs for train_med'):
if qid in med_qids():
plaintext.write_tsv(fout, [(qid, did1, did2)])
def init(self, force=False):
needs_docs = []
for index in [self.index_stem, self.index_stem_2020, self.doc_store]:
if force or not index.built():
needs_docs.append(index)
if needs_docs and self._confirm_dua():
with contextlib.ExitStack() as stack:
doc_iter = self._init_iter_collection()
doc_iter = self.logger.pbar(doc_iter, desc='articles')
doc_iters = util.blocking_tee(doc_iter, len(needs_docs))
for idx, it in zip(needs_docs, doc_iters):
if idx is self.index_stem_2020:
it = (d for d in it if '2020' in d.data['date'])
stack.enter_context(
util.CtxtThread(functools.partial(idx.build, it)))
path = os.path.join(util.path_dataset(self), 'rnd1.tsv')
if not os.path.exists(path) and self._confirm_dua():
with util.download_tmp('https://ir.nist.gov/covidSubmit/data/topics-rnd1.xml', expected_md5="cf1b605222f45f7dbc90ca8e4d9b2c31") as f, \
util.finialized_file(path, 'wt') as fout:
soup = BeautifulSoup(f.read(), 'lxml-xml')
for topic in soup.find_all('topic'):
qid = topic['number']
plaintext.write_tsv(fout, [
(qid, 'query', topic.find('query').get_text()),
(qid, 'quest', topic.find('question').get_text()),
(qid, 'narr', topic.find('narrative').get_text()),
])
udel_flag = path + '.includes_udel'
if not os.path.exists(udel_flag):
with open(path,
'at') as fout, util.finialized_file(udel_flag, 'wt'):
with util.download_tmp(
'https://raw.githubusercontent.com/castorini/anserini/master/src/main/resources/topics-and-qrels/topics.covid-round1-udel.xml',
expected_md5="2915cf59ae222f0aa20b2a671f67fd7a") as f:
soup = BeautifulSoup(f.read(), 'lxml-xml')
for topic in soup.find_all('topic'):
qid = topic['number']
plaintext.write_tsv(fout, [
(qid, 'udel', topic.find('query').get_text()),
])
path = os.path.join(util.path_dataset(self), 'rnd2.tsv')
if not os.path.exists(path) and self._confirm_dua():
with util.download_tmp('https://ir.nist.gov/covidSubmit/data/topics-rnd2.xml', expected_md5="550129e71c83de3fb4d6d29a172c5842") as f, \
util.finialized_file(path, 'wt') as fout:
soup = BeautifulSoup(f.read(), 'lxml-xml')
for topic in soup.find_all('topic'):
qid = topic['number']
plaintext.write_tsv(fout, [
(qid, 'query', topic.find('query').get_text()),
(qid, 'quest', topic.find('question').get_text()),
(qid, 'narr', topic.find('narrative').get_text()),
])
udel_flag = path + '.includes_udel'
if not os.path.exists(udel_flag):
with open(path,
'at') as fout, util.finialized_file(udel_flag, 'wt'):
with util.download_tmp(
'https://raw.githubusercontent.com/castorini/anserini/master/src/main/resources/topics-and-qrels/topics.covid-round2-udel.xml',
expected_md5="a8988734e6f812921d5125249c197985") as f:
soup = BeautifulSoup(f.read(), 'lxml-xml')
for topic in soup.find_all('topic'):
qid = topic['number']
plaintext.write_tsv(fout, [
(qid, 'udel', topic.find('query').get_text()),
])
path = os.path.join(util.path_dataset(self), 'rnd5.tsv')
if not os.path.exists(path) and self._confirm_dua():
with util.download_tmp('https://ir.nist.gov/covidSubmit/data/topics-rnd5.xml', expected_md5="0307a37b6b9f1a5f233340a769d538ea") as f, \
util.finialized_file(path, 'wt') as fout:
soup = BeautifulSoup(f.read(), 'lxml-xml')
for topic in soup.find_all('topic'):
qid = topic['number']
plaintext.write_tsv(fout, [
(qid, 'query', topic.find('query').get_text()),
(qid, 'quest', topic.find('question').get_text()),
(qid, 'narr', topic.find('narrative').get_text()),
])
udel_flag = path + '.includes_udel'
if not os.path.exists(udel_flag):
with open(path,
'at') as fout, util.finialized_file(udel_flag, 'wt'):
with util.download_tmp(
'https://raw.githubusercontent.com/castorini/anserini/master/src/main/resources/topics-and-qrels/topics.covid-round5-udel.xml',
expected_md5="966a49487348dc853634bcdd0829fd26") as f:
soup = BeautifulSoup(f.read(), 'lxml-xml')
for topic in soup.find_all('topic'):
qid = topic['number']
plaintext.write_tsv(fout, [
(qid, 'udel', topic.find('query').get_text()),
])
qrels_file = os.path.join(util.path_dataset(self), 'rnd5.qrels')
if (force or not os.path.exists(qrels_file)) and self._confirm_dua():
util.download(
'https://raw.githubusercontent.com/castorini/anserini/master/src/main/resources/topics-and-qrels/qrels.covid-round4-cumulative.txt',
qrels_file)
for fold in FOLDS:
fold_qrels_file = os.path.join(util.path_dataset(self),
f'{fold}-rnd5.qrels')
if (force or not os.path.exists(fold_qrels_file)):
all_qrels = trec.read_qrels_dict(qrels_file)
fold_qrels = {
qid: dids
for qid, dids in all_qrels.items()
if str(qid) in FOLDS[fold]
}
trec.write_qrels_dict(fold_qrels_file, fold_qrels)