from sentence_transformers import SparseEncoder
from sentence_transformers.sparse_encoder.evaluation import (
    SparseInformationRetrievalEvaluator,
)


model = SparseEncoder(
    "naver/splade-cocondenser-ensembledistil",
    device="cpu",
)

queries = {
    "q1": "How do I reset a forgotten password?",
    "q2": "How can I export invoices as a CSV file?",
}

corpus = {
    "d1": "Reset a lost account password from the profile security page.",
    "d2": "Export paid invoices from the billing dashboard as a CSV file.",
    "d3": "Change the color theme for the analytics workspace.",
    "d4": "Archive an inactive user without deleting historical records.",
}

relevant_docs = {
    "q1": {"d1"},
    "q2": {"d2"},
}


def run_evaluator(name, max_active_dims):
    evaluator = SparseInformationRetrievalEvaluator(
        queries=queries,
        corpus=corpus,
        relevant_docs=relevant_docs,
        name=name,
        accuracy_at_k=[1, 3],
        precision_recall_at_k=[1, 3],
        mrr_at_k=[3],
        ndcg_at_k=[3],
        map_at_k=[3],
        max_active_dims=max_active_dims,
        show_progress_bar=False,
        write_csv=False,
    )

    results = evaluator(model)
    print(f"{name} primary metric: {evaluator.primary_metric}")
    print(f"{name} primary score: {results[evaluator.primary_metric]:.3f}")
    for metric in (
        f"{name}_dot_accuracy@1",
        f"{name}_dot_recall@3",
        f"{name}_dot_mrr@3",
        f"{name}_dot_ndcg@3",
    ):
        print(f"{metric}: {results[metric]:.3f}")


run_evaluator("support-sparse-full", None)
run_evaluator("support-sparse-64", 64)