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

data_designer/config/default_model_settings.py

@@ -6,9 +6,15 @@ import logging
 import os
 from functools import lru_cache
 from pathlib import Path
-from typing import Any, Literal, Optional
-
-from data_designer.config.models import InferenceParameters, ModelConfig, ModelProvider
+from typing import Any, Literal
+
+from data_designer.config.models import (
+    ChatCompletionInferenceParams,
+    EmbeddingInferenceParams,
+    InferenceParamsT,
+    ModelConfig,
+    ModelProvider,
+)
 from data_designer.config.utils.constants import (
     MANAGED_ASSETS_PATH,
     MODEL_CONFIGS_FILE_PATH,
@@ -21,46 +27,32 @@ from data_designer.config.utils.io_helpers import load_config_file, save_config_
 logger = logging.getLogger(__name__)
 
 
-def get_default_text_alias_inference_parameters() -> InferenceParameters:
-    return InferenceParameters(
-        temperature=0.85,
-        top_p=0.95,
-    )
-
-
-def get_default_reasoning_alias_inference_parameters() -> InferenceParameters:
-    return InferenceParameters(
-        temperature=0.35,
-        top_p=0.95,
-    )
-
-
-def get_default_vision_alias_inference_parameters() -> InferenceParameters:
-    return InferenceParameters(
-        temperature=0.85,
-        top_p=0.95,
-    )
-
-
-def get_default_inference_parameters(model_alias: Literal["text", "reasoning", "vision"]) -> InferenceParameters:
+def get_default_inference_parameters(
+    model_alias: Literal["text", "reasoning", "vision", "embedding"],
+    inference_parameters: dict[str, Any],
+) -> InferenceParamsT:
     if model_alias == "reasoning":
-        return get_default_reasoning_alias_inference_parameters()
+        return ChatCompletionInferenceParams(**inference_parameters)
     elif model_alias == "vision":
-        return get_default_vision_alias_inference_parameters()
+        return ChatCompletionInferenceParams(**inference_parameters)
+    elif model_alias == "embedding":
+        return EmbeddingInferenceParams(**inference_parameters)
     else:
-        return get_default_text_alias_inference_parameters()
+        return ChatCompletionInferenceParams(**inference_parameters)
 
 
 def get_builtin_model_configs() -> list[ModelConfig]:
     model_configs = []
     for provider, model_alias_map in PREDEFINED_PROVIDERS_MODEL_MAP.items():
-        for model_alias, model_id in model_alias_map.items():
+        for model_alias, settings in model_alias_map.items():
             model_configs.append(
                 ModelConfig(
                     alias=f"{provider}-{model_alias}",
-                    model=model_id,
+                    model=settings["model"],
                     provider=provider,
-                    inference_parameters=get_default_inference_parameters(model_alias),
+                    inference_parameters=get_default_inference_parameters(
+                        model_alias, settings["inference_parameters"]
+                    ),
                 )
             )
     return model_configs
@@ -93,7 +85,7 @@ def get_default_providers() -> list[ModelProvider]:
     return []
 
 
-def get_default_provider_name() -> Optional[str]:
+def get_default_provider_name() -> str | None:
     return _get_default_providers_file_content(MODEL_PROVIDERS_FILE_PATH).get("default")
 
 
@@ -103,7 +95,8 @@ def resolve_seed_default_model_settings() -> None:
             f"🍾 Default model configs were not found, so writing the following to {str(MODEL_CONFIGS_FILE_PATH)!r}"
         )
         save_config_file(
-            MODEL_CONFIGS_FILE_PATH, {"model_configs": [mc.model_dump() for mc in get_builtin_model_configs()]}
+            MODEL_CONFIGS_FILE_PATH,
+            {"model_configs": [mc.model_dump(mode="json") for mc in get_builtin_model_configs()]},
         )
 
     if not MODEL_PROVIDERS_FILE_PATH.exists():
@@ -111,7 +104,7 @@ def resolve_seed_default_model_settings() -> None:
             f"🪄  Default model providers were not found, so writing the following to {str(MODEL_PROVIDERS_FILE_PATH)!r}"
         )
         save_config_file(
-            MODEL_PROVIDERS_FILE_PATH, {"providers": [p.model_dump() for p in get_builtin_model_providers()]}
+            MODEL_PROVIDERS_FILE_PATH, {"providers": [p.model_dump(mode="json") for p in get_builtin_model_providers()]}
         )
 
     if not MANAGED_ASSETS_PATH.exists():

