diff --git a/scripts/materialize_searchqa.py b/scripts/materialize_searchqa.py index 30838ac0..be390864 100644 --- a/scripts/materialize_searchqa.py +++ b/scripts/materialize_searchqa.py @@ -44,6 +44,27 @@ def load_manifest_ids(manifest_dir: Path) -> dict[str, list[str]]: return split_ids +def _reject_duplicate_manifest_ids(split_ids: Mapping[str, Iterable[str]]) -> None: + seen: dict[str, str] = {} + duplicates: dict[str, list[str]] = {} + for split, ids in split_ids.items(): + for item_id in ids: + if item_id in seen: + duplicates.setdefault(item_id, [seen[item_id]]).append(split) + else: + seen[item_id] = split + + if duplicates: + preview = ", ".join( + f"{item_id} ({'/'.join(splits)})" + for item_id, splits in sorted(duplicates.items())[:5] + ) + raise ValueError( + "SearchQA split manifest contains duplicate IDs across splits. " + f"First IDs: {preview}" + ) + + def _iter_dataset_rows(dataset: Mapping[str, Iterable[dict]]) -> Iterable[dict]: for source_split in dataset.values(): yield from source_split @@ -78,6 +99,7 @@ def materialize_searchqa_splits( manifest_dir = manifest_dir.resolve() output_dir = output_dir.resolve() split_ids = load_manifest_ids(manifest_dir) + _reject_duplicate_manifest_ids(split_ids) wanted_ids = {item_id for ids in split_ids.values() for item_id in ids} selected: dict[str, dict] = {} diff --git a/tests/test_materialize_searchqa.py b/tests/test_materialize_searchqa.py index bbfb2a86..3228ed83 100644 --- a/tests/test_materialize_searchqa.py +++ b/tests/test_materialize_searchqa.py @@ -64,3 +64,16 @@ def test_materialize_searchqa_splits_fails_on_missing_manifest_id(tmp_path): {"train": [_row("a")]}, dataset_name="example/searchqa", ) + + +def test_materialize_searchqa_splits_rejects_duplicate_manifest_ids(tmp_path): + manifest_dir = tmp_path / "manifest" + _write_manifest(manifest_dir, {"train": ["a"], "val": ["a"], "test": []}) + + with pytest.raises(ValueError, match="duplicate IDs"): + materialize_searchqa_splits( + manifest_dir, + tmp_path / "out", + {"train": [_row("a")]}, + dataset_name="example/searchqa", + )