moving to different UX

Author: andreatgretel

examples/.gitignore

@@ -1,2 +1,4 @@
 artifacts
 processor_outputs
+*.jsonl
+*.csv

examples/example.py

@@ -5,7 +5,7 @@
     DataDesigner,
     DataDesignerConfigBuilder,
     LLMTextColumnConfig,
-    OutputFormatProcessorConfig,
+    AncillaryDatasetProcessorConfig,
     PersonSamplerParams,
     SamplerColumnConfig,
     Score,
@@ -180,7 +180,7 @@
 }
 
 config_builder.add_processor(
-    AncillaryDatasetProcessor(
+    AncillaryDatasetProcessorConfig(
         name="jsonl_output",
         template=jsonl_entry_template,
     )
@@ -194,5 +194,7 @@
 preview.display_sample_record()
 
 results = dd.create(config_builder, num_records=20)
-jsonl_output = results.load_processor_artifact("jsonl_output")
-pd.read_parquet(jsonl_output.path_to_parquet_files).to_jsonl(desired_path, lines=True)
+path_to_processor_artifacts = results.get_path_to_processor_artifacts("jsonl_output")
+
+import pandas as pd
+pd.read_parquet(path_to_processor_artifacts).to_json("./output.jsonl", orient="records", lines=True)

src/data_designer/config/processors.py

@@ -3,7 +3,7 @@
 
 from abc import ABC
 from enum import Enum
-from typing import Literal
+from typing import Any, Literal
 
 from pydantic import Field, field_validator
 
@@ -15,7 +15,7 @@
 
 class ProcessorType(str, Enum):
     DROP_COLUMNS = "drop_columns"
-    OUTPUT_FORMAT = "output_format"
+    ANCILLARY_DATASET = "ancillary_dataset"
 
 
 class ProcessorConfig(ConfigBase, ABC):
@@ -39,18 +39,17 @@ 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.OUTPUT_FORMAT:
-        return OutputFormatProcessorConfig(**kwargs)
+    elif processor_type == ProcessorType.ANCILLARY_DATASET:
+        return AncillaryDatasetProcessorConfig(**kwargs)
 
 
 class DropColumnsProcessorConfig(ProcessorConfig):
     column_names: list[str]
     processor_type: Literal[ProcessorType.DROP_COLUMNS] = ProcessorType.DROP_COLUMNS
 
 
-class OutputFormatProcessorConfig(ProcessorConfig):
-    template: str = Field(
-        ..., description="The Jinja template to use for each entry in the dataset, as a single string."
+class AncillaryDatasetProcessorConfig(ProcessorConfig):
+    template: dict[str, Any] = Field(
+        ..., description="Jinja2 template to use for each column of the ancillary dataset. Keys are the column names, values are the Jinja2 templates."
     )
-    extension: str = Field(default="jsonl", description="The extension of the output files, e.g. 'jsonl' or 'csv'.")
-    processor_type: Literal[ProcessorType.OUTPUT_FORMAT] = ProcessorType.OUTPUT_FORMAT
+    processor_type: Literal[ProcessorType.ANCILLARY_DATASET] = ProcessorType.ANCILLARY_DATASET

…e/processing/processors/output_format.py …ocessing/processors/ancillary_dataset.pysrc/data_designer/engine/processing/processors/output_format.py renamed to src/data_designer/engine/processing/processors/ancillary_dataset.py src/data_designer/engine/processing/processors/output_format.py renamed to src/data_designer/engine/processing/processors/ancillary_dataset.py

@@ -2,11 +2,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-from pathlib import Path
+import json
+from typing import Any
 
 import pandas as pd
 
-from data_designer.config.processors import OutputFormatProcessorConfig
+from data_designer.config.processors import AncillaryDatasetProcessorConfig
 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
@@ -16,17 +17,21 @@
 logger = logging.getLogger(__name__)
 
 
-class OutputFormatProcessor(WithJinja2UserTemplateRendering, Processor[OutputFormatProcessorConfig]):
+class AncillaryDatasetProcessor(WithJinja2UserTemplateRendering, Processor[AncillaryDatasetProcessorConfig]):
     @staticmethod
     def metadata() -> ConfigurableTaskMetadata:
         return ConfigurableTaskMetadata(
-            name="output_format",
-            description="Format the dataset using a Jinja2 template.",
+            name="ancillary_dataset",
+            description="Generate an ancillary dataset using a Jinja2 template.",
             required_resources=None,
         )
 
+    @property
+    def template_as_str(self) -> str:
+        return json.dumps(self.config.template)
+
     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())
+        self.prepare_jinja2_template_renderer(self.template_as_str, data.columns.to_list())
         formatted_records = [
             self.render_template(deserialize_json_values(record)).replace("\n", "\\n")
             for record in data.to_dict(orient="records")
@@ -43,16 +48,4 @@ def process(self, data: pd.DataFrame, *, current_batch_number: int | None = None
             # Just preview the first record for now
             self.artifact_storage.processor_artifact_preview[self.config.name] = formatted_records[0]
 
-        return data
-
-    @staticmethod
-    def write_outputs_to_disk(
-        processor_config: OutputFormatProcessorConfig, artifacts_path: Path, output_path: Path
-    ) -> None:
-        output_path.mkdir(parents=True, exist_ok=True)
-        with open(output_path / f"formatted_output.{processor_config.extension}", "w") as f:
-            for file_path in artifacts_path.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"] + "\n")
+        return data

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

@@ -13,8 +13,4 @@ class Processor(ConfigurableTask[TaskConfigT], ABC):
     def metadata() -> ConfigurableTaskMetadata: ...
 
     @abstractmethod
-    def process(self, data: DataT, *, current_batch_number: int | None = None) -> DataT: ...
-
-    @staticmethod
-    @abstractmethod
-    def write_outputs_to_disk(processor_config: TaskConfigT, artifacts_path: Path, output_path: Path) -> None: ...
+    def process(self, data: DataT, *, current_batch_number: int | None = None) -> DataT: ...

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

@@ -34,12 +34,6 @@ def process(self, data: pd.DataFrame, *, current_batch_number: int | None = None
                 logger.warning(f"⚠️ Cannot drop column: `{column}` not found in the dataset.")
         return data
 
-    @staticmethod
-    def write_outputs_to_disk(
-        processor_config: DropColumnsProcessorConfig, artifacts_path: Path, output_path: Path
-    ) -> None:
-        pass
-
     def _save_dropped_columns_if_needed(self, data: pd.DataFrame, current_batch_number: int) -> None:
         logger.debug("📦 Saving dropped columns to dropped-columns directory")
         dropped_column_parquet_file_name = self.artifact_storage.create_batch_file_path(

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

@@ -3,13 +3,13 @@
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.processors import (
+    AncillaryDatasetProcessorConfig,
     DropColumnsProcessorConfig,
-    OutputFormatProcessorConfig,
     ProcessorType,
 )
 from data_designer.engine.processing.processors.base import Processor
+from data_designer.engine.processing.processors.ancillary_dataset import AncillaryDatasetProcessor
 from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
-from data_designer.engine.processing.processors.output_format import OutputFormatProcessor
 from data_designer.engine.registry.base import TaskRegistry
 
 
@@ -18,6 +18,6 @@ class ProcessorRegistry(TaskRegistry[str, Processor, ConfigBase]): ...
 
 def create_default_processor_registry() -> ProcessorRegistry:
     registry = ProcessorRegistry()
+    registry.register(ProcessorType.ANCILLARY_DATASET, AncillaryDatasetProcessor, AncillaryDatasetProcessorConfig, False)
     registry.register(ProcessorType.DROP_COLUMNS, DropColumnsProcessor, DropColumnsProcessorConfig, False)
-    registry.register(ProcessorType.OUTPUT_FORMAT, OutputFormatProcessor, OutputFormatProcessorConfig, False)
     return registry

src/data_designer/essentials/__init__.py

@@ -34,7 +34,7 @@
     UniformDistribution,
     UniformDistributionParams,
 )
-from ..config.processors import DropColumnsProcessorConfig, OutputFormatProcessorConfig, ProcessorType
+from ..config.processors import AncillaryDatasetProcessorConfig, DropColumnsProcessorConfig, ProcessorType
 from ..config.sampler_constraints import ColumnInequalityConstraint, ScalarInequalityConstraint
 from ..config.sampler_params import (
     BernoulliMixtureSamplerParams,
@@ -75,6 +75,7 @@
     pass
 
 __all__ = [
+    "AncillaryDatasetProcessorConfig",
     "BernoulliMixtureSamplerParams",
     "BernoulliSamplerParams",
     "BinomialSamplerParams",
@@ -110,7 +111,6 @@
     "ModalityContext",
     "ModalityDataType",
     "ModelConfig",
-    "OutputFormatProcessorConfig",
     "PartitionBlock",
     "PersonSamplerParams",
     "PersonFromFakerSamplerParams",

src/data_designer/interface/results.py

@@ -11,7 +11,7 @@
 from data_designer.config.config_builder import DataDesignerConfigBuilder
 from data_designer.config.utils.visualization import WithRecordSamplerMixin
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
-from data_designer.engine.processing.processors.registry import ProcessorRegistry
+from data_designer.engine.dataset_builders.errors import ArtifactStorageError
 
 
 class DatasetCreationResults(WithRecordSamplerMixin):
@@ -57,31 +57,15 @@ def load_dataset(self) -> pd.DataFrame:
         """
         return self.artifact_storage.load_dataset()
 
-    def write_processors_outputs_to_disk(
-        self,
-        processors: list[str],
-        output_folder: Path | str,
-    ) -> None:
-        """Write collected artifacts from each processor to disk.
+    def get_path_to_processor_artifacts(self, processor_name: str) -> Path:
+        """Get the path to the artifacts generated by a processor.
 
         Args:
-            processors (list[str]): List of processor names to collect artifacts from.
-            output_folder (Path | str): Path to the output folder.
+            processor_name: The name of the processor to load the artifact from.
 
         Returns:
-            None
+            The path to the artifacts.
         """
-        output_folder = Path(output_folder)
-        output_folder.mkdir(parents=True, exist_ok=True)
-
-        processors = set(processors)
-        for processor_config in self._config_builder.get_processor_configs():
-            if processor_config.name not in processors:
-                continue
-
-            ProcessorClass = ProcessorRegistry.get_for_config_type(type(processor_config))
-            ProcessorClass.write_outputs_to_disk(
-                processor_config=processor_config,
-                artifacts_path=self.artifact_storage.processors_outputs_path / processor_config.name,
-                output_path=output_folder / processor_config.name,
-            )
+        if not self.artifact_storage.processors_outputs_path.exists():
+            raise ArtifactStorageError(f"Processor {processor_name} has no artifacts.")
+        return self.artifact_storage.processors_outputs_path / processor_name

tests/config/test_processors.py

@@ -6,8 +6,8 @@
 
 from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.processors import (
+    AncillaryDatasetProcessorConfig,
     DropColumnsProcessorConfig,
-    OutputFormatProcessorConfig,
     ProcessorConfig,
     ProcessorType,
     get_processor_config_from_kwargs,
@@ -54,7 +54,7 @@ def test_drop_columns_processor_config_serialization():
 
 
 def test_output_format_processor_config_creation():
-    config = OutputFormatProcessorConfig(
+    config = AncillaryDatasetProcessorConfig(
         name="output_format_processor",
         build_stage=BuildStage.POST_BATCH,
         template='{"text": "{{ col1 }}"}',
@@ -69,19 +69,19 @@ def test_output_format_processor_config_creation():
 def test_output_format_processor_config_validation():
     # Test unsupported stage raises error
     with pytest.raises(ValidationError, match="Invalid dataset builder stage"):
-        OutputFormatProcessorConfig(
+        AncillaryDatasetProcessorConfig(
             name="output_format_processor",
             build_stage=BuildStage.PRE_BATCH,
             template='{"text": "{{ col1 }}"}',
         )
 
     # Test missing required field raises error
     with pytest.raises(ValidationError, match="Field required"):
-        OutputFormatProcessorConfig(name="output_format_processor", build_stage=BuildStage.POST_BATCH)
+        AncillaryDatasetProcessorConfig(name="output_format_processor", build_stage=BuildStage.POST_BATCH)
 
 
 def test_output_format_processor_config_serialization():
-    config = OutputFormatProcessorConfig(
+    config = AncillaryDatasetProcessorConfig(
         name="output_format_processor",
         build_stage=BuildStage.POST_BATCH,
         template='{"text": "{{ col1 }}"}',
@@ -93,7 +93,7 @@ def test_output_format_processor_config_serialization():
     assert config_dict["template"] == '{"text": "{{ col1 }}"}'
 
     # Deserialize from dict
-    config_restored = OutputFormatProcessorConfig.model_validate(config_dict)
+    config_restored = AncillaryDatasetProcessorConfig.model_validate(config_dict)
     assert config_restored.build_stage == config.build_stage
     assert config_restored.template == config.template
 
@@ -116,7 +116,7 @@ def test_get_processor_config_from_kwargs():
         build_stage=BuildStage.POST_BATCH,
         template='{"text": "{{ col1 }}"}',
     )
-    assert isinstance(config_output_format, OutputFormatProcessorConfig)
+    assert isinstance(config_output_format, AncillaryDatasetProcessorConfig)
     assert config_output_format.template == '{"text": "{{ col1 }}"}'
     assert config_output_format.processor_type == ProcessorType.OUTPUT_FORMAT