data-designer 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

data_designer/config/models.py

@@ -1,9 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import logging
 from abc import ABC, abstractmethod
 from enum import Enum
-import logging
 from pathlib import Path
 from typing import Any, Generic, List, Optional, TypeVar, Union
 
@@ -11,29 +11,35 @@ import numpy as np
 from pydantic import BaseModel, Field, model_validator
 from typing_extensions import Self, TypeAlias
 
-from .base import ConfigBase
-from .errors import InvalidConfigError
-from .utils.constants import (
+from data_designer.config.base import ConfigBase
+from data_designer.config.errors import InvalidConfigError
+from data_designer.config.utils.constants import (
     MAX_TEMPERATURE,
     MAX_TOP_P,
     MIN_TEMPERATURE,
     MIN_TOP_P,
 )
-from .utils.io_helpers import smart_load_yaml
+from data_designer.config.utils.io_helpers import smart_load_yaml
 
 logger = logging.getLogger(__name__)
 
 
 class Modality(str, Enum):
+    """Supported modality types for multimodal model data."""
+
     IMAGE = "image"
 
 
 class ModalityDataType(str, Enum):
+    """Data type formats for multimodal data."""
+
     URL = "url"
     BASE64 = "base64"
 
 
 class ImageFormat(str, Enum):
+    """Supported image formats for image modality."""
+
     PNG = "png"
     JPG = "jpg"
     JPEG = "jpeg"
@@ -42,6 +48,8 @@ class ImageFormat(str, Enum):
 
 
 class DistributionType(str, Enum):
+    """Types of distributions for sampling inference parameters."""
+
     UNIFORM = "uniform"
     MANUAL = "manual"
 
@@ -56,10 +64,27 @@ class ModalityContext(ABC, BaseModel):
 
 
 class ImageContext(ModalityContext):
+    """Configuration for providing image context to multimodal models.
+
+    Attributes:
+        modality: The modality type (always "image").
+        column_name: Name of the column containing image data.
+        data_type: Format of the image data ("url" or "base64").
+        image_format: Image format (required for base64 data).
+    """
+
     modality: Modality = Modality.IMAGE
     image_format: Optional[ImageFormat] = None
 
     def get_context(self, record: dict) -> dict[str, Any]:
+        """Get the context for the image modality.
+
+        Args:
+            record: The record containing the image data.
+
+        Returns:
+            The context for the image modality.
+        """
         context = dict(type="image_url")
         context_value = record[self.column_name]
         if self.data_type == ModalityDataType.URL:
@@ -90,6 +115,13 @@ class Distribution(ABC, ConfigBase, Generic[DistributionParamsT]):
 
 
 class ManualDistributionParams(ConfigBase):
+    """Parameters for manual distribution sampling.
+
+    Attributes:
+        values: List of possible values to sample from.
+        weights: Optional list of weights for each value. If not provided, all values have equal probability.
+    """
+
     values: List[float] = Field(min_length=1)
     weights: Optional[List[float]] = None
 
@@ -107,14 +139,36 @@ class ManualDistributionParams(ConfigBase):
 
 
 class ManualDistribution(Distribution[ManualDistributionParams]):
+    """Manual (discrete) distribution for sampling inference parameters.
+
+    Samples from a discrete set of values with optional weights. Useful for testing
+    specific values or creating custom probability distributions for temperature or top_p.
+
+    Attributes:
+        distribution_type: Type of distribution ("manual").
+        params: Distribution parameters (values, weights).
+    """
+
     distribution_type: Optional[DistributionType] = "manual"
     params: ManualDistributionParams
 
     def sample(self) -> float:
+        """Sample a value from the manual distribution.
+
+        Returns:
+            A float value sampled from the manual distribution.
+        """
         return float(np.random.choice(self.params.values, p=self.params.weights))
 
 
 class UniformDistributionParams(ConfigBase):
+    """Parameters for uniform distribution sampling.
+
+    Attributes:
+        low: Lower bound (inclusive).
+        high: Upper bound (exclusive).
+    """
+
     low: float
     high: float
 
@@ -126,10 +180,25 @@ class UniformDistributionParams(ConfigBase):
 
 
 class UniformDistribution(Distribution[UniformDistributionParams]):
+    """Uniform distribution for sampling inference parameters.
+
+    Samples values uniformly between low and high bounds. Useful for exploring
+    a continuous range of values for temperature or top_p.
+
+    Attributes:
+        distribution_type: Type of distribution ("uniform").
+        params: Distribution parameters (low, high).
+    """
+
     distribution_type: Optional[DistributionType] = "uniform"
     params: UniformDistributionParams
 
     def sample(self) -> float:
+        """Sample a value from the uniform distribution.
+
+        Returns:
+            A float value sampled from the uniform distribution.
+        """
         return float(np.random.uniform(low=self.params.low, high=self.params.high, size=1)[0])
 
 
@@ -137,6 +206,17 @@ DistributionT: TypeAlias = Union[UniformDistribution, ManualDistribution]
 
 
 class InferenceParameters(ConfigBase):
+    """Configuration for LLM inference parameters.
+
+    Attributes:
+        temperature: Sampling temperature (0.0-2.0). Can be a fixed value or a distribution for dynamic sampling.
+        top_p: Nucleus sampling probability (0.0-1.0). Can be a fixed value or a distribution for dynamic sampling.
+        max_tokens: Maximum number of tokens (includes both input and output tokens).
+        max_parallel_requests: Maximum number of parallel requests to the model API.
+        timeout: Timeout in seconds for each request.
+        extra_body: Additional parameters to pass to the model API.
+    """
+
     temperature: Optional[Union[float, DistributionT]] = None
     top_p: Optional[Union[float, DistributionT]] = None
     max_tokens: Optional[int] = Field(default=None, ge=1)
@@ -146,6 +226,11 @@ class InferenceParameters(ConfigBase):
 
     @property
     def generate_kwargs(self) -> dict[str, Union[float, int]]:
+        """Get the generate kwargs for the inference parameters.
+
+        Returns:
+            A dictionary of the generate kwargs.
+        """
         result = {}
         if self.temperature is not None:
             result["temperature"] = (
@@ -206,6 +291,15 @@ class InferenceParameters(ConfigBase):
 
 
 class ModelConfig(ConfigBase):
+    """Configuration for a model used for generation.
+
+    Attributes:
+        alias: User-defined alias to reference in column configurations.
+        model: Model identifier (e.g., from build.nvidia.com or other providers).
+        inference_parameters: Inference parameters for the model (temperature, top_p, max_tokens, etc.).
+        provider: Optional model provider name if using custom providers.
+    """
+
     alias: str
     model: str
     inference_parameters: InferenceParameters = Field(default_factory=InferenceParameters)
@@ -213,6 +307,16 @@ class ModelConfig(ConfigBase):
 
 
 class ModelProvider(ConfigBase):
+    """Configuration for a custom model provider.
+
+    Attributes:
+        name: Name of the model provider.
+        endpoint: API endpoint URL for the provider.
+        provider_type: Provider type (default: "openai"). Determines the API format to use.
+        api_key: Optional API key for authentication.
+        extra_body: Additional parameters to pass in API requests.
+    """
+
     name: str
     endpoint: str
     provider_type: str = "openai"

data_designer/config/preview_results.py

@@ -3,13 +3,13 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, Union
 
 import pandas as pd
 
-from .analysis.dataset_profiler import DatasetProfilerResults
-from .config_builder import DataDesignerConfigBuilder
-from .utils.visualization import WithRecordSamplerMixin
+from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
+from data_designer.config.config_builder import DataDesignerConfigBuilder
+from data_designer.config.utils.visualization import WithRecordSamplerMixin
 
 
 class PreviewResults(WithRecordSamplerMixin):
@@ -19,6 +19,7 @@ class PreviewResults(WithRecordSamplerMixin):
         config_builder: DataDesignerConfigBuilder,
         dataset: Optional[pd.DataFrame] = None,
         analysis: Optional[DatasetProfilerResults] = None,
+        processor_artifacts: Optional[dict[str, Union[list[str], str]]] = None,
     ):
         """Creates a new instance with results from a Data Designer preview run.
 
@@ -26,7 +27,9 @@ class PreviewResults(WithRecordSamplerMixin):
             config_builder: Data Designer configuration builder.
             dataset: Dataset of the preview run.
             analysis: Analysis of the preview run.
+            processor_artifacts: Artifacts generated by the processors.
         """
-        self.dataset: pd.DataFrame | None = dataset
-        self.analysis: DatasetProfilerResults | None = analysis
+        self.dataset: Optional[pd.DataFrame] = dataset
+        self.analysis: Optional[DatasetProfilerResults] = analysis
+        self.processor_artifacts: Optional[dict[str, Union[list[str], str]]] = processor_artifacts
         self._config_builder = config_builder

data_designer/config/processors.py

@@ -1,25 +1,32 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import json
 from abc import ABC
 from enum import Enum
-from typing import Literal
+from typing import Any, Literal
 
 from pydantic import Field, field_validator
 
-from .base import ConfigBase
-from .dataset_builders import BuildStage
+from data_designer.config.base import ConfigBase
+from data_designer.config.dataset_builders import BuildStage
+from data_designer.config.errors import InvalidConfigError
 
 SUPPORTED_STAGES = [BuildStage.POST_BATCH]
 
 
 class ProcessorType(str, Enum):
     DROP_COLUMNS = "drop_columns"
+    SCHEMA_TRANSFORM = "schema_transform"
 
 
 class ProcessorConfig(ConfigBase, ABC):
+    name: str = Field(
+        description="The name of the processor, used to identify the processor in the results and to write the artifacts to disk.",
+    )
     build_stage: BuildStage = Field(
-        ..., description=f"The stage at which the processor will run. Supported stages: {', '.join(SUPPORTED_STAGES)}"
+        default=BuildStage.POST_BATCH,
+        description=f"The stage at which the processor will run. Supported stages: {', '.join(SUPPORTED_STAGES)}",
     )
 
     @field_validator("build_stage")
@@ -34,8 +41,45 @@ class ProcessorConfig(ConfigBase, ABC):
 def get_processor_config_from_kwargs(processor_type: ProcessorType, **kwargs) -> ProcessorConfig:
     if processor_type == ProcessorType.DROP_COLUMNS:
         return DropColumnsProcessorConfig(**kwargs)
+    elif processor_type == ProcessorType.SCHEMA_TRANSFORM:
+        return SchemaTransformProcessorConfig(**kwargs)
 
 
 class DropColumnsProcessorConfig(ProcessorConfig):
     column_names: list[str]
     processor_type: Literal[ProcessorType.DROP_COLUMNS] = ProcessorType.DROP_COLUMNS
+
+
+class SchemaTransformProcessorConfig(ProcessorConfig):
+    template: dict[str, Any] = Field(
+        ...,
+        description="""
+        Dictionary specifying columns and templates to use in the new dataset with transformed schema.
+
+        Each key is a new column name, and each value is an object containing Jinja2 templates - for instance, a string or a list of strings.
+        Values must be JSON-serializable.
+
+        Example:
+
+        ```python
+        template = {
+            "list_of_strings": ["{{ col1 }}", "{{ col2 }}"],
+            "uppercase_string": "{{ col1 | upper }}",
+            "lowercase_string": "{{ col2 | lower }}",
+        }
+        ```
+
+        The above templates will create an new dataset with three columns: "list_of_strings", "uppercase_string", and "lowercase_string".
+        References to columns "col1" and "col2" in the templates will be replaced with the actual values of the columns in the dataset.
+        """,
+    )
+    processor_type: Literal[ProcessorType.SCHEMA_TRANSFORM] = ProcessorType.SCHEMA_TRANSFORM
+
+    @field_validator("template")
+    def validate_template(cls, v: dict[str, Any]) -> dict[str, Any]:
+        try:
+            json.dumps(v)
+        except TypeError as e:
+            if "not JSON serializable" in str(e):
+                raise InvalidConfigError("Template must be JSON serializable")
+        return v

data_designer/config/sampler_constraints.py

@@ -7,7 +7,7 @@ from typing import Union
 
 from typing_extensions import TypeAlias
 
-from .base import ConfigBase
+from data_designer.config.base import ConfigBase
 
 
 class ConstraintType(str, Enum):

data_designer/config/sampler_params.py

@@ -8,8 +8,8 @@ import pandas as pd
 from pydantic import Field, field_validator, model_validator
 from typing_extensions import Self, TypeAlias
 
-from .base import ConfigBase
-from .utils.constants import (
+from data_designer.config.base import ConfigBase
+from data_designer.config.utils.constants import (
     AVAILABLE_LOCALES,
     DEFAULT_AGE_RANGE,
     LOCALES_WITH_MANAGED_DATASETS,

data_designer/config/seed.py

@@ -8,9 +8,9 @@ from typing import Optional, Union
 from pydantic import Field, field_validator, model_validator
 from typing_extensions import Self
 
-from .base import ConfigBase
-from .datastore import DatastoreSettings
-from .utils.io_helpers import (
+from data_designer.config.base import ConfigBase
+from data_designer.config.datastore import DatastoreSettings
+from data_designer.config.utils.io_helpers import (
     VALID_DATASET_FILE_EXTENSIONS,
     validate_dataset_file_path,
     validate_path_contains_files_of_type,

data_designer/config/utils/constants.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from enum import Enum
 import os
+from enum import Enum
 from pathlib import Path
 
 from rich.theme import Theme

data_designer/config/utils/errors.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from ...errors import DataDesignerError
+from data_designer.errors import DataDesignerError
 
 
 class UserJinjaTemplateSyntaxError(DataDesignerError): ...

data_designer/config/utils/info.py

@@ -5,10 +5,14 @@ from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Literal, TypeVar
 
-from ..models import ModelConfig, ModelProvider
-from ..sampler_params import SamplerType
-from .type_helpers import get_sampler_params
-from .visualization import display_model_configs_table, display_model_providers_table, display_sampler_table
+from data_designer.config.models import ModelConfig, ModelProvider
+from data_designer.config.sampler_params import SamplerType
+from data_designer.config.utils.type_helpers import get_sampler_params
+from data_designer.config.utils.visualization import (
+    display_model_configs_table,
+    display_model_providers_table,
+    display_sampler_table,
+)
 
 
 class InfoType(str, Enum):

data_designer/config/utils/io_helpers.py

@@ -1,12 +1,12 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from datetime import date, datetime, timedelta
-from decimal import Decimal
 import json
 import logging
-from numbers import Number
 import os
+from datetime import date, datetime, timedelta
+from decimal import Decimal
+from numbers import Number
 from pathlib import Path
 from typing import Any, Union
 
@@ -14,7 +14,7 @@ import numpy as np
 import pandas as pd
 import yaml
 
-from ..errors import InvalidFileFormatError, InvalidFilePathError
+from data_designer.config.errors import InvalidFileFormatError, InvalidFilePathError
 
 logger = logging.getLogger(__name__)
 
@@ -44,7 +44,7 @@ def load_config_file(file_path: Path) -> dict:
         InvalidFileFormatError: If YAML is malformed
         InvalidConfigError: If file is empty
     """
-    from ..errors import InvalidConfigError
+    from data_designer.config.errors import InvalidConfigError
 
     if not file_path.exists():
         raise InvalidFilePathError(f"Configuration file not found: {file_path}")

data_designer/config/utils/misc.py

@@ -3,14 +3,14 @@
 
 from __future__ import annotations
 
-from contextlib import contextmanager
 import json
+from contextlib import contextmanager
 from typing import Optional, Union
 
 from jinja2 import TemplateSyntaxError, meta
 from jinja2.sandbox import ImmutableSandboxedEnvironment
 
-from .errors import UserJinjaTemplateSyntaxError
+from data_designer.config.utils.errors import UserJinjaTemplateSyntaxError
 
 REPR_LIST_LENGTH_USE_JSON = 4
 
@@ -43,7 +43,7 @@ def assert_valid_jinja2_template(template: str) -> None:
 def can_run_data_designer_locally() -> bool:
     """Returns True if Data Designer can be run locally, False otherwise."""
     try:
-        from ... import engine  # noqa: F401
+        from ... import engine  # noqa: F401, TID252
     except ImportError:
         return False
     return True

data_designer/config/utils/numerical_helpers.py

@@ -5,7 +5,7 @@ import numbers
 from numbers import Number
 from typing import Any, Type
 
-from .constants import REPORTING_PRECISION
+from data_designer.config.utils.constants import REPORTING_PRECISION
 
 
 def is_int(val: Any) -> bool:

data_designer/config/utils/type_helpers.py

@@ -1,14 +1,18 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from enum import Enum
 import inspect
+from enum import Enum
 from typing import Any, Literal, Type, get_args, get_origin
 
 from pydantic import BaseModel
 
-from .. import sampler_params
-from .errors import InvalidDiscriminatorFieldError, InvalidEnumValueError, InvalidTypeUnionError
+from data_designer.config import sampler_params
+from data_designer.config.utils.errors import (
+    InvalidDiscriminatorFieldError,
+    InvalidEnumValueError,
+    InvalidTypeUnionError,
+)
 
 
 class StrEnum(str, Enum):

data_designer/config/utils/validation.py

@@ -15,11 +15,14 @@ from rich.console import Console, Group
 from rich.padding import Padding
 from rich.panel import Panel
 
-from ..column_types import ColumnConfigT, DataDesignerColumnType, column_type_is_llm_generated
-from ..processors import ProcessorConfig, ProcessorType
-from ..validator_params import ValidatorType
-from .constants import RICH_CONSOLE_THEME
-from .misc import can_run_data_designer_locally
+from data_designer.config.column_types import ColumnConfigT, DataDesignerColumnType, column_type_is_llm_generated
+from data_designer.config.processors import ProcessorConfig, ProcessorType
+from data_designer.config.utils.constants import RICH_CONSOLE_THEME
+from data_designer.config.utils.misc import (
+    can_run_data_designer_locally,
+    get_prompt_template_keywords,
+)
+from data_designer.config.validator_params import ValidatorType
 
 
 class ViolationType(str, Enum):
@@ -63,6 +66,7 @@ def validate_data_designer_config(
     violations.extend(validate_expression_references(columns=columns, allowed_references=allowed_references))
     violations.extend(validate_columns_not_all_dropped(columns=columns))
     violations.extend(validate_drop_columns_processor(columns=columns, processor_configs=processor_configs))
+    violations.extend(validate_schema_transform_processor(columns=columns, processor_configs=processor_configs))
     if not can_run_data_designer_locally():
         violations.extend(validate_local_only_columns(columns=columns))
     return violations
@@ -271,7 +275,7 @@ def validate_drop_columns_processor(
     columns: list[ColumnConfigT],
     processor_configs: list[ProcessorConfig],
 ) -> list[Violation]:
-    all_column_names = set([c.name for c in columns])
+    all_column_names = {c.name for c in columns}
     for processor_config in processor_configs:
         if processor_config.processor_type == ProcessorType.DROP_COLUMNS:
             invalid_columns = set(processor_config.column_names) - all_column_names
@@ -288,6 +292,33 @@ def validate_drop_columns_processor(
     return []
 
 
+def validate_schema_transform_processor(
+    columns: list[ColumnConfigT],
+    processor_configs: list[ProcessorConfig],
+) -> list[Violation]:
+    violations = []
+
+    all_column_names = {c.name for c in columns}
+    for processor_config in processor_configs:
+        if processor_config.processor_type == ProcessorType.SCHEMA_TRANSFORM:
+            for col, template in processor_config.template.items():
+                template_keywords = get_prompt_template_keywords(template)
+                invalid_keywords = set(template_keywords) - all_column_names
+                if len(invalid_keywords) > 0:
+                    invalid_keywords = ", ".join([f"'{k}'" for k in invalid_keywords])
+                    message = f"Ancillary dataset processor attempts to reference columns {invalid_keywords} in the template for '{col}', but the columns are not defined in the dataset."
+                    violations.append(
+                        Violation(
+                            column=None,
+                            type=ViolationType.INVALID_REFERENCE,
+                            message=message,
+                            level=ViolationLevel.ERROR,
+                        )
+                    )
+
+    return violations
+
+
 def validate_expression_references(
     columns: list[ColumnConfigT],
     allowed_references: list[str],