from sentence_transformers import SentenceTransformer
from sentence_transformers.evaluation import InformationRetrievalEvaluator

model = SentenceTransformer("all-MiniLM-L6-v2")

queries = {
    "q1": "who is alfie",
    "q2": "which pet is maple",
    "q3": "tell me about poplar",
}

corpus = {
    "d1": "Alfie is a playful dog who likes long walks and tennis balls.",
    "d2": "Maple is a calm orange cat who naps in sunny windows.",
    "d3": "Poplar is a curious gray cat who chases toy mice at night.",
    "d4": "Alfie the dog enjoys treats after learning new tricks.",
    "d5": "Maple the cat loves soft blankets and quiet afternoons.",
    "d6": "Poplar the cat likes climbing shelves and watching birds.",
}

relevant_docs = {
    "q1": {"d1", "d4"},
    "q2": {"d2", "d5"},
    "q3": {"d3", "d6"},
}

ir_evaluator = InformationRetrievalEvaluator(
    queries=queries,
    corpus=corpus,
    relevant_docs=relevant_docs,
    write_predictions=True,
)

out_dir = "./output"
ir_evaluator(model, output_path=out_dir)  # this would fail before