import torch
import mteb

prompts_dict = {
    "BrightRetrieval": "Given a Post, retrieve relevant passages that help answer the post",
}

tasks = mteb.get_tasks(tasks=["BrightRetrieval"])
evaluation = mteb.MTEB(tasks=tasks)

model = mteb.get_model(
    "ReasonIR/ReasonIR-8B",
    model_kwargs={"torch_dtype": torch.bfloat16},
    prompts_dict=prompts_dict,
)

evaluation.run(
    model,
    save_predictions=True,
    output_folder="results",
    encode_kwargs={"batch_size": 1},
)