data_designer/config/exports.py

@@ -3,6 +3,7 @@
 
 from data_designer.config.analysis.column_profilers import JudgeScoreProfilerConfig
 from data_designer.config.column_configs import (
+    EmbeddingColumnConfig,
     ExpressionColumnConfig,
     LLMCodeColumnConfig,
     LLMJudgeColumnConfig,
@@ -19,6 +20,9 @@ from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.datastore import DatastoreSettings
 from data_designer.config.models import (
+    ChatCompletionInferenceParams,
+    EmbeddingInferenceParams,
+    GenerationType,
     ImageContext,
     ImageFormat,
     InferenceParameters,
@@ -81,6 +85,7 @@ def get_config_exports() -> list[str]:
         CodeLang.__name__,
         CodeValidatorParams.__name__,
         ColumnInequalityConstraint.__name__,
+        ChatCompletionInferenceParams.__name__,
         DataDesignerColumnType.__name__,
         DataDesignerConfig.__name__,
         DataDesignerConfigBuilder.__name__,
@@ -89,8 +94,11 @@ def get_config_exports() -> list[str]:
         DatastoreSettings.__name__,
         DatetimeSamplerParams.__name__,
         DropColumnsProcessorConfig.__name__,
+        EmbeddingColumnConfig.__name__,
+        EmbeddingInferenceParams.__name__,
         ExpressionColumnConfig.__name__,
         GaussianSamplerParams.__name__,
+        GenerationType.__name__,
         IndexRange.__name__,
         InfoType.__name__,
         ImageContext.__name__,

data_designer/config/models.py

@@ -5,10 +5,10 @@ import logging
 from abc import ABC, abstractmethod
 from enum import Enum
 from pathlib import Path
-from typing import Any, Generic, List, Optional, TypeVar, Union
+from typing import Any, Generic, Literal, TypeVar
 
 import numpy as np
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field, field_validator, model_validator
 from typing_extensions import Self, TypeAlias
 
 from data_designer.config.base import ConfigBase
@@ -74,7 +74,7 @@ class ImageContext(ModalityContext):
     """
 
     modality: Modality = Modality.IMAGE
-    image_format: Optional[ImageFormat] = None
+    image_format: ImageFormat | None = None
 
     def get_context(self, record: dict) -> dict[str, Any]:
         """Get the context for the image modality.
@@ -122,8 +122,8 @@ class ManualDistributionParams(ConfigBase):
         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
+    values: list[float] = Field(min_length=1)
+    weights: list[float] | None = None
 
     @model_validator(mode="after")
     def _normalize_weights(self) -> Self:
@@ -149,7 +149,7 @@ class ManualDistribution(Distribution[ManualDistributionParams]):
         params: Distribution parameters (values, weights).
     """
 
-    distribution_type: Optional[DistributionType] = "manual"
+    distribution_type: DistributionType | None = "manual"
     params: ManualDistributionParams
 
     def sample(self) -> float:
@@ -190,7 +190,7 @@ class UniformDistribution(Distribution[UniformDistributionParams]):
         params: Distribution parameters (low, high).
     """
 
-    distribution_type: Optional[DistributionType] = "uniform"
+    distribution_type: DistributionType | None = "uniform"
     params: UniformDistributionParams
 
     def sample(self) -> float:
@@ -202,36 +202,93 @@ class UniformDistribution(Distribution[UniformDistributionParams]):
         return float(np.random.uniform(low=self.params.low, high=self.params.high, size=1)[0])
 
 
-DistributionT: TypeAlias = Union[UniformDistribution, ManualDistribution]
+DistributionT: TypeAlias = UniformDistribution | ManualDistribution
 
 
-class InferenceParameters(ConfigBase):
-    """Configuration for LLM inference parameters.
+class GenerationType(str, Enum):
+    CHAT_COMPLETION = "chat-completion"
+    EMBEDDING = "embedding"
+
+
+class BaseInferenceParams(ConfigBase, ABC):
+    """Base configuration for 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).
+        generation_type: Type of generation (chat-completion or embedding). Acts as discriminator.
         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)
