add software issue localization datasets

Author: tarsur909

mteb/evaluation/evaluators/RerankingEvaluator.py

@@ -20,7 +20,7 @@
 class RerankingEvaluator(Evaluator):
     """This class evaluates a SentenceTransformer model for the task of re-ranking.
     Given a query and a list of documents, it computes the score [query, doc_i] for all possible
-    documents and sorts them in decreasing order. Then, MRR@10 and MAP is compute to measure the quality of the ranking.
+    documents and sorts them in decreasing order. Then, MRR@k, MAP, and Recall@k are computed to measure the quality of the ranking.
     :param samples: Must be a list and each element is of the form:
         - {'query': '', 'positive': [], 'negative': []}. Query is the search query, positive is a list of positive
         (relevant) documents, negative is a list of negative (irrelevant) documents.
@@ -143,6 +143,7 @@ def compute_metrics_individual(self, model: Encoder):
     def _encode_candidates(self, model: Encoder, batched: bool, all_query_embs=None):
         all_mrr_scores = []
         all_ap_scores = []
+        all_recall_scores = []
         all_conf_scores = []
         logger.info("Encoding candidates...")
         if batched:
@@ -151,16 +152,18 @@ def _encode_candidates(self, model: Encoder, batched: bool, all_query_embs=None)
                 all_query_embs=all_query_embs,
                 all_mrr_scores=all_mrr_scores,
                 all_ap_scores=all_ap_scores,
+                all_recall_scores=all_recall_scores,
                 all_conf_scores=all_conf_scores,
             )
         else:
             self._encode_candidates_individual(
                 model=model,
                 all_mrr_scores=all_mrr_scores,
                 all_ap_scores=all_ap_scores,
+                all_recall_scores=all_recall_scores,
                 all_conf_scores=all_conf_scores,
             )
-        scores = self._collect_results(all_mrr_scores, all_ap_scores, all_conf_scores)
+        scores = self._collect_results(all_mrr_scores, all_ap_scores, all_recall_scores, all_conf_scores)
         return scores
 
     def _encode_candidates_batched(
@@ -169,6 +172,7 @@ def _encode_candidates_batched(
         model: Encoder,
         all_mrr_scores,
         all_ap_scores,
+        all_recall_scores,
         all_conf_scores,
     ):
         all_docs = []
@@ -208,6 +212,7 @@ def _encode_candidates_batched(
                 is_relevant,
                 all_mrr_scores,
                 all_ap_scores,
+                all_recall_scores,
                 all_conf_scores,
                 model,
             )
@@ -217,6 +222,7 @@ def _encode_candidates_individual(
         model: Encoder,
         all_mrr_scores,
         all_ap_scores,
+        all_recall_scores,
         all_conf_scores,
     ):
         for instance in tqdm.tqdm(self.samples, desc="Samples"):
@@ -255,19 +261,22 @@ def _encode_candidates_individual(
                 is_relevant,
                 all_mrr_scores,
                 all_ap_scores,
+                all_recall_scores,
                 all_conf_scores,
                 model,
             )
 
-    def _collect_results(self, all_mrr_scores, all_ap_scores, all_conf_scores):
+    def _collect_results(self, all_mrr_scores, all_ap_scores, all_recall_scores, all_conf_scores):
         mean_ap = np.mean(all_ap_scores)
         mean_mrr = np.mean(all_mrr_scores)
+        mean_recall = np.mean(all_recall_scores)
 
         # Compute nAUCs
         naucs_map = self.nAUC_scores(all_conf_scores, all_ap_scores, "map")
         naucs_mrr = self.nAUC_scores(all_conf_scores, all_mrr_scores, "mrr")
+        naucs_recall = self.nAUC_scores(all_conf_scores, all_recall_scores, f"recall_at_{self.mrr_at_k}")
 
-        return {**{"map": mean_ap, "mrr": mean_mrr}, **naucs_map, **naucs_mrr}
+        return {**{"map": mean_ap, "mrr": mean_mrr, f"recall_at_{self.mrr_at_k}": mean_recall}, **naucs_map, **naucs_mrr, **naucs_recall}
 
     def _encode_candidates_miracl(
         self,
@@ -408,6 +417,7 @@ def _apply_sim_scores(
         is_relevant,
         all_mrr_scores,
         all_ap_scores,
+        all_recall_scores,
         all_conf_scores,
         model: Encoder,
     ):
@@ -417,6 +427,7 @@ def _apply_sim_scores(
 
         all_mrr_scores.append(scores["mrr"])
         all_ap_scores.append(scores["ap"])
+        all_recall_scores.append(scores["recall"])
         all_conf_scores.append(conf_scores)
 
     @staticmethod
@@ -483,11 +494,13 @@ def _compute_metrics_instance(
             scores:
                 - `mrr`: Mean Reciprocal Rank @ `self.mrr_at_k`
                 - `ap`: Average Precision
+                - `recall`: Recall @ `self.mrr_at_k`
         """
         pred_scores_argsort = torch.argsort(-sim_scores)  # Sort in decreasing order
         mrr = self.mrr_at_k_score(is_relevant, pred_scores_argsort, self.mrr_at_k)
         ap = self.ap_score(is_relevant, sim_scores.cpu().tolist())
