tune_bm25_with_skopt.py

#
# 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 sys, os
import re
import subprocess
from tqdm import tqdm
from skopt import gp_minimize
from skopt.space import Real, Integer

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)')
parser.add_argument('--skopt-iters', type=int, default=20, help='Iteration of bayesian optimization')
parser.add_argument('--hits', type=int, default=1000, help='Number of hits')
parser.add_argument('--threads', type=int, default=1, help='Number of threads')
parser.add_argument('--metric', type=str, default='recall', help='Metric for tuning')
parser.add_argument('--seed', type=int, default=0, help='Random seed')

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
iters = args.skopt_iters
hits = args.hits
threads = args.threads
metric = args.metric
seed = args.seed

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')

def objective(weight):
    k1, b = weight
    print(f'Trying... k1 = {k1:.2f}, b = {b:.2f}')
    filename = f'run.bm25.k1_{k1:.2f}.b_{b:.2f}.txt'
    if not os.path.isfile(f'{base_directory}/{filename}'):
        subprocess.call(f'python tools/scripts/msmarco/retrieve.py --index {index} --queries {queries} \
            --output {base_directory}/{filename} --k1 {k1:.2f} --b {b:.2f} --hits {hits} --threads {threads}', shell=True)
        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', f'-mrecall.{hits}', '-mmap'])
    match = re.search('map +\tall\t([0-9.]+)', results.decode('utf-8'))
    ap = float(match.group(1))
    match = re.search(f'recall_{hits} +\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@{hits} = {recall}')
    d = dict(recall=recall, mrr=rr, ap=ap)
    return -d[metric]

space  = [Real(0, 2, name='k1'),
          Real(0, 1, name='b')]
res_gp = gp_minimize(objective, space, n_calls=iters, random_state=seed)
print(f'\n\nBest parameters: {res_gp.x}, Best {metric}: {res_gp.fun}')