data-designer 0.3.8rc2__py3-none-any.whl → 0.4.0__py3-none-any.whl

data_designer/config/config_builder.py

@@ -1,595 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import json
-import logging
-from pathlib import Path
-
-from pygments import highlight
-from pygments.formatters import HtmlFormatter
-from pygments.lexers import PythonLexer
-from typing_extensions import Self
-
-from data_designer.config.analysis.column_profilers import ColumnProfilerConfigT
-from data_designer.config.base import ExportableConfigBase
-from data_designer.config.column_configs import SeedDatasetColumnConfig
-from data_designer.config.column_types import (
-    ColumnConfigT,
-    DataDesignerColumnType,
-    get_column_config_from_kwargs,
-    get_column_display_order,
-)
-from data_designer.config.data_designer_config import DataDesignerConfig
-from data_designer.config.dataset_builders import BuildStage
-from data_designer.config.default_model_settings import get_default_model_configs
-from data_designer.config.errors import BuilderConfigurationError, BuilderSerializationError, InvalidColumnTypeError
-from data_designer.config.models import ModelConfig, load_model_configs
-from data_designer.config.processors import ProcessorConfigT, ProcessorType, get_processor_config_from_kwargs
-from data_designer.config.sampler_constraints import (
-    ColumnConstraintT,
-    ColumnInequalityConstraint,
-    ConstraintType,
-    ScalarInequalityConstraint,
-)
-from data_designer.config.seed import (
-    IndexRange,
-    PartitionBlock,
-    SamplingStrategy,
-    SeedConfig,
-)
-from data_designer.config.seed_source import DataFrameSeedSource
-from data_designer.config.seed_source_types import SeedSourceT
-from data_designer.config.utils.constants import DEFAULT_REPR_HTML_STYLE, REPR_HTML_TEMPLATE
-from data_designer.config.utils.info import ConfigBuilderInfo
-from data_designer.config.utils.io_helpers import serialize_data, smart_load_yaml
-from data_designer.config.utils.misc import can_run_data_designer_locally, json_indent_list_of_strings, kebab_to_snake
-from data_designer.config.utils.type_helpers import resolve_string_enum
-
-logger = logging.getLogger(__name__)
-
-
-class BuilderConfig(ExportableConfigBase):
-    """Configuration container for Data Designer builder.
-
-    This class holds the main Data Designer configuration along with optional
-    datastore settings needed for seed dataset operations.
-
-    Attributes:
-        data_designer: The main Data Designer configuration containing columns,
-            constraints, profilers, and other settings.
-    """
-
-    data_designer: DataDesignerConfig
-
-
-class DataDesignerConfigBuilder:
-    """Config builder for Data Designer configurations.
-
-    This class provides a high-level interface for building Data Designer configurations.
-    """
-
-    @classmethod
-    def from_config(cls, config: dict | str | Path | BuilderConfig) -> Self:
-        """Create a DataDesignerConfigBuilder from an existing configuration.
-
-        Args:
-            config: Configuration source. Can be:
-                - A dictionary containing the configuration
-                - A string or Path to a YAML/JSON configuration file
-                - A BuilderConfig object
-
-        Returns:
-            A new instance populated with the configuration from the provided source.
-
-        Raises:
-            ValueError: If the config format is invalid.
-            ValidationError: If the builder config loaded from the config is invalid.
-        """
-        if isinstance(config, BuilderConfig):
-            builder_config = config
-        else:
-            json_config = json.loads(serialize_data(smart_load_yaml(config)))
-            builder_config = BuilderConfig.model_validate(json_config)
-
-        builder = cls(model_configs=builder_config.data_designer.model_configs)
-        data_designer_config = builder_config.data_designer
-
-        for col in data_designer_config.columns:
-            builder.add_column(col)
-
-        for constraint in data_designer_config.constraints or []:
-            builder.add_constraint(constraint=constraint)
-
-        if (seed_config := data_designer_config.seed_config) is not None:
-            builder.with_seed_dataset(
-                seed_config.source,
-                sampling_strategy=seed_config.sampling_strategy,
-                selection_strategy=seed_config.selection_strategy,
-            )
-
-        return builder
-
-    def __init__(self, model_configs: list[ModelConfig] | str | Path | None = None):
-        """Initialize a new DataDesignerConfigBuilder instance.
-
-        Args:
-            model_configs: Model configurations. Can be:
-                - None to use default model configurations in local mode
-                - A list of ModelConfig objects
-                - A string or Path to a model configuration file
-        """
-        self._column_configs = {}
-        self._model_configs = _load_model_configs(model_configs)
-        self._processor_configs: list[ProcessorConfigT] = []
-        self._seed_config: SeedConfig | None = None
-        self._constraints: list[ColumnConstraintT] = []
-        self._profilers: list[ColumnProfilerConfigT] = []
-
-    @property
-    def model_configs(self) -> list[ModelConfig]:
-        """Get the model configurations for this builder.
-
-        Returns:
-            A list of ModelConfig objects used for data generation.
-        """
-        return self._model_configs
-
-    @property
-    def allowed_references(self) -> list[str]:
-        """Get all referenceable variables allowed in prompt templates and expressions.
-
-        This includes all column names and their side effect columns that can be
-        referenced in prompt templates and expressions within the configuration.
-
-        Returns:
-            A list of variable names that can be referenced in templates and expressions.
-        """
-        side_effect_columns = sum([[c.name] + c.side_effect_columns for c in self._column_configs.values()], [])
-        return list(self._column_configs.keys()) + list(set(side_effect_columns))
-
-    @property
-    def info(self) -> ConfigBuilderInfo:
-        """Get the ConfigBuilderInfo object for this builder.
-
-        Returns:
-            An object containing information about the configuration.
-        """
-        return ConfigBuilderInfo(model_configs=self._model_configs)
-
-    def add_model_config(self, model_config: ModelConfig) -> Self:
-        """Add a model configuration to the current Data Designer configuration.
-
-        Args:
-            model_config: The model configuration to add.
-        """
-        if model_config.alias in [mc.alias for mc in self._model_configs]:
-            raise BuilderConfigurationError(
-                f"🛑 Model configuration with alias {model_config.alias} already exists. Please delete the existing model configuration or choose a different alias."
-            )
-        self._model_configs.append(model_config)
-        return self
-
-    def delete_model_config(self, alias: str) -> Self:
-        """Delete a model configuration from the current Data Designer configuration by alias.
-
-        Args:
-            alias: The alias of the model configuration to delete.
-        """
-        self._model_configs = [mc for mc in self._model_configs if mc.alias != alias]
-        if len(self._model_configs) == 0:
-            logger.warning(
-                f"⚠️ No model configurations found after deleting model configuration with alias {alias}. Please add a model configuration before building the configuration."
-            )
-        return self
-
-    def add_column(
-        self,
-        column_config: ColumnConfigT | None = None,
-        *,
-        name: str | None = None,
-        column_type: DataDesignerColumnType | None = None,
-        **kwargs,
-    ) -> Self:
-        """Add a Data Designer column configuration to the current Data Designer configuration.
-
-        If no column config object is provided, you must provide the `name`, `column_type`, and any
-        additional keyword arguments that are required by the column config constructor.
-
-        Args:
-            column_config: Data Designer column config object to add.
-            name: Name of the column to add. This is only used if `column_config` is not provided.
-            column_type: Column type to add. This is only used if `column_config` is not provided.
-            **kwargs: Additional keyword arguments to pass to the column constructor.
-
-        Returns:
-            The current Data Designer config builder instance.
-
-        Raises:
-            BuilderConfigurationError: If the column name collides with an existing seed dataset column.
-        """
-        if column_config is None:
-            if name is None or column_type is None:
-                raise BuilderConfigurationError(
-                    "🛑 You must provide either a 'column_config' object or 'name' *and* 'column_type' "
-                    f"with additional keyword arguments. You provided {column_config=}, {name=}, and {column_type=}."
-                )
-            column_config = get_column_config_from_kwargs(name=name, column_type=column_type, **kwargs)
-
-        allowed_column_configs = ColumnConfigT.__args__
-        if not any(isinstance(column_config, t) for t in allowed_column_configs):
-            raise InvalidColumnTypeError(
-                f"🛑 Invalid column config object: '{column_config}'. Valid column config options are: "
-                f"{', '.join([t.__name__ for t in allowed_column_configs])}"
-            )
-
-        self._column_configs[column_config.name] = column_config
-        return self
-
-    def add_constraint(
-        self,
-        constraint: ColumnConstraintT | None = None,
-        *,
-        constraint_type: ConstraintType | None = None,
-        **kwargs,
-    ) -> Self:
-        """Add a constraint to the current Data Designer configuration.
-
-        Currently, constraints are only supported for numerical samplers.
-
-        You can either provide a constraint object directly, or provide a constraint type and
-        additional keyword arguments to construct the constraint object. Valid constraint types are:
-            - "scalar_inequality": Constraint between a column and a scalar value.
-            - "column_inequality": Constraint between two columns.
-
-        Args:
-            constraint: Constraint object to add.
-            constraint_type: Constraint type to add. Ignored when `constraint` is provided.
-            **kwargs: Additional keyword arguments to pass to the constraint constructor.
-
-        Returns:
-            The current Data Designer config builder instance.
-        """
-        if constraint is None:
-            if constraint_type is None:
-                raise BuilderConfigurationError(
-                    "🛑 You must provide either a 'constraint' object or 'constraint_type' "
-                    "with additional keyword arguments."
-                )
-            try:
-                constraint_type = ConstraintType(constraint_type)
-            except Exception:
-                raise BuilderConfigurationError(
-                    f"🛑 Invalid constraint type: {constraint_type}. Valid options are: "
-                    f"{', '.join([t.value for t in ConstraintType])}"
-                )
-            if constraint_type == ConstraintType.SCALAR_INEQUALITY:
-                constraint = ScalarInequalityConstraint(**kwargs)
-            elif constraint_type == ConstraintType.COLUMN_INEQUALITY:
-                constraint = ColumnInequalityConstraint(**kwargs)
-
-        allowed_constraint_types = ColumnConstraintT.__args__
-        if not any(isinstance(constraint, t) for t in allowed_constraint_types):
-            raise BuilderConfigurationError(
-                "🛑 Invalid constraint object. Valid constraint options are: "
-                f"{', '.join([t.__name__ for t in allowed_constraint_types])}"
-            )
-
-        self._constraints.append(constraint)
-        return self
-
-    def add_processor(
-        self,
-        processor_config: ProcessorConfigT | None = None,
-        *,
-        processor_type: ProcessorType | None = None,
-        **kwargs,
-    ) -> Self:
-        """Add a processor to the current Data Designer configuration.
-
-        You can either provide a processor config object directly, or provide a processor type and
-        additional keyword arguments to construct the processor config object.
-
-        Args:
-            processor_config: The processor configuration object to add.
-            processor_type: The type of processor to add.
-            **kwargs: Additional keyword arguments to pass to the processor constructor.
-
-        Returns:
-            The current Data Designer config builder instance.
-        """
-        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)
-
-        # Checks elsewhere fail if DropColumnsProcessor drops a column but it is not marked for drop
-        if processor_config.processor_type == ProcessorType.DROP_COLUMNS:
-            for column in processor_config.column_names:
-                if column in self._column_configs:
-                    self._column_configs[column].drop = True
-
-        self._processor_configs.append(processor_config)
-        return self
-
-    def add_profiler(self, profiler_config: ColumnProfilerConfigT) -> Self:
-        """Add a profiler to the current Data Designer configuration.
-
-        Args:
-            profiler_config: The profiler configuration object to add.
-
-        Returns:
-            The current Data Designer config builder instance.
-
-        Raises:
-            BuilderConfigurationError: If the profiler configuration is of an invalid type.
-        """
-        if not isinstance(profiler_config, ColumnProfilerConfigT):
-            if hasattr(ColumnProfilerConfigT, "__args__"):
-                valid_options = ", ".join([t.__name__ for t in ColumnProfilerConfigT.__args__])
-            else:
-                valid_options = ColumnProfilerConfigT.__name__
-            raise BuilderConfigurationError(f"🛑 Invalid profiler object. Valid profiler options are: {valid_options}")
-        self._profilers.append(profiler_config)
-        return self
-
-    def get_profilers(self) -> list[ColumnProfilerConfigT]:
-        """Get all profilers.
-
-        Returns:
-            A list of profiler configuration objects.
-        """
-        return self._profilers
-
-    def build(self) -> DataDesignerConfig:
-        """Build a DataDesignerConfig instance based on the current builder configuration.
-
-        Returns:
-            The current Data Designer config object.
-        """
-        return DataDesignerConfig(
-            model_configs=self._model_configs,
-            seed_config=self._seed_config,
-            columns=list(self._column_configs.values()),
-            constraints=self._constraints or None,
-            profilers=self._profilers or None,
-            processors=self._processor_configs or None,
-        )
-
-    def delete_constraints(self, target_column: str) -> Self:
-        """Delete all constraints for the given target column.
-
-        Args:
-            target_column: Name of the column to remove constraints for.
-
-        Returns:
-            The current Data Designer config builder instance.
-        """
-        self._constraints = [c for c in self._constraints if c.target_column != target_column]
-        return self
-
-    def delete_column(self, column_name: str) -> Self:
-        """Delete the column with the given name.
-
-        Args:
-            column_name: Name of the column to delete.
-
-        Returns:
-            The current Data Designer config builder instance.
-
-        Raises:
-            BuilderConfigurationError: If trying to delete a seed dataset column.
-        """
-        if isinstance(self._column_configs.get(column_name), SeedDatasetColumnConfig):
-            raise BuilderConfigurationError("Seed columns cannot be deleted. Please update the seed dataset instead.")
-        self._column_configs.pop(column_name, None)
-        return self
-
-    def get_column_config(self, name: str) -> ColumnConfigT:
-        """Get a column configuration by name.
-
-        Args:
-            name: Name of the column to retrieve the config for.
-
-        Returns:
-            The column configuration object.
-
-        Raises:
-            KeyError: If no column with the given name exists.
-        """
-        return self._column_configs[name]
-
-    def get_column_configs(self) -> list[ColumnConfigT]:
-        """Get all column configurations.
-
-        Returns:
-            A list of all column configuration objects.
-        """
-        return list(self._column_configs.values())
-
-    def get_constraints(self, target_column: str) -> list[ColumnConstraintT]:
-        """Get all constraints for the given target column.
-
-        Args:
-            target_column: Name of the column to get constraints for.
-
-        Returns:
-            A list of constraint objects targeting the specified column.
-        """
-        return [c for c in self._constraints if c.target_column == target_column]
-
-    def get_columns_of_type(self, column_type: DataDesignerColumnType) -> list[ColumnConfigT]:
-        """Get all column configurations of the specified type.
-
-        Args:
-            column_type: The type of columns to filter by.
-
-        Returns:
-            A list of column configurations matching the specified type.
-        """
-        column_type = resolve_string_enum(column_type, DataDesignerColumnType)
-        return [c for c in self._column_configs.values() if c.column_type == column_type]
-
-    def get_columns_excluding_type(self, column_type: DataDesignerColumnType) -> list[ColumnConfigT]:
-        """Get all column configurations excluding the specified type.
-
-        Args:
-            column_type: The type of columns to exclude.
-
-        Returns:
-            A list of column configurations that do not match the specified type.
-        """
-        column_type = resolve_string_enum(column_type, DataDesignerColumnType)
-        return [c for c in self._column_configs.values() if c.column_type != column_type]
-
-    def get_processor_configs(self) -> dict[BuildStage, list[ProcessorConfigT]]:
-        """Get processor configuration objects.
-
-        Returns:
-            A dictionary of processor configuration objects by dataset builder stage.
-        """
-        return self._processor_configs
-
-    def get_seed_config(self) -> SeedConfig | None:
-        """Get the seed config for the current Data Designer configuration.
-
-        Returns:
-            The seed config if configured, None otherwise.
-        """
-        return self._seed_config
-
-    def num_columns_of_type(self, column_type: DataDesignerColumnType) -> int:
-        """Get the count of columns of the specified type.
-
-        Args:
-            column_type: The type of columns to count.
-
-        Returns:
-            The number of columns matching the specified type.
-        """
-        return len(self.get_columns_of_type(column_type))
-
-    def with_seed_dataset(
-        self,
-        seed_source: SeedSourceT,
-        *,
-        sampling_strategy: SamplingStrategy = SamplingStrategy.ORDERED,
-        selection_strategy: IndexRange | PartitionBlock | None = None,
-    ) -> Self:
-        """Add a seed dataset to the current Data Designer configuration.
-
-        This method sets the seed dataset for the configuration, but columns are not resolved until
-        compilation (including validation) is performed by the engine using a SeedReader.
-
-        Args:
-            seed_source: The pointer to the seed dataset.
-            sampling_strategy: The sampling strategy to use when generating data from the seed dataset.
-                Defaults to ORDERED sampling.
-            selection_strategy: An optional selection strategy to use when generating data from the seed dataset.
-                Defaults to None.
-
-        Returns:
-            The current Data Designer config builder instance.
-        """
-        self._seed_config = SeedConfig(
-            source=seed_source,
-            sampling_strategy=sampling_strategy,
-            selection_strategy=selection_strategy,
-        )
-        return self
-
-    def write_config(self, path: str | Path, indent: int | None = 2, **kwargs) -> None:
-        """Write the current configuration to a file.
-
-        Args:
-            path: Path to the file to write the configuration to.
-            indent: Indentation level for the output file (default: 2).
-            **kwargs: Additional keyword arguments passed to the serialization methods used.
-
-        Raises:
-            BuilderConfigurationError: If the file format is unsupported.
-            BuilderSerializationError: If the configuration cannot be serialized.
-        """
-        if (seed_config := self.get_seed_config()) is not None and isinstance(seed_config.source, DataFrameSeedSource):
-            raise BuilderSerializationError(
-                "This builder was configured with a DataFrame seed dataset. "
-                "DataFrame seeds cannot be serialized to config files. "
-                "To serialize this configuration, change your seed dataset to a more persistent, serializable source format. "
-                "For example, you could make a local file seed source from the dataframe:\n\n"
-                "LocalFileSeedSource.from_dataframe(my_dataframe, '/path/to/data.parquet')"
-            )
-
-        cfg = self.get_builder_config()
-        suffix = Path(path).suffix
-        if suffix in {".yaml", ".yml"}:
-            cfg.to_yaml(path, indent=indent, **kwargs)
-        elif suffix == ".json":
-            cfg.to_json(path, indent=indent, **kwargs)
-        else:
-            raise BuilderConfigurationError(f"🛑 Unsupported file type: {suffix}. Must be `.yaml`, `.yml` or `.json`.")
-
-    def get_builder_config(self) -> BuilderConfig:
-        """Get the builder config for the current Data Designer configuration.
-
-        Returns:
-            The builder config.
-        """
-        return BuilderConfig(data_designer=self.build())
-
-    def __repr__(self) -> str:
-        """Generates a string representation of the DataDesignerConfigBuilder instance.
-
-        Returns:
-            A formatted string showing the builder's configuration including seed dataset and column information grouped by type.
-        """
-        if len(self._column_configs) == 0:
-            return f"{self.__class__.__name__}()"
-
-        props_to_repr = {
-            "seed_dataset": (None if self._seed_config is None else f"{self._seed_config.source.seed_type} seed"),
-        }
-
-        for column_type in get_column_display_order():
-            columns = self.get_columns_of_type(column_type)
-            if len(columns) > 0:
-                column_label = f"{kebab_to_snake(column_type.value)}_columns"
-                props_to_repr[column_label] = json_indent_list_of_strings([c.name for c in columns], indent=8)
-
-        repr_string = f"{self.__class__.__name__}(\n"
-        for k, v in props_to_repr.items():
-            if v is not None:
-                v_indented = v if "[" not in v else f"{v[:-1]}" + "    " + v[-1]
-                repr_string += f"    {k}: {v_indented}\n"
-        repr_string += ")"
-        return repr_string
-
-    def _repr_html_(self) -> str:
-        """Return an HTML representation of the DataDesignerConfigBuilder instance..
-
-        This method provides a syntax-highlighted HTML representation of the
-        builder's string representation.
-
-        Returns:
-            HTML string with syntax highlighting for the builder representation.
-        """
-        repr_string = self.__repr__()
-        formatter = HtmlFormatter(style=DEFAULT_REPR_HTML_STYLE, cssclass="code")
-        highlighted_html = highlight(repr_string, PythonLexer(), formatter)
-        css = formatter.get_style_defs(".code")
-        return REPR_HTML_TEMPLATE.format(css=css, highlighted_html=highlighted_html)
-
-
-def _load_model_configs(model_configs: list[ModelConfig] | str | Path | None = None) -> list[ModelConfig]:
-    """Resolves the provided model_configs, which may be a string or Path to a model configuration file.
-    If None or empty, returns default model configurations if possible, otherwise raises an error.
-    """
-    if model_configs:
-        return load_model_configs(model_configs)
-    elif can_run_data_designer_locally():
-        return get_default_model_configs()
-    else:
-        raise BuilderConfigurationError("🛑 Model configurations are required!")