-        return {"mrr": mrr, "ap": ap}
+        recall = self.recall_at_k_score(is_relevant, pred_scores_argsort, self.mrr_at_k)
+        return {"mrr": mrr, "ap": ap, "recall": recall}
 
     @staticmethod
     def conf_scores(sim_scores: torch.Tensor) -> dict[str, float]:
@@ -570,3 +583,29 @@ def ap_score(is_relevant, pred_scores):
         # ap = np.mean([np.mean(preds[: k + 1]) for k in range(len(preds)) if preds[k]])
         ap = average_precision_score(is_relevant, pred_scores)
         return ap
+
+    @staticmethod
+    def recall_at_k_score(
+        is_relevant: list[bool], pred_ranking: list[int], k: int
+    ) -> float:
+        """Computes Recall@k score
+
+        Args:
+            is_relevant: True if the document is relevant
+            pred_ranking: Indices of the documents sorted in decreasing order
+                of the similarity score
+            k: Top-k documents to consider
+
+        Returns:
+            The Recall@k score
+        """
+        total_relevant = sum(is_relevant)
+        if total_relevant == 0:
+            return 0.0
+        
+        relevant_retrieved = 0
+        for rank, index in enumerate(pred_ranking[:k]):
+            if is_relevant[index]:
+                relevant_retrieved += 1
+
+        return relevant_retrieved / total_relevant

mteb/tasks/Reranking/__init__.py

@@ -7,6 +7,9 @@
 from .eng.SciDocsReranking import *
 from .eng.StackOverflowDupQuestions import *
 from .eng.WebLINXCandidatesReranking import *
+from .eng.SWEbenchLiteReranking import *
+from .eng.SWEbenchVerifiedReranking import *
+from .eng.LocBenchReranking import *
 from .fra.AlloprofReranking import *
 from .fra.SyntecReranking import *
 from .jpn.JaCWIRReranking import *

mteb/tasks/Reranking/eng/LocBenchReranking.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from mteb.abstasks.TaskMetadata import TaskMetadata
+
+from ....abstasks.AbsTaskReranking import AbsTaskReranking
+
+
+class SWEbenchLiteReranking(AbsTaskReranking):
+    metadata = TaskMetadata(
+        name="LocBenchRR",
+        description="Software Issue Localization.",
+        reference="https://arxiv.org/abs/2503.09089",
+        dataset={
+            "path": "tarsur909/mteb-loc-bench-reranking",
+            "revision": "6741d68550b3793b45b18b6aaf981d00f33971cb",
+        },
+        type="Reranking",
+        category="p2p",
+        modalities=["text"],
+        eval_splits=["train"],
+        eval_langs=["eng-Latn", "python-Code"],
+        main_score="recall_at_10",
+        date=("2025-03-12", "2025-03-12"),  # arxiv v1 submission date
+        domains=["Programming", "Written"],
+        task_subtypes=["Code retrieval"],
+        license="mit",
+        annotations_creators="derived",
+        dialect=[],
+        sample_creation="found",
+        prompt="Instruct: Given a github issue, identify the code that needs to be changed to fix the issue. Query: ",
+        bibtex_citation=r"""
+@misc{chen2025locagentgraphguidedllmagents,
+      title={LocAgent: Graph-Guided LLM Agents for Code Localization}, 
+      author={Zhaoling Chen and Xiangru Tang and Gangda Deng and Fang Wu and Jialong Wu and Zhiwei Jiang and Viktor Prasanna and Arman Cohan and Xingyao Wang},
+      year={2025},
+      eprint={2503.09089},
+      archivePrefix={arXiv},
+      primaryClass={cs.SE},
+      url={https://arxiv.org/abs/2503.09089}, 
+}
+"""
+    )
+

