from sentence_transformers import SentenceTransformer
from sentence_transformers.sentence_transformer.evaluation import InformationRetrievalEvaluator

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

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"},
}

evaluator = InformationRetrievalEvaluator(
    queries=queries,
    corpus=corpus,
    relevant_docs=relevant_docs,
    name="support-search-dev",
    accuracy_at_k=[1, 3],
    precision_recall_at_k=[1, 3],
    mrr_at_k=[3],
    ndcg_at_k=[3],
    map_at_k=[3],
    show_progress_bar=False,
    write_csv=False,
)

results = evaluator(model)

for metric in (
    "support-search-dev_cosine_accuracy@1",
    "support-search-dev_cosine_recall@3",
    "support-search-dev_cosine_mrr@3",
    "support-search-dev_cosine_ndcg@3",
):
    print(f"{metric}: {results[metric]:.3f}")

print(f"primary metric: {evaluator.primary_metric}")
print(f"primary score: {results[evaluator.primary_metric]:.3f}")