data_designer/config/data_designer_config.py

@@ -1,40 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from typing import Annotated
-
-from pydantic import Field
-
-from data_designer.config.analysis.column_profilers import ColumnProfilerConfigT
-from data_designer.config.base import ExportableConfigBase
-from data_designer.config.column_types import ColumnConfigT
-from data_designer.config.models import ModelConfig
-from data_designer.config.processors import ProcessorConfigT
-from data_designer.config.sampler_constraints import ColumnConstraintT
-from data_designer.config.seed import SeedConfig
-
-
-class DataDesignerConfig(ExportableConfigBase):
-    """Configuration for NeMo Data Designer.
-
-    This class defines the main configuration structure for NeMo Data Designer,
-    which orchestrates the generation of synthetic data.
-
-    Attributes:
-        columns: Required list of column configurations defining how each column
-            should be generated. Must contain at least one column.
-        model_configs: Optional list of model configurations for LLM-based generation.
-            Each model config defines the model, provider, and inference parameters.
-        seed_config: Optional seed dataset settings to use for generation.
-        constraints: Optional list of column constraints.
-        profilers: Optional list of column profilers for analyzing generated data characteristics.
-    """
-
-    columns: list[Annotated[ColumnConfigT, Field(discriminator="column_type")]] = Field(min_length=1)
-    model_configs: list[ModelConfig] | None = None
-    seed_config: SeedConfig | None = None
-    constraints: list[ColumnConstraintT] | None = None
-    profilers: list[ColumnProfilerConfigT] | None = None
-    processors: list[Annotated[ProcessorConfigT, Field(discriminator="processor_type")]] | None = None

data_designer/config/dataset_builders.py

@@ -1,13 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from enum import Enum
-
-
-class BuildStage(str, Enum):
-    PRE_BATCH = "pre_batch"
-    POST_BATCH = "post_batch"
-    PRE_GENERATION = "pre_generation"
-    POST_GENERATION = "post_generation"

data_designer/config/dataset_metadata.py

@@ -1,18 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from pydantic import BaseModel
-
-
-class DatasetMetadata(BaseModel):
-    """Metadata about a generated dataset.
-
-    This object is created by the engine and passed to results objects for use
-    in visualization and other client-side utilities. It is designed to be
-    serializable so it can be sent over the wire in a client-server architecture.
-
-    Attributes:
-        seed_column_names: Names of columns from the seed dataset. Empty list if no seed dataset.
-    """
-
-    seed_column_names: list[str] = []