From dde0e600f3f68f344d0c0f1063e897e6e998c16a Mon Sep 17 00:00:00 2001 From: Andrew White Date: Fri, 24 Jul 2026 08:41:44 -0500 Subject: [PATCH 1/2] fix: read qrels corpus-id column without trailing space The qrels TSV is written with header 'corpus-id' (no trailing space), but load_positive_docs_with_modality looked up row['corpus-id '], which raised KeyError when the function was used. Signed-off-by: Andrew White --- .../postprocess.py | 4 +- .../tests/test_postprocess.py | 38 ++++++++++++++++++- 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/postprocess.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/postprocess.py index 1bbff7f..7b74070 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/postprocess.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/postprocess.py @@ -202,7 +202,7 @@ def load_positive_docs_with_modality( doc_to_modality: dict[str, set[str]] = defaultdict(set) for _, row in qrels_df.iterrows(): query_id = row["query-id"] - corpus_id = row["corpus-id "] # trailing space in column name + corpus_id = row["corpus-id"] if query_id in query_to_modality: doc_to_modality[corpus_id].add(query_to_modality[query_id]) @@ -212,7 +212,7 @@ def load_positive_docs_with_modality( doc_to_modality_final[doc_id] = next(iter(modalities)) else: modality_counts: dict[str, int] = defaultdict(int) - for _, r in qrels_df[qrels_df["corpus-id "] == doc_id].iterrows(): + for _, r in qrels_df[qrels_df["corpus-id"] == doc_id].iterrows(): qid = r["query-id"] if qid in query_to_modality: modality_counts[query_to_modality[qid]] += 1 diff --git a/plugins/data-designer-retrieval-sdg/tests/test_postprocess.py b/plugins/data-designer-retrieval-sdg/tests/test_postprocess.py index f73b24d..880c63d 100644 --- a/plugins/data-designer-retrieval-sdg/tests/test_postprocess.py +++ b/plugins/data-designer-retrieval-sdg/tests/test_postprocess.py @@ -1,9 +1,16 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import json +from pathlib import Path + import pandas as pd -from data_designer_retrieval_sdg.postprocess import filter_qa_pairs_by_quality, postprocess_retriever_data +from data_designer_retrieval_sdg.postprocess import ( + filter_qa_pairs_by_quality, + load_positive_docs_with_modality, + postprocess_retriever_data, +) # --------------------------------------------------------------------------- # postprocess_retriever_data @@ -85,3 +92,32 @@ def test_filter_skips_mismatched() -> None: filtered_df, skipped = filter_qa_pairs_by_quality(df, quality_threshold=5.0) assert len(filtered_df) == 0 assert len(skipped) == 1 + + +# --------------------------------------------------------------------------- +# load_positive_docs_with_modality +# --------------------------------------------------------------------------- + + +def test_load_positive_docs_with_modality_reads_qrels_without_trailing_space(tmp_path: Path) -> None: + """Regression test for reading 'corpus-id' without a trailing space. + + The qrels TSV is written with header ``corpus-id`` (no trailing space), but + the loader previously looked up ``row["corpus-id "]``. This test would have + raised a KeyError before the fix. + """ + qrels_path = tmp_path / "test.tsv" + qrels_path.write_text("query-id\tcorpus-id\tscore\nq1\td1\t1\n") + + corpus_path = tmp_path / "corpus.jsonl" + corpus_path.write_text(json.dumps({"_id": "d1", "text": "hello world", "title": "T"}) + "\n") + + split_path = tmp_path / "split.json" + split_path.write_text(json.dumps({"text": ["q1"]})) + + docs_df, modality_map = load_positive_docs_with_modality(qrels_path, corpus_path, split_path) + + assert len(docs_df) == 1 + assert docs_df.iloc[0]["doc_id"] == "d1" + assert docs_df.iloc[0]["modality"] == "text" + assert modality_map == {"d1": "text"} From f099514347648331a1c8e5e07b2d25b94d76f03f Mon Sep 17 00:00:00 2001 From: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:45:53 +0100 Subject: [PATCH 2/2] fix: preserve legacy qrels corpus ID headers Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com> --- .../postprocess.py | 2 ++ .../tests/test_postprocess.py | 34 +++++++++++++------ 2 files changed, 26 insertions(+), 10 deletions(-) diff --git a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/postprocess.py b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/postprocess.py index 7b74070..e31085f 100644 --- a/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/postprocess.py +++ b/plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/postprocess.py @@ -190,6 +190,8 @@ def load_positive_docs_with_modality( Tuple of ``(positive_docs_df, doc_to_modality_final)``. """ qrels_df = pd.read_csv(test_tsv_path, sep="\t") + if "corpus-id" not in qrels_df.columns and "corpus-id " in qrels_df.columns: + qrels_df = qrels_df.rename(columns={"corpus-id ": "corpus-id"}) with open(split_json_path, encoding="utf-8") as f: splits = json.load(f) diff --git a/plugins/data-designer-retrieval-sdg/tests/test_postprocess.py b/plugins/data-designer-retrieval-sdg/tests/test_postprocess.py index 880c63d..8cc9a25 100644 --- a/plugins/data-designer-retrieval-sdg/tests/test_postprocess.py +++ b/plugins/data-designer-retrieval-sdg/tests/test_postprocess.py @@ -5,6 +5,7 @@ from pathlib import Path import pandas as pd +import pytest from data_designer_retrieval_sdg.postprocess import ( filter_qa_pairs_by_quality, @@ -99,21 +100,17 @@ def test_filter_skips_mismatched() -> None: # --------------------------------------------------------------------------- -def test_load_positive_docs_with_modality_reads_qrels_without_trailing_space(tmp_path: Path) -> None: - """Regression test for reading 'corpus-id' without a trailing space. - - The qrels TSV is written with header ``corpus-id`` (no trailing space), but - the loader previously looked up ``row["corpus-id "]``. This test would have - raised a KeyError before the fix. - """ +@pytest.mark.parametrize("corpus_id_header", ["corpus-id", "corpus-id "]) +def test_load_positive_docs_with_modality_accepts_corpus_id_headers(tmp_path: Path, corpus_id_header: str) -> None: + """Canonical and legacy corpus ID headers produce the same result.""" qrels_path = tmp_path / "test.tsv" - qrels_path.write_text("query-id\tcorpus-id\tscore\nq1\td1\t1\n") + qrels_path.write_text(f"query-id\t{corpus_id_header}\tscore\nq1\td1\t1\n", encoding="utf-8") corpus_path = tmp_path / "corpus.jsonl" - corpus_path.write_text(json.dumps({"_id": "d1", "text": "hello world", "title": "T"}) + "\n") + corpus_path.write_text(json.dumps({"_id": "d1", "text": "hello world", "title": "T"}) + "\n", encoding="utf-8") split_path = tmp_path / "split.json" - split_path.write_text(json.dumps({"text": ["q1"]})) + split_path.write_text(json.dumps({"text": ["q1"]}), encoding="utf-8") docs_df, modality_map = load_positive_docs_with_modality(qrels_path, corpus_path, split_path) @@ -121,3 +118,20 @@ def test_load_positive_docs_with_modality_reads_qrels_without_trailing_space(tmp assert docs_df.iloc[0]["doc_id"] == "d1" assert docs_df.iloc[0]["modality"] == "text" assert modality_map == {"d1": "text"} + + +def test_load_positive_docs_with_modality_uses_most_common_modality(tmp_path: Path) -> None: + """A document shared across modalities uses its most frequent modality.""" + qrels_path = tmp_path / "test.tsv" + qrels_path.write_text("query-id\tcorpus-id\tscore\nq1\td1\t1\nq2\td1\t1\nq3\td1\t1\n", encoding="utf-8") + + corpus_path = tmp_path / "corpus.jsonl" + corpus_path.write_text(json.dumps({"_id": "d1", "text": "hello world", "title": "T"}) + "\n", encoding="utf-8") + + split_path = tmp_path / "split.json" + split_path.write_text(json.dumps({"text": ["q1", "q2"], "image": ["q3"]}), encoding="utf-8") + + docs_df, modality_map = load_positive_docs_with_modality(qrels_path, corpus_path, split_path) + + assert docs_df.iloc[0]["modality"] == "text" + assert modality_map == {"d1": "text"}