# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.8.12 64-bit (''re_dev'': conda)'
#     name: python3
# ---

# # This notebook contains the code to reproduce the experiment in `ranx` paper.

# +
import random

import numpy as np
import pytrec_eval
from tabulate import tabulate

from ranx import evaluate, Qrels, Run


# +
# Wrapper for pytrec_eval
def run_trec_metrics(qrels, run, metrics):
    evaluator = pytrec_eval.RelevanceEvaluator(qrels, metrics)
    results = evaluator.evaluate(run)
    return {m: np.mean([v[m] for v in results.values()]) for m in list(metrics)}


def run_single_trec_metric(qrels, run, metric):
    return run_trec_metrics(qrels, run, {metric})[metric]


# +
def generate_qrels(query_count, max_relevant_per_query):
    qrels = {}
    for i in range(query_count):
        y_t = {}
        k = random.choice(range(1, max_relevant_per_query))
        for j in range(k):
            y_t[f"d{j}"] = random.choice([1, 2, 3, 4, 5])

        qrels[f"q{i}"] = y_t

    return qrels


def generate_run(query_count, result_count):
    run = {}
    for i in range(query_count):
        y_p = {}
        for j in range(result_count):
            y_p[f"d{j}"] = random.uniform(0.0, 1.0)

        run[f"q{i}"] = y_p

    return run


# -

random.seed = 42
np.random.seed(42)

# +
results = []
result_count = 100
max_relevant_per_query = 10

for query_count in [1_000_000]:
# for query_count in [1_000, 10_000, 100_000]:
    print(f"Queries: {query_count}")
    # Generate Qrels and Run
    trec_qrels = generate_qrels(query_count, max_relevant_per_query)
    trec_run = generate_run(query_count, result_count)

    re_qrels = Qrels.from_dict(trec_qrels).to_typed_list()
    re_run = Run.from_dict(trec_run).to_typed_list()

    # x = %timeit -o -q run_single_trec_metric(trec_qrels, trec_run, "map")
    map_avg_time = round(x.average, 3) * 1000
    results.append([query_count, "pytrec_eval", 1, "map", map_avg_time, 1.0])

    # x = %timeit -o -q run_single_trec_metric(trec_qrels, trec_run, "recip_rank")
    mrr_avg_time = round(x.average, 3) * 1000
    results.append([query_count, "pytrec_eval", 1, "mrr", mrr_avg_time, 1.0])

    # x = %timeit -o -q run_single_trec_metric(trec_qrels, trec_run, "ndcg")
    ndcg_avg_time = round(x.average, 3) * 1000
    results.append([query_count, "pytrec_eval", 1, "ndcg", ndcg_avg_time, 1.0])

    for threads in [1, 2, 4, 8]:
        # Run metrics once to ensure they have been compiled
        evaluate(re_qrels, re_run, [f"map", f"mrr", f"ndcg"], threads=threads)

        # x = %timeit -o -q evaluate(re_qrels, re_run, f"map", threads=threads)
        avg_time = max(round(x.average, 3) * 1000, 1)
        results.append([query_count, "ranx", threads, "map", avg_time, round(map_avg_time / avg_time, 1)])

        # x = %timeit -o -q evaluate(re_qrels, re_run, f"mrr", threads=threads)
        avg_time = max(round(x.average, 3) * 1000, 1)
        results.append([query_count, "ranx", threads, "mrr", avg_time, round(mrr_avg_time / avg_time, 1)])

        # x = %timeit -o -q evaluate(re_qrels, re_run, f"ndcg", threads=threads)
        avg_time = max(round(x.average, 3) * 1000, 1)
        results.append([query_count, "ranx", threads, "ndcg", avg_time, round(ndcg_avg_time / avg_time, 1)])
# -

print(tabulate(results, headers=["Query count", "Approach", "Threads", "Metric", "Avg. (ms)", "Speed-Up"]))



StarcoderdataJupyterScriptsDedupFiltered