feat: retrival detail add matched_count (#422)

* fix: update semantic_scholar cli example

* feat: retrival detail add matched_count

* fix comment

Author: e06084

dingo/exec/retrieval.py

@@ -105,6 +105,7 @@ def execute(self) -> SummaryModel:
                 continue
 
             try:
+                self._attach_relevant_docs(model, tasks)
                 results = mteb.evaluate(
                     model,
                     tasks=tasks,
@@ -161,6 +162,81 @@ def execute(self) -> SummaryModel:
         self.summary = summary
         return summary
 
+    @staticmethod
+    def _attach_relevant_docs(model: SearchClientModel, tasks: list[Any]) -> None:
+        """Load task qrels into the search adapter for detailed trace annotation."""
+        for task in tasks:
+            task.load_data()
+            if hasattr(task, "convert_v1_dataset_format_to_v2"):
+                task.convert_v1_dataset_format_to_v2(num_proc=None)
+
+            task_name = task.metadata.name
+            attached = False
+            for hf_subset, splits in getattr(task, "dataset", {}).items():
+                if not isinstance(splits, dict):
+                    continue
+                for hf_split, data_split in splits.items():
+                    if not isinstance(data_split, dict):
+                        continue
+                    relevant_docs = data_split.get("relevant_docs")
+                    if relevant_docs is None:
+                        continue
+                    model.set_relevant_docs(
+                        task_name,
+                        hf_split,
+                        hf_subset,
+                        relevant_docs,
+                    )
+                    attached = True
+
+            if attached:
+                continue
+
+            hf_subset = getattr(task, "hf_subset", "default")
+            relevant_docs_dict = getattr(task, "relevant_docs", {})
+            for (
+                hf_subset,
+                hf_split,
+                relevant_docs,
+            ) in RetrievalExecutor._iter_legacy_qrels(relevant_docs_dict, hf_subset):
+                model.set_relevant_docs(
+                    task_name,
+                    hf_split,
+                    hf_subset,
+                    relevant_docs,
+                )
+
+    @staticmethod
+    def _iter_legacy_qrels(
+        relevant_docs_dict: Any,
+        default_subset: str,
+    ):
+        """Yield qrels from older MTEB task.relevant_docs layouts."""
+        if not isinstance(relevant_docs_dict, dict):
+            return
+
+        for key, value in relevant_docs_dict.items():
+            if RetrievalExecutor._looks_like_qrels(value):
+                yield default_subset, key, value
+            elif isinstance(value, dict):
+                for split, qrels in value.items():
+                    if RetrievalExecutor._looks_like_qrels(qrels):
+                        yield key, split, qrels
+
+    @staticmethod
+    def _looks_like_qrels(value: Any) -> bool:
+        if not isinstance(value, dict):
+            return False
+        if not value:
+            return True
+        sample = next(iter(value.values()))
+        if isinstance(sample, dict):
+            if not sample:
+                return True
+            nested_sample = next(iter(sample.values()))
+            return not isinstance(nested_sample, (dict, list, tuple, set))
+        return isinstance(sample, (list, tuple, set))
+
     def _extract_metrics(self, model_result) -> dict[str, float]:
         """Extract metrics of interest from MTEB ModelResult."""
         metrics: dict[str, float] = {}

dingo/retrieval/backends/semantic_scholar.py

@@ -11,15 +11,7 @@
   --backend semantic_scholar \
   --tasks SciFact \
   --api-url https://api.semanticscholar.org \
-  --limit 100 \
-  --max-queries 5 \
-  --rate-limit 1.1 \
-  -o outputs/retrieval_eval
-
-S2_API_KEY=your_key dingo eval-retrieval \
-  --backend semantic_scholar \
-  --tasks SciFact \
-  --api-url https://api.semanticscholar.org \
+  --api-token YOUR_S2_API_KEY \
   --limit 100 \
   --max-queries 5 \
   -o outputs/retrieval_eval

dingo/retrieval/mteb_adapter.py

@@ -70,6 +70,9 @@ def __init__(
         self._corpus_size = 0
         self._collisions = 0
         self._search_traces: list[dict[str, Any]] = []
+        self._relevant_docs_by_context: dict[
+            tuple[str, str, str], dict[str, set[str]]
+        ] = {}
 
         safe_name = client.name.replace(" ", "-")
         self._mteb_model_meta = ModelMeta(
@@ -99,6 +102,31 @@ def __init__(
     def mteb_model_meta(self) -> ModelMeta:
         return self._mteb_model_meta
 
+    def set_relevant_docs(
+        self,
+        task_name: str,
+        hf_split: str,
+        hf_subset: str,
+        relevant_docs: dict[str, Any],
+    ) -> None:
+        """Attach qrels for richer debug traces.
+
+        MTEB's SearchProtocol does not pass qrels into ``search()``, but Dingo's
+        detailed traces are easier to inspect when mapped hits are annotated as
+        relevant or not.
+        """
+        normalized: dict[str, set[str]] = {}
+        for qid, docs in (relevant_docs or {}).items():
+            if isinstance(docs, dict):
+                normalized[str(qid)] = {
+                    str(doc_id) for doc_id, score in docs.items() if score
+                }
+            else:
+                normalized[str(qid)] = {str(doc_id) for doc_id in docs}
+        self._relevant_docs_by_context[
+            (task_name, hf_split, hf_subset)
+        ] = normalized
+
     def index(
         self,
         corpus: "CorpusDatasetType",
@@ -158,6 +186,9 @@ def search(
         errors = 0
         total_matched = 0
         query_details: list[dict[str, Any]] = []
+        relevant_docs_by_qid = self._relevant_docs_by_context.get(
+            (task_metadata.name, hf_split, hf_subset)
+        )
 
         def _process_query(idx_qid_text):
             idx, qid, q_text = idx_qid_text
@@ -168,13 +199,18 @@ def _process_query(idx_qid_text):
                     query=q_text, results=[], response_time_ms=0.0,
                     status_code=0, error=str(e),
                 )
-                return idx, qid, q_text, error_resp, None, None, None
+                return idx, qid, q_text, error_resp, None, None, None, None
 
             if response.error:
-                return idx, qid, q_text, response, None, None, None
+                return idx, qid, q_text, response, None, None, None, None
 
             doc_scores: dict[str, float] = {}
             top_api_results: list[dict[str, Any]] = []
+            relevant_doc_ids = (
+                relevant_docs_by_qid.get(str(qid))
+                if relevant_docs_by_qid is not None
+                else None
+            )
             mapping_stats: dict[str, int] = {
                 "doc_id_exact": 0,
                 "title_fallback": 0,
@@ -195,13 +231,33 @@ def _process_query(idx_qid_text):
                         "score": paper.score,
                         "resolved_corpus_id": resolved_id,
                         "mapping_source": src,
+                        "is_relevant": (
+                            bool(resolved_id and resolved_id in relevant_doc_ids)
+                            if relevant_doc_ids is not None
+                            else None
+                        ),
                     }
                 )
                 if not resolved_id or resolved_id in doc_scores:
                     continue
                 doc_scores[resolved_id] = 1.0 / (rank + 1)
 
-            return idx, qid, q_text, response, doc_scores, top_api_results, mapping_stats
+            relevant_matched_count = (
+                sum(1 for doc_id in doc_scores if doc_id in relevant_doc_ids)
+                if relevant_doc_ids is not None
+                else None
+            )
+
+            return (
+                idx,
+                qid,
+                q_text,
+                response,
+                doc_scores,
+                top_api_results,
+                mapping_stats,
+                relevant_matched_count,
+            )
 
         items = [(i, qid, qt) for i, (qid, qt) in enumerate(zip(query_ids, query_texts))]
 
@@ -215,7 +271,16 @@ def _process_query(idx_qid_text):
                 unit="query",
             )
             for future in concurrent.futures.as_completed(futures):
-                idx, qid, q_text, response, doc_scores, top_api_results, mapping_stats = future.result()
+                (
+                    idx,
+                    qid,
+                    q_text,
+                    response,
+                    doc_scores,
+                    top_api_results,
+                    mapping_stats,
+                    relevant_matched_count,
+                ) = future.result()
 
                 if doc_scores is None:
                     errors += 1
@@ -231,19 +296,44 @@ def _process_query(idx_qid_text):
                             "response_time_ms": response.response_time_ms,
                             "api_results_count": 0,
                             "matched_count": 0,
+                            "mapped_count": 0,
+                            "relevant_matched_count": 0,
+                            "relevant_total": 0,
                         }
                     )
                 else:
                     results[qid] = doc_scores
                     total_matched += len(doc_scores)
+                    matched_count = (
+                        relevant_matched_count
+                        if relevant_matched_count is not None
+                        else len(doc_scores)
+                    )
+                    relevant_doc_ids = (
+                        relevant_docs_by_qid.get(str(qid))
+                        if relevant_docs_by_qid is not None
+                        else None
+                    )
                     query_details.append(
                         {
                             "qid": qid,
                             "query_text": q_text,
                             "error": "",
                             "response_time_ms": response.response_time_ms,
                             "api_results_count": len(response.results),
-                            "matched_count": len(doc_scores),
+                            "matched_count": matched_count,
+                            "mapped_count": len(doc_scores),
+                            "relevant_matched_count": relevant_matched_count,
+                            "relevant_total": (
+                                len(relevant_doc_ids)
+                                if relevant_doc_ids is not None
+                                else None
+                            ),
+                            "gold_doc_ids": (
+                                sorted(relevant_doc_ids)
+                                if relevant_doc_ids is not None
+                                else None
+                            ),
                             "top_api_results": top_api_results,
                             "retrieved_doc_ids": list(doc_scores.keys()),
                             "mapping_stats": mapping_stats,

test/scripts/retrieval/test_mteb_adapter.py

@@ -0,0 +1,66 @@
+# """Unit tests for retrieval MTEB adapter traces."""
+
+# from types import SimpleNamespace
+
+# from dingo.retrieval.mteb_adapter import SearchClientModel
+# from dingo.retrieval.search_client import PaperResult, SearchClient, SearchResponse
+
+
+# class FakeSearchClient(SearchClient):
+#     name = "fake-search"
+
+#     def search(self, query: str, limit: int = 100) -> SearchResponse:
+#         return SearchResponse(
+#             query=query,
+#             results=[
+#                 PaperResult(
+#                     paper_id="external-1",
+#                     title="Mapped Non Relevant Paper",
+#                     score=0.9,
+#                 )
+#             ],
+#             response_time_ms=12.3,
+#             status_code=200,
+#         )
+
+
+# def test_trace_distinguishes_mapped_from_relevant_matches():
+#     model = SearchClientModel(FakeSearchClient(), search_limit=10)
+#     task_metadata = SimpleNamespace(name="FakeTask")
+#     corpus = [
+#         {"id": "doc-1", "title": "Mapped Non Relevant Paper", "text": ""},
+#         {"id": "doc-2", "title": "Gold Paper", "text": ""},
+#     ]
+#     queries = {"id": ["q1"], "text": ["test query"]}
+
+#     model.set_relevant_docs(
+#         "FakeTask",
+#         "test",
+#         "default",
+#         {"q1": {"doc-2": 1}},
+#     )
+#     model.index(
+#         corpus,
+#         task_metadata=task_metadata,
+#         hf_split="test",
+#         hf_subset="default",
+#         encode_kwargs={},
+#     )
+
+#     results = model.search(
+#         queries,
+#         task_metadata=task_metadata,
+#         hf_split="test",
+#         hf_subset="default",
+#         top_k=10,
+#         encode_kwargs={},
+#     )
+
+#     assert results == {"q1": {"doc-1": 1.0}}
+#     trace_query = model.get_search_traces()[0]["queries"][0]
+#     assert trace_query["mapped_count"] == 1
+#     assert trace_query["matched_count"] == 0
+#     assert trace_query["relevant_matched_count"] == 0
+#     assert trace_query["relevant_total"] == 1
+#     assert trace_query["gold_doc_ids"] == ["doc-2"]
+#     assert trace_query["top_api_results"][0]["is_relevant"] is False