Spaces:
Runtime error
Runtime error
# | |
# Pyserini: Python interface to the Anserini IR toolkit built on Lucene | |
# | |
# 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. | |
# | |
# Simple script for tuning BM25 parameters (k1 and b) for MS MARCO | |
import argparse | |
import os | |
import re | |
import subprocess | |
parser = argparse.ArgumentParser(description='Tunes BM25 parameters for MS MARCO Passages') | |
parser.add_argument('--base-directory', required=True, help='base directory for storing runs') | |
parser.add_argument('--index', required=True, help='index to use') | |
parser.add_argument('--queries', required=True, help='queries for evaluation') | |
parser.add_argument('--qrels-trec', required=True, help='qrels for evaluation (TREC format)') | |
parser.add_argument('--qrels-tsv', required=True, help='qrels for evaluation (MS MARCO format)') | |
args = parser.parse_args() | |
base_directory = args.base_directory | |
index = args.index | |
queries = args.queries | |
qrels_trec = args.qrels_trec | |
qrels_tsv = args.qrels_tsv | |
if not os.path.exists(base_directory): | |
os.makedirs(base_directory) | |
print('# Settings') | |
print(f'base directory: {base_directory}') | |
print(f'index: {index}') | |
print(f'queries: {queries}') | |
print(f'qrels (TREC): {qrels_trec}') | |
print(f'qrels (MS MARCO): {qrels_tsv}') | |
print('\n') | |
for k1 in [0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2]: | |
for b in [0.5, 0.6, 0.7, 0.8, 0.9]: | |
print(f'Trying... k1 = {k1}, b = {b}') | |
filename = f'run.bm25.k1_{k1}.b_{b}.txt' | |
if os.path.isfile(f'{base_directory}/{filename}'): | |
print('Run already exists, skipping!') | |
else: | |
subprocess.call(f'python tools/scripts/msmarco/retrieve.py --index {index} --queries {queries} \ | |
--output {base_directory}/{filename} --k1 {k1} --b {b} --hits 1000', shell=True) | |
print('\n\nStarting evaluation...') | |
# We're going to be tuning to maximize recall, although we'll compute MRR and MAP also just for reference. | |
max_score = 0 | |
max_file = '' | |
for filename in sorted(os.listdir(base_directory)): | |
# TREC output run file, perhaps left over from a previous tuning run: skip. | |
if filename.endswith('trec'): | |
continue | |
# Convert to a TREC run and evaluate with trec_eval: | |
subprocess.call(f'python tools/scripts/msmarco/convert_msmarco_to_trec_run.py \ | |
--input {base_directory}/{filename} --output {base_directory}/{filename}.trec', shell=True) | |
results = subprocess.check_output(['tools/eval/trec_eval.9.0.4/trec_eval', qrels_trec, | |
f'{base_directory}/{filename}.trec', '-mrecall.1000', '-mmap']) | |
match = re.search('map +\tall\t([0-9.]+)', results.decode('utf-8')) | |
ap = float(match.group(1)) | |
match = re.search('recall_1000 +\tall\t([0-9.]+)', results.decode('utf-8')) | |
recall = float(match.group(1)) | |
# Evaluate with official scoring script | |
results = subprocess.check_output(['python', 'tools/scripts/msmarco/msmarco_passage_eval.py', | |
'collections/msmarco-passage/qrels.train.tsv', | |
f'{base_directory}/{filename}']) | |
match = re.search(r'MRR @10: ([\d.]+)', results.decode('utf-8')) | |
rr = float(match.group(1)) | |
print(f'{filename}: MRR@10 = {rr}, MAP = {ap}, R@1000 = {recall}') | |
if recall > max_score: | |
max_score = recall | |
max_file = filename | |
print(f'\n\nBest parameters: {max_file}: R@1000 = {max_score}') | |