advancing in the discussed direction

Author: andreatgretel

examples/.gitignore

@@ -1 +1,2 @@
 artifacts
+processor_outputs

examples/example.py

@@ -1,37 +1,20 @@
 import json
 
 from data_designer.essentials import (
-    BuildStage,
     CategorySamplerParams,
     DataDesigner,
     DataDesignerConfigBuilder,
-    InferenceParameters,
-    JsonlExportProcessorConfig,
     LLMTextColumnConfig,
-    ModelConfig,
+    OutputFormatProcessorConfig,
     PersonSamplerParams,
-    ProcessorType,
     SamplerColumnConfig,
     Score,
     SubcategorySamplerParams,
 )
 
 # define model aliases
-model_alias_generator = "content_generator"
-model_configs = [
-    ModelConfig(
-        alias=model_alias_generator,
-        provider="nvidia",
-        model="deepseek-ai/deepseek-r1-distill-qwen-14b",
-        inference_parameters=InferenceParameters(
-            max_tokens=8000,
-            temperature=0.7,
-            top_p=0.95,
-        ),
-    )
-]
-
-config_builder = DataDesignerConfigBuilder(model_configs=model_configs)
+model_alias_generator = "nvidia-text"
+config_builder = DataDesignerConfigBuilder()
 
 # ESI levels
 ESI_LEVELS = [
@@ -198,14 +181,9 @@
 
 template_as_str = json.dumps(jsonl_entry_template)
 config_builder.add_processor(
-    JsonlExportProcessorConfig(
-        processor_type=ProcessorType.JSONL_EXPORT,
-        build_stage=BuildStage.POST_BATCH,
+    OutputFormatProcessorConfig(
+        name="jsonl_output",
         template=template_as_str,
-        fraction_per_file={
-            "train.jsonl": 0.8,
-            "validation.jsonl": 0.2,
-        },
     )
 )
 
@@ -214,4 +192,5 @@
 )
 preview = dd.preview(config_builder, num_records=10)
 
-dd.create(config_builder, num_records=20)
+results = dd.create(config_builder, num_records=20)
+results.write_processor_outputs_to_disk("./processor_outputs", "jsonl")

src/data_designer/config/processors.py

@@ -15,7 +15,7 @@
 
 class ProcessorType(str, Enum):
     DROP_COLUMNS = "drop_columns"
-    JSONL_EXPORT = "jsonl_export"
+    OUTPUT_FORMAT = "output_format"
 
 
 class ProcessorConfig(ConfigBase, ABC):
@@ -39,25 +39,15 @@ 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.JSONL_EXPORT:
-        return JsonlExportProcessorConfig(**kwargs)
+    elif processor_type == ProcessorType.OUTPUT_FORMAT:
+        return OutputFormatProcessorConfig(**kwargs)
 
 
 class DropColumnsProcessorConfig(ProcessorConfig):
     column_names: list[str]
     processor_type: Literal[ProcessorType.DROP_COLUMNS] = ProcessorType.DROP_COLUMNS
 
 
-class JsonlExportProcessorConfig(ProcessorConfig):
+class OutputFormatProcessorConfig(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.JSONL_EXPORT] = ProcessorType.JSONL_EXPORT
-
-    @field_validator("fraction_per_file")
-    def validate_fraction_per_file(cls, v: dict[str, float]) -> dict[str, float]:
-        if sum(v.values()) != 1:
-            raise ValueError("The fractions must sum to 1.")
-        return v
+    processor_type: Literal[ProcessorType.OUTPUT_FORMAT] = ProcessorType.OUTPUT_FORMAT

src/data_designer/engine/dataset_builders/artifact_storage.py

@@ -23,6 +23,7 @@ class BatchStage(StrEnum):
     PARTIAL_RESULT = "partial_results_path"
     FINAL_RESULT = "final_dataset_path"
     DROPPED_COLUMNS = "dropped_columns_dataset_path"
+    PROCESSORS_OUTPUTS = "processors_outputs_path"
 
 
 class ArtifactStorage(BaseModel):
@@ -75,6 +76,7 @@ def validate_folder_names(self):
             self.final_dataset_folder_name,
             self.partial_results_folder_name,
             self.dropped_columns_folder_name,
+            self.processors_outputs_folder_name,
         ]
 
         for name in folder_names:
@@ -160,19 +162,21 @@ def write_batch_to_parquet_file(
         batch_number: int,
         dataframe: pd.DataFrame,
         batch_stage: BatchStage,
+        subfolder: str = "",
     ) -> Path:
         file_path = self.create_batch_file_path(batch_number, batch_stage=batch_stage)
