addressing comments pt1

Author: andreatgretel

examples/example.py

@@ -1,15 +1,19 @@
+import json
+
 from data_designer.essentials import (
+    BuildStage,
     CategorySamplerParams,
     DataDesigner,
     DataDesignerConfigBuilder,
     InferenceParameters,
+    JsonlExportProcessorConfig,
     LLMTextColumnConfig,
     ModelConfig,
     PersonSamplerParams,
+    ProcessorType,
     SamplerColumnConfig,
     Score,
     SubcategorySamplerParams,
-    ToJsonlProcessorConfig,
 )
 
 # define model aliases
@@ -192,10 +196,12 @@
     ],
 }
 
+template_as_str = json.dumps(jsonl_entry_template)
 config_builder.add_processor(
-    ToJsonlProcessorConfig(
-        template=jsonl_entry_template,
-        folder_name="jsonl_files",
+    JsonlExportProcessorConfig(
+        processor_type=ProcessorType.JSONL_EXPORT,
+        build_stage=BuildStage.POST_BATCH,
+        template=template_as_str,
         fraction_per_file={
             "train.jsonl": 0.8,
             "validation.jsonl": 0.2,

src/data_designer/config/config_builder.py

@@ -314,13 +314,19 @@ def add_processor(
         Returns:
             The current Data Designer config builder instance.
         """
+        num_processors_of_type = len([p for p in self._processor_configs if p.processor_type == processor_type])
+
         if processor_config is None:
             if processor_type is None:
                 raise BuilderConfigurationError(
                     "🛑 You must provide either a 'processor_config' object or 'processor_type' "
                     "with additional keyword arguments."
                 )
-            processor_config = get_processor_config_from_kwargs(processor_type=processor_type, **kwargs)
+            processor_config = get_processor_config_from_kwargs(
+                processor_type=processor_type, name=f"{processor_type.value}-{num_processors_of_type + 1}", **kwargs
+            )
+        elif processor_config.name is None:
+            processor_config.name = f"{processor_config.processor_type}-{num_processors_of_type + 1}"
 
         # Checks elsewhere fail if DropColumnsProcessor drops a column but it is not marked for drop
         if processor_config.processor_type == ProcessorType.DROP_COLUMNS:

src/data_designer/config/processors.py

@@ -15,10 +15,13 @@
 
 class ProcessorType(str, Enum):
     DROP_COLUMNS = "drop_columns"
-    TO_JSONL = "to_jsonl"
+    JSONL_EXPORT = "jsonl_export"
 
 
 class ProcessorConfig(ConfigBase, ABC):
+    name: str | None = Field(
+        default=None, description="The name of the processor. If not provided, a default name will be generated."
+    )
     build_stage: BuildStage = Field(
         default=BuildStage.POST_BATCH,
         description=f"The stage at which the processor will run. Supported stages: {', '.join(SUPPORTED_STAGES)}",
@@ -36,23 +39,22 @@ def validate_build_stage(cls, v: BuildStage) -> BuildStage:
 def get_processor_config_from_kwargs(processor_type: ProcessorType, **kwargs) -> ProcessorConfig:
     if processor_type == ProcessorType.DROP_COLUMNS:
         return DropColumnsProcessorConfig(**kwargs)
-    elif processor_type == ProcessorType.TO_JSONL:
-        return ToJsonlProcessorConfig(**kwargs)
+    elif processor_type == ProcessorType.JSONL_EXPORT:
+        return JsonlExportProcessorConfig(**kwargs)
 
 
 class DropColumnsProcessorConfig(ProcessorConfig):
     column_names: list[str]
     processor_type: Literal[ProcessorType.DROP_COLUMNS] = ProcessorType.DROP_COLUMNS
 
 
-class ToJsonlProcessorConfig(ProcessorConfig):
-    template: dict = Field(..., description="The template to use for each entry in the dataset.")
-    folder_name: str = Field(..., description="Folder where JSONL files will be saved.")
+class JsonlExportProcessorConfig(ProcessorConfig):
+    template: str = Field(..., description="The template to use for each entry in the dataset, as a single string.")
     fraction_per_file: dict[str, float] = Field(
         default={"train.jsonl": 0.8, "validation.jsonl": 0.2},
         description="Fraction of the dataset to save in each file. The keys are the filenames and the values are the fractions.",
     )
-    processor_type: Literal[ProcessorType.TO_JSONL] = ProcessorType.TO_JSONL
+    processor_type: Literal[ProcessorType.JSONL_EXPORT] = ProcessorType.JSONL_EXPORT
 
     @field_validator("fraction_per_file")
     def validate_fraction_per_file(cls, v: dict[str, float]) -> dict[str, float]:

src/data_designer/engine/dataset_builders/artifact_storage.py

@@ -32,7 +32,7 @@ class ArtifactStorage(BaseModel):
     final_dataset_folder_name: str = "parquet-files"
     partial_results_folder_name: str = "tmp-partial-parquet-files"
     dropped_columns_folder_name: str = "dropped-columns-parquet-files"
-    outputs_folder_name: str = "outputs"
+    processors_outputs_folder_name: str = "processors-files"
 
     @property
     def artifact_path_exists(self) -> bool:
@@ -59,8 +59,8 @@ def partial_results_path(self) -> Path:
         return self.base_dataset_path / self.partial_results_folder_name
 
     @property
-    def outputs_path(self) -> Path:
-        return self.base_dataset_path / self.outputs_folder_name
+    def processors_outputs_path(self) -> Path:
+        return self.base_dataset_path / self.processors_outputs_folder_name
 
     @field_validator("artifact_path")
     def validate_artifact_path(cls, v: Union[Path, str]) -> Path:
@@ -183,10 +183,10 @@ def write_metadata(self, metadata: dict) -> Path:
             json.dump(metadata, file)
         return self.metadata_file_path
 
-    def move_to_outputs(self, from_path: Path, to_folder_name: str) -> Path:
-        self.mkdir_if_needed(self.outputs_path / to_folder_name)
-        shutil.move(from_path, self.outputs_path / to_folder_name / from_path.name)
-        return self.outputs_path / to_folder_name / from_path.name
+    def move_processor_output(self, from_path: Path, folder_name: str) -> Path:
+        self.mkdir_if_needed(self.processors_outputs_path / folder_name)
+        shutil.move(from_path, self.processors_outputs_path / folder_name / from_path.name)
+        return self.processors_outputs_path / folder_name / from_path.name
 
     def _get_stage_path(self, stage: BatchStage) -> Path:
         return getattr(self, resolve_string_enum(stage, BatchStage).value)

…engine/processing/processors/to_jsonl.py …ne/processing/processors/jsonl_export.pysrc/data_designer/engine/processing/processors/to_jsonl.py renamed to src/data_designer/engine/processing/processors/jsonl_export.py src/data_designer/engine/processing/processors/to_jsonl.py renamed to src/data_designer/engine/processing/processors/jsonl_export.py

@@ -1,14 +1,13 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import json
 import logging
 from pathlib import Path
 import tempfile
 
 import pandas as pd
 
-from data_designer.config.processors import ToJsonlProcessorConfig
+from data_designer.config.processors import JsonlExportProcessorConfig
 from data_designer.engine.configurable_task import ConfigurableTaskMetadata
 from data_designer.engine.processing.ginja.environment import WithJinja2UserTemplateRendering
 from data_designer.engine.processing.processors.base import Processor
@@ -17,19 +16,15 @@
 logger = logging.getLogger(__name__)
 
 
-class ToJsonlProcessor(WithJinja2UserTemplateRendering, Processor[ToJsonlProcessorConfig]):
+class JsonlExportProcessor(WithJinja2UserTemplateRendering, Processor[JsonlExportProcessorConfig]):
     @staticmethod
     def metadata() -> ConfigurableTaskMetadata:
         return ConfigurableTaskMetadata(
-            name="to_jsonl",
+            name="jsonl_export",
             description="Save formatted dataset as JSONL files.",
             required_resources=None,
         )
 
-    @property
-    def template_as_string(self) -> str:
-        return json.dumps(self.config.template)
-
     def _get_stop_index_per_file(self, dataset_size: int) -> dict[str, int]:
         """Helper function to get the end index for each file of the split."""
         stop_index_per_file = {}
@@ -42,7 +37,7 @@ def _get_stop_index_per_file(self, dataset_size: int) -> dict[str, int]:
         return stop_index_per_file
 
     def process(self, data: pd.DataFrame, *, current_batch_number: int | None = None) -> pd.DataFrame:
-        self.prepare_jinja2_template_renderer(self.template_as_string, data.columns.to_list())
+        self.prepare_jinja2_template_renderer(self.config.template, data.columns.to_list())
 
         stop_index_per_file = self._get_stop_index_per_file(len(data))
         with tempfile.TemporaryDirectory() as temp_dir:
@@ -60,6 +55,8 @@ def process(self, data: pd.DataFrame, *, current_batch_number: int | None = None
                             f.write("\n")
                 start_index = stop_index
 
-                self.artifact_storage.move_to_outputs(Path(temp_dir) / filename, self.config.folder_name)
+                self.artifact_storage.move_processor_output(
+                    from_path=Path(temp_dir) / filename, folder_name=self.config.name
+                )
 
         return data

src/data_designer/engine/processing/processors/registry.py

@@ -4,12 +4,12 @@
 from data_designer.config.base import ConfigBase
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
+    JsonlExportProcessorConfig,
     ProcessorType,
-    ToJsonlProcessorConfig,
 )
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
-from data_designer.engine.processing.processors.to_jsonl import ToJsonlProcessor
+from data_designer.engine.processing.processors.jsonl_export import JsonlExportProcessor
 from data_designer.engine.registry.base import TaskRegistry
 
 
@@ -19,5 +19,5 @@ class ProcessorRegistry(TaskRegistry[str, Processor, ConfigBase]): ...
 def create_default_processor_registry() -> ProcessorRegistry:
     registry = ProcessorRegistry()
     registry.register(ProcessorType.DROP_COLUMNS, DropColumnsProcessor, DropColumnsProcessorConfig, False)
-    registry.register(ProcessorType.TO_JSONL, ToJsonlProcessor, ToJsonlProcessorConfig, False)
+    registry.register(ProcessorType.JSONL_EXPORT, JsonlExportProcessor, JsonlExportProcessorConfig, False)
     return registry

src/data_designer/essentials/__init__.py

@@ -34,7 +34,7 @@
     UniformDistribution,
     UniformDistributionParams,
 )
-from ..config.processors import DropColumnsProcessorConfig, ProcessorType, ToJsonlProcessorConfig
+from ..config.processors import DropColumnsProcessorConfig, JsonlExportProcessorConfig, ProcessorType
 from ..config.sampler_constraints import ColumnInequalityConstraint, ScalarInequalityConstraint
 from ..config.sampler_params import (
     BernoulliMixtureSamplerParams,
@@ -98,6 +98,7 @@
     "ImageFormat",
     "InferenceParameters",
     "JudgeScoreProfilerConfig",
+    "JsonlExportProcessorConfig",
     "LLMCodeColumnConfig",
     "LLMJudgeColumnConfig",
     "LLMStructuredColumnConfig",
@@ -124,7 +125,6 @@
     "SeedDatasetColumnConfig",
     "SubcategorySamplerParams",
     "TimeDeltaSamplerParams",
-    "ToJsonlProcessorConfig",
     "UniformDistribution",
     "UniformDistributionParams",
     "UniformSamplerParams",

tests/config/test_processors.py

@@ -7,9 +7,9 @@
 from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
+    JsonlExportProcessorConfig,
     ProcessorConfig,
     ProcessorType,
-    ToJsonlProcessorConfig,
     get_processor_config_from_kwargs,
 )
 
@@ -47,66 +47,59 @@ def test_drop_columns_processor_config_serialization():
     assert config_restored.column_names == config.column_names
 
 
-def test_to_jsonl_processor_config_creation():
-    config = ToJsonlProcessorConfig(
+def test_jsonl_export_processor_config_creation():
+    config = JsonlExportProcessorConfig(
         build_stage=BuildStage.POST_BATCH,
-        template={"text": "{{ col1 }}"},
-        folder_name="jsonl_output",
+        template='{"text": "{{ col1 }}"}',
         fraction_per_file={"train.jsonl": 0.8, "validation.jsonl": 0.2},
     )
 
     assert config.build_stage == BuildStage.POST_BATCH
-    assert config.template == {"text": "{{ col1 }}"}
-    assert config.folder_name == "jsonl_output"
+    assert config.template == '{"text": "{{ col1 }}"}'
     assert config.fraction_per_file == {"train.jsonl": 0.8, "validation.jsonl": 0.2}
-    assert config.processor_type == ProcessorType.TO_JSONL
+    assert config.processor_type == ProcessorType.JSONL_EXPORT
     assert isinstance(config, ProcessorConfig)
 
 
-def test_to_jsonl_processor_config_validation():
+def test_jsonl_export_processor_config_validation():
     # Test unsupported stage raises error
     with pytest.raises(ValidationError, match="Invalid dataset builder stage"):
-        ToJsonlProcessorConfig(
+        JsonlExportProcessorConfig(
             build_stage=BuildStage.PRE_BATCH,
-            template={"text": "{{ col1 }}"},
-            folder_name="jsonl_output",
+            template='{"text": "{{ col1 }}"}',
             fraction_per_file={"train.jsonl": 0.8, "validation.jsonl": 0.2},
         )
 
     # Test missing required field raises error
     with pytest.raises(ValidationError, match="Field required"):
-        ToJsonlProcessorConfig(build_stage=BuildStage.POST_BATCH, template={"text": "{{ col1 }}"})
+        JsonlExportProcessorConfig(build_stage=BuildStage.POST_BATCH)
 
     # Test invalid fraction per file raises error
     with pytest.raises(ValidationError, match="The fractions must sum to 1."):
-        ToJsonlProcessorConfig(
+        JsonlExportProcessorConfig(
             build_stage=BuildStage.POST_BATCH,
-            template={"text": "{{ col1 }}"},
-            folder_name="jsonl_output",
+            template='{"text": "{{ col1 }}"}',
             fraction_per_file={"train.jsonl": 0.8, "validation.jsonl": 0.3},
         )
 
 
-def test_to_jsonl_processor_config_serialization():
-    config = ToJsonlProcessorConfig(
+def test_jsonl_export_processor_config_serialization():
+    config = JsonlExportProcessorConfig(
         build_stage=BuildStage.POST_BATCH,
-        template={"text": "{{ col1 }}"},
-        folder_name="jsonl_output",
+        template='{"text": "{{ col1 }}"}',
         fraction_per_file={"train.jsonl": 0.8, "validation.jsonl": 0.2},
     )
 
     # Serialize to dict
     config_dict = config.model_dump()
     assert config_dict["build_stage"] == "post_batch"
-    assert config_dict["template"] == {"text": "{{ col1 }}"}
-    assert config_dict["folder_name"] == "jsonl_output"
+    assert config_dict["template"] == '{"text": "{{ col1 }}"}'
     assert config_dict["fraction_per_file"] == {"train.jsonl": 0.8, "validation.jsonl": 0.2}
 
     # Deserialize from dict
-    config_restored = ToJsonlProcessorConfig.model_validate(config_dict)
+    config_restored = JsonlExportProcessorConfig.model_validate(config_dict)
     assert config_restored.build_stage == config.build_stage
     assert config_restored.template == config.template
-    assert config_restored.folder_name == config.folder_name
     assert config_restored.fraction_per_file == config.fraction_per_file
 
 
@@ -119,18 +112,16 @@ def test_get_processor_config_from_kwargs():
     assert config_drop_columns.column_names == ["col1"]
     assert config_drop_columns.processor_type == ProcessorType.DROP_COLUMNS
 
-    config_to_jsonl = get_processor_config_from_kwargs(
-        ProcessorType.TO_JSONL,
+    config_jsonl_export = get_processor_config_from_kwargs(
+        ProcessorType.JSONL_EXPORT,
         build_stage=BuildStage.POST_BATCH,
-        template={"text": "{{ col1 }}"},
-        folder_name="jsonl_output",
+        template='{"text": "{{ col1 }}"}',
         fraction_per_file={"train.jsonl": 0.8, "validation.jsonl": 0.2},
     )
-    assert isinstance(config_to_jsonl, ToJsonlProcessorConfig)
-    assert config_to_jsonl.template == {"text": "{{ col1 }}"}
-    assert config_to_jsonl.folder_name == "jsonl_output"
-    assert config_to_jsonl.fraction_per_file == {"train.jsonl": 0.8, "validation.jsonl": 0.2}
-    assert config_to_jsonl.processor_type == ProcessorType.TO_JSONL
+    assert isinstance(config_jsonl_export, JsonlExportProcessorConfig)
+    assert config_jsonl_export.template == '{"text": "{{ col1 }}"}'
+    assert config_jsonl_export.fraction_per_file == {"train.jsonl": 0.8, "validation.jsonl": 0.2}
+    assert config_jsonl_export.processor_type == ProcessorType.JSONL_EXPORT
 
     # Test with unknown processor type returns None
     from enum import Enum