mteb/tasks/Reranking/eng/SWEbenchLiteReranking.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from mteb.abstasks.TaskMetadata import TaskMetadata
+
+from ....abstasks.AbsTaskReranking import AbsTaskReranking
+
+
+class SWEbenchLiteReranking(AbsTaskReranking):
+    metadata = TaskMetadata(
+        name="SWEbenchLiteRR",
+        description="Software Issue Localization.",
+        reference="https://www.swebench.com/",
+        dataset={
+            "path": "tarsur909/mteb-swe-bench-lite-reranking",
+            "revision": "9020779825304b569312509a068219d1771bae7d",
+        },
+        type="Reranking",
+        category="p2p",
+        modalities=["text"],
+        eval_splits=["train"],
+        eval_langs=["eng-Latn", "python-Code"],
+        main_score="recall_at_10",
+        date=("2023-10-10", "2023-10-10"), 
+        domains=["Programming", "Written"],
+        task_subtypes=["Code retrieval"],
+        license="mit",
+        annotations_creators="derived",
+        dialect=[],
+        sample_creation="found",
+        prompt="Instruct: Given a github issue, identify the code that needs to be changed to fix the issue. Query: ",
+        bibtex_citation=r"""
+@misc{jimenez2024swebenchlanguagemodelsresolve,
+      title={SWE-bench: Can Language Models Resolve Real-World GitHub Issues?}, 
+      author={Carlos E. Jimenez and John Yang and Alexander Wettig and Shunyu Yao and Kexin Pei and Ofir Press and Karthik Narasimhan},
+      year={2024},
+      eprint={2310.06770},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2310.06770}, 
+}
+"""
+    )

mteb/tasks/Reranking/eng/SWEbenchVerifiedReranking.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from mteb.abstasks.TaskMetadata import TaskMetadata
+
+from ....abstasks.AbsTaskReranking import AbsTaskReranking
+
+
+class SWEbenchVerifiedReranking(AbsTaskReranking):
+    metadata = TaskMetadata(
+        name="SWEbenchVerifiedRR",
+        description="Software Issue Localization for SWE-bench Verified",
+        reference="https://openai.com/index/introducing-swe-bench-verified/",
+        dataset={
+            "path": "tarsur909/mteb-swe-bench-verified-reranking",
+            "revision": "796ae0b4b187e5c0533a12411dee0d8e34eaf0b5",
+        },
+        type="Reranking",
+        category="p2p",
+        modalities=["text"],
+        eval_splits=["train"],
+        eval_langs=["eng-Latn", "python-Code"],
+        main_score="recall_at_10",
+        date=("2024-08-13", "2024-08-13"),  # arxiv v1 submission date
+        domains=["Programming", "Written"],
+        task_subtypes=["Code retrieval"],
+        license="mit",
+        annotations_creators="derived",
+        dialect=[],
+        sample_creation="found",
+        prompt="Instruct: Given a github issue, identify the code that needs to be changed to fix the issue. Query: ",
+        bibtex_citation=r"""
+        @misc{openai2024swebenchverified,
+              title={Introducing swe-bench verified}, 
+              author={OpenAI},
+              year={2024},
+              url={https://openai.com/index/introducing-swe-bench-verified/}, 
+        }
+  """
+    )