+    generation_type: GenerationType
     max_parallel_requests: int = Field(default=4, ge=1)
-    timeout: Optional[int] = Field(default=None, ge=1)
-    extra_body: Optional[dict[str, Any]] = None
+    timeout: int | None = Field(default=None, ge=1)
+    extra_body: dict[str, Any] | None = None
 
     @property
-    def generate_kwargs(self) -> dict[str, Union[float, int]]:
+    def generate_kwargs(self) -> dict[str, Any]:
         """Get the generate kwargs for the inference parameters.
 
         Returns:
             A dictionary of the generate kwargs.
         """
         result = {}
+        if self.timeout is not None:
+            result["timeout"] = self.timeout
+        if self.extra_body is not None and self.extra_body != {}:
+            result["extra_body"] = self.extra_body
+        return result
+
+    def format_for_display(self) -> str:
+        """Format inference parameters for display.
+
+        Returns:
+            Formatted string of inference parameters
+        """
+        params_dict = self.model_dump(exclude_none=True, mode="json")
+
+        if not params_dict:
+            return "(none)"
+
+        parts = []
+        for key, value in params_dict.items():
+            formatted_value = self._format_value(key, value)
+            parts.append(f"{key}={formatted_value}")
+        return ", ".join(parts)
+
+    def _format_value(self, key: str, value: Any) -> str:
+        """Format a single parameter value. Override in subclasses for custom formatting.
+
+        Args:
+            key: Parameter name
+            value: Parameter value
+
+        Returns:
+            Formatted string representation of the value
+        """
+        if isinstance(value, float):
+            return f"{value:.2f}"
+        return str(value)
+
+
+class ChatCompletionInferenceParams(BaseInferenceParams):
+    """Configuration for LLM inference parameters.
+
+    Attributes:
+        generation_type: Type of generation, always "chat-completion" for this class.
+        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).
+    """
+
+    generation_type: Literal[GenerationType.CHAT_COMPLETION] = GenerationType.CHAT_COMPLETION
+    temperature: float | DistributionT | None = None
+    top_p: float | DistributionT | None = None
+    max_tokens: int | None = Field(default=None, ge=1)
+
+    @property
+    def generate_kwargs(self) -> dict[str, Any]:
+        result = super().generate_kwargs
         if self.temperature is not None:
             result["temperature"] = (
                 self.temperature.sample() if hasattr(self.temperature, "sample") else self.temperature
@@ -240,10 +297,6 @@ class InferenceParameters(ConfigBase):
             result["top_p"] = self.top_p.sample() if hasattr(self.top_p, "sample") else self.top_p
         if self.max_tokens is not None:
             result["max_tokens"] = self.max_tokens
-        if self.timeout is not None:
-            result["timeout"] = self.timeout
-        if self.extra_body is not None and self.extra_body != {}:
-            result["extra_body"] = self.extra_body
         return result
 
     @model_validator(mode="after")
@@ -266,7 +319,7 @@ class InferenceParameters(ConfigBase):
 
     def _run_validation(
         self,
-        value: Union[float, DistributionT, None],
+        value: float | DistributionT | None,
         param_name: str,
         min_value: float,
         max_value: float,
@@ -289,6 +342,61 @@ class InferenceParameters(ConfigBase):
     def _is_value_in_range(self, value: float, min_value: float, max_value: float) -> bool:
         return min_value <= value <= max_value
 
+    def _format_value(self, key: str, value: Any) -> str:
+        """Format chat completion parameter values, including distributions.
+
+        Args:
+            key: Parameter name
+            value: Parameter value
+
+        Returns:
+            Formatted string representation of the value
+        """
+        if isinstance(value, dict) and "distribution_type" in value:
+            return "dist"
+        return super()._format_value(key, value)
+
+
+# Maintain backwards compatibility with a deprecation warning
+class InferenceParameters(ChatCompletionInferenceParams):
+    """
+    Deprecated: Use ChatCompletionInferenceParams instead.
+    This alias will be removed in a future version.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        logger.warning(
+            "InferenceParameters is deprecated and will be removed in a future version. "
+            "Use ChatCompletionInferenceParams instead."
+        )
+        super().__init__(*args, **kwargs)
+
+
+class EmbeddingInferenceParams(BaseInferenceParams):
+    """Configuration for embedding generation parameters.
+
+    Attributes:
+        generation_type: Type of generation, always "embedding" for this class.
+        encoding_format: Format of the embedding encoding ("float" or "base64").
+        dimensions: Number of dimensions for the embedding.
+    """
+
+    generation_type: Literal[GenerationType.EMBEDDING] = GenerationType.EMBEDDING
+    encoding_format: Literal["float", "base64"] = "float"
+    dimensions: int | None = None
+
+    @property
+    def generate_kwargs(self) -> dict[str, float | int]:
+        result = super().generate_kwargs
+        if self.encoding_format is not None:
+            result["encoding_format"] = self.encoding_format
+        if self.dimensions is not None:
+            result["dimensions"] = self.dimensions
+        return result
+
+
+InferenceParamsT: TypeAlias = ChatCompletionInferenceParams | EmbeddingInferenceParams | InferenceParameters
+
 
 class ModelConfig(ConfigBase):
     """Configuration for a model used for generation.
@@ -297,13 +405,31 @@ class ModelConfig(ConfigBase):
         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.).
+            The generation_type is determined by the type of inference_parameters.
         provider: Optional model provider name if using custom providers.
     """
 
     alias: str
     model: str
-    inference_parameters: InferenceParameters = Field(default_factory=InferenceParameters)
-    provider: Optional[str] = None
+    inference_parameters: InferenceParamsT = Field(default_factory=ChatCompletionInferenceParams)
+    provider: str | None = None
+
+    @property
+    def generation_type(self) -> GenerationType:
+        """Get the generation type from the inference parameters."""
+        return self.inference_parameters.generation_type
+
+    @field_validator("inference_parameters", mode="before")
+    @classmethod
+    def _convert_inference_parameters(cls, value: Any) -> Any:
+        """Convert raw dict to appropriate inference parameters type based on field presence."""
+        if isinstance(value, dict):
+            # Infer type from presence of embedding-specific fields
+            if "encoding_format" in value or "dimensions" in value:
+                return EmbeddingInferenceParams(**value)
+            else:
+                return ChatCompletionInferenceParams(**value)
+        return value
 
 
 class ModelProvider(ConfigBase):
@@ -320,11 +446,11 @@ class ModelProvider(ConfigBase):
     name: str
     endpoint: str
     provider_type: str = "openai"
-    api_key: Optional[str] = None
-    extra_body: Optional[dict[str, Any]] = None
+    api_key: str | None = None
+    extra_body: dict[str, Any] | None = None
 
 
-def load_model_configs(model_configs: Union[list[ModelConfig], str, Path]) -> list[ModelConfig]:
+def load_model_configs(model_configs: list[ModelConfig] | str | Path) -> list[ModelConfig]:
     if isinstance(model_configs, list) and all(isinstance(mc, ModelConfig) for mc in model_configs):
         return model_configs
     json_config = smart_load_yaml(model_configs)

data_designer/config/preview_results.py

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

data_designer/config/processors.py

@@ -7,6 +7,7 @@ from enum import Enum
 from typing import Any, Literal
 
 from pydantic import Field, field_validator
+from typing_extensions import TypeAlias
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.dataset_builders import BuildStage
@@ -16,11 +17,30 @@ SUPPORTED_STAGES = [BuildStage.POST_BATCH]
 
 
 class ProcessorType(str, Enum):
+    """Enumeration of available processor types.
+
+    Attributes:
+        DROP_COLUMNS: Processor that removes specified columns from the output dataset.
+        SCHEMA_TRANSFORM: Processor that creates a new dataset with a transformed schema using Jinja2 templates.
+    """
+
     DROP_COLUMNS = "drop_columns"
     SCHEMA_TRANSFORM = "schema_transform"
 
 
 class ProcessorConfig(ConfigBase, ABC):
+    """Abstract base class for all processor configuration types.
+
+    Processors are transformations that run before or after columns are generated.
+    They can modify, reshape, or augment the dataset before it's saved.
+
+    Attributes:
+        name: Unique name of the processor, used to identify the processor in results
+            and to name output artifacts on disk.
+        build_stage: The stage at which the processor runs. Currently only `POST_BATCH`
+            is supported, meaning processors run after each batch of columns is generated.
+    """
+
     name: str = Field(
         description="The name of the processor, used to identify the processor in the results and to write the artifacts to disk.",
     )
@@ -28,6 +48,7 @@ class ProcessorConfig(ConfigBase, ABC):
         default=BuildStage.POST_BATCH,
         description=f"The stage at which the processor will run. Supported stages: {', '.join(SUPPORTED_STAGES)}",
     )
+    processor_type: str
 
     @field_validator("build_stage")
     def validate_build_stage(cls, v: BuildStage) -> BuildStage:
@@ -38,7 +59,16 @@ class ProcessorConfig(ConfigBase, ABC):
         return v
 
 
-def get_processor_config_from_kwargs(processor_type: ProcessorType, **kwargs) -> ProcessorConfig:
+def get_processor_config_from_kwargs(processor_type: ProcessorType, **kwargs: Any) -> ProcessorConfig:
+    """Create a processor configuration from a processor type and keyword arguments.
+
+    Args:
+        processor_type: The type of processor to create.
+        **kwargs: Additional keyword arguments passed to the processor constructor.
+
+    Returns:
+        A processor configuration object of the specified type.
+    """
     if processor_type == ProcessorType.DROP_COLUMNS:
         return DropColumnsProcessorConfig(**kwargs)
     elif processor_type == ProcessorType.SCHEMA_TRANSFORM:
@@ -46,11 +76,39 @@ def get_processor_config_from_kwargs(processor_type: ProcessorType, **kwargs) ->
 
 
 class DropColumnsProcessorConfig(ProcessorConfig):
-    column_names: list[str]
+    """Configuration for dropping columns from the output dataset.
+
+    This processor removes specified columns from the generated dataset. The dropped
+    columns are saved separately in a `dropped-columns` directory for reference.
+    When this processor is added via the config builder, the corresponding column
+    configs are automatically marked with `drop = True`.
+
+    Alternatively, you can set `drop = True` when configuring a column.
+
+    Attributes:
+        column_names: List of column names to remove from the output dataset.
+        processor_type: Discriminator field, always `ProcessorType.DROP_COLUMNS` for this configuration type.
+    """
+
+    column_names: list[str] = Field(description="List of column names to drop from the output dataset.")
     processor_type: Literal[ProcessorType.DROP_COLUMNS] = ProcessorType.DROP_COLUMNS
 
 
 class SchemaTransformProcessorConfig(ProcessorConfig):
+    """Configuration for transforming the dataset schema using Jinja2 templates.
+
+    This processor creates a new dataset with a transformed schema. Each key in the
+    template becomes a column in the output, and values are Jinja2 templates that
+    can reference any column in the batch. The transformed dataset is written to
+    a `processors-outputs/{processor_name}/` directory alongside the main dataset.
+
+    Attributes:
+        template: Dictionary defining the output schema. Keys are new column names,
+            values are Jinja2 templates (strings, lists, or nested structures).
+            Must be JSON-serializable.
+        processor_type: Discriminator field, always `ProcessorType.SCHEMA_TRANSFORM` for this configuration type.
+    """
+
     template: dict[str, Any] = Field(
         ...,
         description="""
@@ -83,3 +141,6 @@ class SchemaTransformProcessorConfig(ProcessorConfig):
             if "not JSON serializable" in str(e):
                 raise InvalidConfigError("Template must be JSON serializable")
         return v
+
+
+ProcessorConfigT: TypeAlias = DropColumnsProcessorConfig | SchemaTransformProcessorConfig

data_designer/config/sampler_constraints.py

@@ -3,7 +3,6 @@
 
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Union
 
 from typing_extensions import TypeAlias
 
@@ -48,4 +47,4 @@ class ColumnInequalityConstraint(Constraint):
         return ConstraintType.COLUMN_INEQUALITY
 
 
-ColumnConstraintT: TypeAlias = Union[ScalarInequalityConstraint, ColumnInequalityConstraint]
+ColumnConstraintT: TypeAlias = ScalarInequalityConstraint | ColumnInequalityConstraint