address review comments

Signed-off-by: Steve Han <sthan@nvidia.com>

Author: shan-nvidia

.gitignore

@@ -35,3 +35,6 @@ htmlcov/
 
 # Distribution
 *.tar.gz
+
+# CI artifacts
+*artifacts/

plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/chunking.py

@@ -21,6 +21,7 @@
 from typing import Literal
 
 import nltk
+import yaml
 from nltk.tokenize import sent_tokenize
 
 logger = logging.getLogger(__name__)
@@ -40,8 +41,6 @@ def load_multi_doc_manifest(manifest_path: Path | None) -> list[list[str]]:
     Returns:
         List of bundles, each a list of file-path strings.
     """
-    import yaml
-
     if not manifest_path:
         return []
 
@@ -160,19 +159,19 @@ def build_bundles(
     return [b for b in bundles if b]
 
 
-def group_chunks_by_doc(chunks: list[dict]) -> dict[str, list[tuple[int, dict]]]:
+def group_chunks_by_doc(chunks: list[dict]) -> dict[str, list[dict]]:
     """Group chunks by their ``doc_id`` field."""
-    grouped: dict[str, list[tuple[int, dict]]] = defaultdict(list)
-    for idx, chunk in enumerate(chunks):
+    grouped: dict[str, list[dict]] = defaultdict(list)
+    for chunk in chunks:
         doc_id = chunk.get("doc_id", "default")
-        grouped[doc_id].append((idx, chunk))
+        grouped[doc_id].append(chunk)
     return dict(grouped)
 
 
-def format_section_chunks(indexed_chunks: list[tuple[int, dict]], section_number: int) -> str:
-    """Render a list of indexed chunks into a section string."""
+def format_section_chunks(section_chunks: list[dict], section_number: int) -> str:
+    """Render a list of chunks into a section string."""
     section_lines: list[str] = []
-    for _, chunk in indexed_chunks:
+    for chunk in section_chunks:
         text = chunk.get("text", "").strip()
         if not text:
             continue
@@ -203,8 +202,7 @@ def chunks_to_sections_sequential(chunks: list[dict], num_sections: int = 1) ->
     for i in range(num_sections):
         start_idx = i * section_size
         end_idx = (i + 1) * section_size if i < num_sections - 1 else total
-        indexed_chunks = [(j, chunks[j]) for j in range(start_idx, end_idx)]
-        section_text = format_section_chunks(indexed_chunks, i + 1)
+        section_text = format_section_chunks(chunks[start_idx:end_idx], i + 1)
         if section_text:
             formatted_sections.append(section_text)
 
@@ -222,9 +220,9 @@ def chunks_to_sections_doc_balanced(chunks: list[dict], num_sections: int = 1) -
 
     chunk_sizes = {doc_id: max(1, math.ceil(len(entries) / num_sections)) for doc_id, entries in grouped.items()}
 
-    sections: list[list[tuple[int, dict]]] = []
+    sections: list[list[dict]] = []
     for part_idx in range(num_sections):
-        part_entries: list[tuple[int, dict]] = []
+        part_entries: list[dict] = []
         for doc_id, entries in grouped.items():
             chunk_size = chunk_sizes[doc_id]
             start = part_idx * chunk_size
@@ -235,8 +233,8 @@ def chunks_to_sections_doc_balanced(chunks: list[dict], num_sections: int = 1) -
             sections.append(part_entries)
 
     formatted_sections: list[str] = []
-    for i, indexed_chunks in enumerate(sections):
-        section_text = format_section_chunks(indexed_chunks, i + 1)
+    for i, section_chunks in enumerate(sections):
+        section_text = format_section_chunks(section_chunks, i + 1)
         if section_text:
             formatted_sections.append(section_text)
 
@@ -254,7 +252,7 @@ def chunks_to_sections_interleaved(chunks: list[dict], num_sections: int = 1) ->
 
     doc_iterators = {doc_id: deque(entries) for doc_id, entries in grouped.items()}
     doc_order = list(grouped.keys())
-    interleaved: list[tuple[int, dict]] = []
+    interleaved: list[dict] = []
 
     while True:
         added = False
@@ -276,8 +274,7 @@ def chunks_to_sections_interleaved(chunks: list[dict], num_sections: int = 1) ->
     for i in range(num_sections):
         start_idx = i * section_size
         end_idx = (i + 1) * section_size if i < num_sections - 1 else total
-        indexed_chunks = interleaved[start_idx:end_idx]
-        section_text = format_section_chunks(indexed_chunks, i + 1)
+        section_text = format_section_chunks(interleaved[start_idx:end_idx], i + 1)
         if section_text:
             formatted_sections.append(section_text)
 

plugins/data-designer-retrieval-sdg/src/data_designer_retrieval_sdg/dedup.py

@@ -15,12 +15,12 @@
 from typing import Any
 
 import numpy as np
+from data_designer.config.errors import BuilderConfigurationError
 from data_designer.config.models import GenerationType
 from data_designer.engine.column_generators.generators.base import (
     ColumnGeneratorWithModelRegistry,
     GenerationStrategy,
 )
-from data_designer.engine.dataset_builders.errors import DatasetGenerationError
 from data_designer.engine.models.facade import ModelFacade
 
 from data_designer_retrieval_sdg.config import EmbeddingDedupColumnConfig
@@ -65,14 +65,14 @@ def _validate(self) -> None:
         from the facade or a 400 from the embeddings endpoint.
 
         Raises:
-            DatasetGenerationError: When ``self.config.model_alias`` resolves
+            BuilderConfigurationError: When ``self.config.model_alias`` resolves
                 to a :class:`ModelConfig` whose inference parameters are not
                 ``EmbeddingInferenceParams``.
         """
         super()._validate()
         model_config = self.get_model_config(model_alias=self.config.model_alias)
         if model_config.generation_type != GenerationType.EMBEDDING:
-            raise DatasetGenerationError(
+            raise BuilderConfigurationError(
                 f"EmbeddingDedupColumnGenerator requires an embedding model, "
                 f"but model alias {self.config.model_alias!r} resolves to a "
                 f"{model_config.generation_type.value!r} model. Configure a "

plugins/data-designer-retrieval-sdg/tests/test_dedup.py

@@ -7,12 +7,12 @@
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+from data_designer.config.errors import BuilderConfigurationError
 from data_designer.config.models import (
     ChatCompletionInferenceParams,
     EmbeddingInferenceParams,
     ModelConfig,
 )
-from data_designer.engine.dataset_builders.errors import DatasetGenerationError
 
 from data_designer_retrieval_sdg.config import EmbeddingDedupColumnConfig
 from data_designer_retrieval_sdg.dedup import EmbeddingDedupColumnGenerator
@@ -189,7 +189,7 @@ def test_validate_rejects_chat_model() -> None:
         model="some/chat-model",
         inference_parameters=ChatCompletionInferenceParams(),
     )
-    with pytest.raises(DatasetGenerationError, match="embed"):
+    with pytest.raises(BuilderConfigurationError, match="embed"):
         gen._validate()