-        self.write_parquet_file(file_path.name, dataframe, batch_stage)
+        self.write_parquet_file(file_path.name, dataframe, batch_stage, subfolder=subfolder)
         return file_path
 
     def write_parquet_file(
         self,
         parquet_file_name: str,
         dataframe: pd.DataFrame,
         batch_stage: BatchStage,
+        subfolder: str = "",
     ) -> Path:
-        self.mkdir_if_needed(self._get_stage_path(batch_stage))
-        file_path = self._get_stage_path(batch_stage) / parquet_file_name
+        self.mkdir_if_needed(self._get_stage_path(batch_stage) / subfolder)
+        file_path = self._get_stage_path(batch_stage) / subfolder / parquet_file_name
         dataframe.to_parquet(file_path, index=False)
         return file_path
 
@@ -182,10 +186,5 @@ def write_metadata(self, metadata: dict) -> Path:
             json.dump(metadata, file)
         return self.metadata_file_path
 
-    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)

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

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import pandas as pd
+
+from data_designer.config.processors import OutputFormatProcessorConfig
+from data_designer.engine.configurable_task import ConfigurableTaskMetadata
+from data_designer.engine.dataset_builders.artifact_storage import BatchStage
+from data_designer.engine.processing.ginja.environment import WithJinja2UserTemplateRendering
+from data_designer.engine.processing.processors.base import Processor
+from data_designer.engine.processing.utils import deserialize_json_values
+
+logger = logging.getLogger(__name__)
+
+
+class OutputFormatProcessor(WithJinja2UserTemplateRendering, Processor[OutputFormatProcessorConfig]):
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="output_format",
+            description="Format the dataset using a Jinja2 template.",
+            required_resources=None,
+        )
+
+    def process(self, data: pd.DataFrame, *, current_batch_number: int | None = None) -> pd.DataFrame:
+        self.prepare_jinja2_template_renderer(self.config.template, data.columns.to_list())
+        formatted_records = [self.render_template(deserialize_json_values(record)) for record in data.to_dict(orient="records")]
+        formatted_data = pd.DataFrame(formatted_records, columns=["formatted_output"])
+        if current_batch_number is not None:
+            self.artifact_storage.write_batch_to_parquet_file(
+                batch_number=current_batch_number,
+                dataframe=formatted_data,
+                batch_stage=BatchStage.PROCESSORS_OUTPUTS,
+                subfolder=self.config.name,
+            )
+        else:
+            logger.warning("⚠️ Cannot write processor outputs to disk in preview mode.")
+
+        return data

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

@@ -4,12 +4,12 @@
 from data_designer.config.base import ConfigBase
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
-    JsonlExportProcessorConfig,
+    OutputFormatProcessorConfig,
     ProcessorType,
 )
 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.jsonl_export import JsonlExportProcessor
+from data_designer.engine.processing.processors.output_format import OutputFormatProcessor
 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.JSONL_EXPORT, JsonlExportProcessor, JsonlExportProcessorConfig, False)
+    registry.register(ProcessorType.OUTPUT_FORMAT, OutputFormatProcessor, OutputFormatProcessorConfig, False)
     return registry

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

@@ -34,7 +34,7 @@
     UniformDistribution,
     UniformDistributionParams,
 )
-from ..config.processors import DropColumnsProcessorConfig, JsonlExportProcessorConfig, ProcessorType
+from ..config.processors import DropColumnsProcessorConfig, OutputFormatProcessorConfig, ProcessorType
 from ..config.sampler_constraints import ColumnInequalityConstraint, ScalarInequalityConstraint
 from ..config.sampler_params import (
     BernoulliMixtureSamplerParams,

src/data_designer/essentials/__init__.py

@@ -2,6 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
+from pathlib import Path
+from typing import Literal
 
 import pandas as pd
 
@@ -53,3 +55,20 @@ def load_dataset(self) -> pd.DataFrame:
             A pandas DataFrame containing the full generated dataset.
         """
         return self.artifact_storage.load_dataset()
+
+    def write_processor_outputs_to_disk(self, output_folder: Path | str, extension: Literal["jsonl", "csv"]) -> None:
+        """Write the processor outputs to disk.
+
+        Returns:
+            None
+        """
+        output_folder = Path(output_folder)
+        output_folder.mkdir(parents=True, exist_ok=True)
+        for subfolder in self.artifact_storage.processors_outputs_path.iterdir():
+            output_file_path = output_folder / f"{subfolder.name}.{extension}"
+            with open(output_file_path, "w") as f:
+                for file_path in subfolder.glob("*.parquet"):
+                    # TODO: faster way to convert than reading and writing row by row?
+                    dataframe = pd.read_parquet(file_path)
+                    for _, row in dataframe.iterrows():
+                        f.write(row["formatted_output"].replace("\n", "\\n") + "\n")