data-designer 0.1.4__py3-none-any.whl → 0.2.0__py3-none-any.whl

data_designer/config/analysis/dataset_profiler.py

@@ -3,7 +3,7 @@
 
 from functools import cached_property
 from pathlib import Path
-from typing import Annotated, Optional, Union
+from typing import Annotated
 
 from pydantic import BaseModel, Field, field_validator
 
@@ -16,11 +16,26 @@ from data_designer.config.utils.numerical_helpers import prepare_number_for_repo
 
 
 class DatasetProfilerResults(BaseModel):
+    """Container for complete dataset profiling and analysis results.
+
+    Stores profiling results for a generated dataset, including statistics for all columns,
+    dataset-level metadata, and optional advanced profiler results. Provides methods for
+    computing derived metrics and generating formatted reports.
+
+    Attributes:
+        num_records: Actual number of records successfully generated in the dataset.
+        target_num_records: Target number of records that were requested to be generated.
+        column_statistics: List of statistics objects for all columns in the dataset. Each
+            column has statistics appropriate to its type. Must contain at least one column.
+        side_effect_column_names: Column names that were generated as side effects of other columns.
+        column_profiles: Column profiler results for specific columns when configured.
+    """
+
     num_records: int
     target_num_records: int
     column_statistics: list[Annotated[ColumnStatisticsT, Field(discriminator="column_type")]] = Field(..., min_length=1)
-    side_effect_column_names: Optional[list[str]] = None
-    column_profiles: Optional[list[ColumnProfilerResultsT]] = None
+    side_effect_column_names: list[str] | None = None
+    column_profiles: list[ColumnProfilerResultsT] | None = None
 
     @field_validator("num_records", "target_num_records", mode="before")
     def ensure_python_integers(cls, v: int) -> int:
@@ -28,10 +43,12 @@ class DatasetProfilerResults(BaseModel):
 
     @property
     def percent_complete(self) -> float:
+        """Returns the completion percentage of the dataset."""
         return 100 * self.num_records / (self.target_num_records + EPSILON)
 
     @cached_property
     def column_types(self) -> list[str]:
+        """Returns a sorted list of unique column types present in the dataset."""
         display_order = get_column_display_order()
         return sorted(
             list(set([c.column_type for c in self.column_statistics])),
@@ -39,12 +56,13 @@ class DatasetProfilerResults(BaseModel):
         )
 
     def get_column_statistics_by_type(self, column_type: DataDesignerColumnType) -> list[ColumnStatisticsT]:
+        """Filters column statistics to return only those of the specified type."""
         return [c for c in self.column_statistics if c.column_type == column_type]
 
     def to_report(
         self,
-        save_path: Optional[Union[str, Path]] = None,
-        include_sections: Optional[list[Union[ReportSection, DataDesignerColumnType]]] = None,
+        save_path: str | Path | None = None,
+        include_sections: list[ReportSection | DataDesignerColumnType] | None = None,
     ) -> None:
         """Generate and print an analysis report based on the dataset profiling results.
 

data_designer/config/analysis/utils/reporting.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 from rich.align import Align
 from rich.console import Console, Group
@@ -48,8 +48,8 @@ DEFAULT_INCLUDE_SECTIONS = [
 
 def generate_analysis_report(
     analysis: DatasetProfilerResults,
-    save_path: Optional[Union[str, Path]] = None,
-    include_sections: Optional[list[Union[ReportSection, DataDesignerColumnType]]] = None,
+    save_path: str | Path | None = None,
+    include_sections: list[ReportSection | DataDesignerColumnType] | None = None,
 ) -> None:
     """Generate an analysis report for dataset profiling results.
 

data_designer/config/base.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any
 
 import yaml
 from pydantic import BaseModel, ConfigDict
@@ -31,7 +31,7 @@ class ExportableConfigBase(ConfigBase):
         """
         return self.model_dump(mode="json")
 
-    def to_yaml(self, path: Optional[Union[str, Path]] = None, *, indent: Optional[int] = 2, **kwargs) -> Optional[str]:
+    def to_yaml(self, path: str | Path | None = None, *, indent: int | None = 2, **kwargs) -> str | None:
         """Convert the configuration to a YAML string or file.
 
         Args:
@@ -49,7 +49,7 @@ class ExportableConfigBase(ConfigBase):
         with open(path, "w") as f:
             f.write(yaml_str)
 
-    def to_json(self, path: Optional[Union[str, Path]] = None, *, indent: Optional[int] = 2, **kwargs) -> Optional[str]:
+    def to_json(self, path: str | Path | None = None, *, indent: int | None = 2, **kwargs) -> str | None:
         """Convert the configuration to a JSON string or file.
 
         Args:

data_designer/config/column_configs.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC
-from typing import Annotated, Literal, Optional, Type, Union
+from typing import Annotated, Literal
 
 from pydantic import BaseModel, Discriminator, Field, model_validator
 from typing_extensions import Self
@@ -91,7 +91,7 @@ class SamplerColumnConfig(SingleColumnConfig):
     sampler_type: SamplerType
     params: Annotated[SamplerParamsT, Discriminator("sampler_type")]
     conditional_params: dict[str, Annotated[SamplerParamsT, Discriminator("sampler_type")]] = {}
-    convert_to: Optional[str] = None
+    convert_to: str | None = None
     column_type: Literal["sampler"] = "sampler"
 
     @model_validator(mode="before")
@@ -146,8 +146,8 @@ class LLMTextColumnConfig(SingleColumnConfig):
 
     prompt: str
     model_alias: str
-    system_prompt: Optional[str] = None
-    multi_modal_context: Optional[list[ImageContext]] = None
+    system_prompt: str | None = None
+    multi_modal_context: list[ImageContext] | None = None
     column_type: Literal["llm-text"] = "llm-text"
 
     @property
@@ -222,7 +222,7 @@ class LLMStructuredColumnConfig(LLMTextColumnConfig):
         column_type: Discriminator field, always "llm-structured" for this configuration type.
     """
 
-    output_format: Union[dict, Type[BaseModel]]
+    output_format: dict | type[BaseModel]
     column_type: Literal["llm-structured"] = "llm-structured"
 
     @model_validator(mode="after")
@@ -255,7 +255,7 @@ class Score(ConfigBase):
 
     name: str = Field(..., description="A clear name for this score.")
     description: str = Field(..., description="An informative and detailed assessment guide for using this score.")
-    options: dict[Union[int, str], str] = Field(..., description="Score options in the format of {score: description}.")
+    options: dict[int | str, str] = Field(..., description="Score options in the format of {score: description}.")
 
 
 class LLMJudgeColumnConfig(LLMTextColumnConfig):
@@ -377,3 +377,24 @@ class SeedDatasetColumnConfig(SingleColumnConfig):
     """
 
     column_type: Literal["seed-dataset"] = "seed-dataset"
+
+
+class EmbeddingColumnConfig(SingleColumnConfig):
+    """Configuration for embedding generation columns.
+
+    Embedding columns generate embeddings for text input using a specified model.
+
+    Attributes:
+        target_column: The column to generate embeddings for. The column could be a single text string or a list of text strings in stringified JSON format.
+            If it is a list of text strings in stringified JSON format, the embeddings will be generated for each text string.
+        model_alias: The model to use for embedding generation.
+        column_type: Discriminator field, always "embedding" for this configuration type.
+    """
+
+    target_column: str
+    model_alias: str
+    column_type: Literal["embedding"] = "embedding"
+
+    @property
+    def required_columns(self) -> list[str]:
+        return [self.target_column]

data_designer/config/column_types.py

@@ -1,11 +1,11 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Union
 
 from typing_extensions import TypeAlias
 
 from data_designer.config.column_configs import (
+    EmbeddingColumnConfig,
     ExpressionColumnConfig,
     LLMCodeColumnConfig,
     LLMJudgeColumnConfig,
@@ -26,16 +26,17 @@ from data_designer.plugin_manager import PluginManager
 
 plugin_manager = PluginManager()
 
-ColumnConfigT: TypeAlias = Union[
-    ExpressionColumnConfig,
-    LLMCodeColumnConfig,
-    LLMJudgeColumnConfig,
-    LLMStructuredColumnConfig,
-    LLMTextColumnConfig,
-    SamplerColumnConfig,
-    SeedDatasetColumnConfig,
-    ValidationColumnConfig,
-]
+ColumnConfigT: TypeAlias = (
+    ExpressionColumnConfig
+    | LLMCodeColumnConfig
+    | LLMJudgeColumnConfig
+    | LLMStructuredColumnConfig
+    | LLMTextColumnConfig
+    | SamplerColumnConfig
+    | SeedDatasetColumnConfig
+    | ValidationColumnConfig
+    | EmbeddingColumnConfig
+)
 ColumnConfigT = plugin_manager.inject_into_column_config_type_union(ColumnConfigT)
 
 DataDesignerColumnType = create_str_enum_from_discriminated_type_union(
@@ -54,13 +55,14 @@ COLUMN_TYPE_EMOJI_MAP = {
     DataDesignerColumnType.SEED_DATASET: "🌱",
     DataDesignerColumnType.SAMPLER: "🎲",
     DataDesignerColumnType.VALIDATION: "🔍",
+    DataDesignerColumnType.EMBEDDING: "🧬",
 }
 COLUMN_TYPE_EMOJI_MAP.update(
     {DataDesignerColumnType(p.name): p.emoji for p in plugin_manager.get_column_generator_plugins()}
 )
 
 
-def column_type_used_in_execution_dag(column_type: Union[str, DataDesignerColumnType]) -> bool:
+def column_type_used_in_execution_dag(column_type: str | DataDesignerColumnType) -> bool:
     """Return True if the column type is used in the workflow execution DAG."""
     column_type = resolve_string_enum(column_type, DataDesignerColumnType)
     dag_column_types = {
@@ -70,27 +72,29 @@ def column_type_used_in_execution_dag(column_type: Union[str, DataDesignerColumn
         DataDesignerColumnType.LLM_STRUCTURED,
         DataDesignerColumnType.LLM_TEXT,
         DataDesignerColumnType.VALIDATION,
+        DataDesignerColumnType.EMBEDDING,
     }
     dag_column_types.update(plugin_manager.get_plugin_column_types(DataDesignerColumnType))
     return column_type in dag_column_types
 
 
-def column_type_is_llm_generated(column_type: Union[str, DataDesignerColumnType]) -> bool:
-    """Return True if the column type is an LLM-generated column."""
+def column_type_is_model_generated(column_type: str | DataDesignerColumnType) -> bool:
+    """Return True if the column type is a model-generated column."""
     column_type = resolve_string_enum(column_type, DataDesignerColumnType)
-    llm_generated_column_types = {
+    model_generated_column_types = {
         DataDesignerColumnType.LLM_TEXT,
         DataDesignerColumnType.LLM_CODE,
         DataDesignerColumnType.LLM_STRUCTURED,
         DataDesignerColumnType.LLM_JUDGE,
+        DataDesignerColumnType.EMBEDDING,
     }
-    llm_generated_column_types.update(
+    model_generated_column_types.update(
         plugin_manager.get_plugin_column_types(
             DataDesignerColumnType,
             required_resources=["model_registry"],
         )
     )
-    return column_type in llm_generated_column_types
+    return column_type in model_generated_column_types
 
 
 def get_column_config_from_kwargs(name: str, column_type: DataDesignerColumnType, **kwargs) -> ColumnConfigT:
@@ -121,6 +125,8 @@ def get_column_config_from_kwargs(name: str, column_type: DataDesignerColumnType
         return SamplerColumnConfig(name=name, **_resolve_sampler_kwargs(name, kwargs))
     if column_type == DataDesignerColumnType.SEED_DATASET:
         return SeedDatasetColumnConfig(name=name, **kwargs)
+    if column_type == DataDesignerColumnType.EMBEDDING:
+        return EmbeddingColumnConfig(name=name, **kwargs)
     if plugin := plugin_manager.get_column_generator_plugin_if_exists(column_type.value):
         return plugin.config_cls(name=name, **kwargs)
     raise InvalidColumnTypeError(f"🛑 {column_type} is not a valid column type.")  # pragma: no cover
@@ -135,6 +141,7 @@ def get_column_display_order() -> list[DataDesignerColumnType]:
         DataDesignerColumnType.LLM_CODE,
         DataDesignerColumnType.LLM_STRUCTURED,
         DataDesignerColumnType.LLM_JUDGE,
+        DataDesignerColumnType.EMBEDDING,
         DataDesignerColumnType.VALIDATION,
         DataDesignerColumnType.EXPRESSION,
     ]

data_designer/config/config_builder.py

@@ -6,7 +6,6 @@ from __future__ import annotations
 import json
 import logging
 from pathlib import Path
-from typing import Optional, Union
 
 from pygments import highlight
 from pygments.formatters import HtmlFormatter
@@ -19,7 +18,7 @@ from data_designer.config.column_configs import SeedDatasetColumnConfig
 from data_designer.config.column_types import (
     ColumnConfigT,
     DataDesignerColumnType,
-    column_type_is_llm_generated,
+    column_type_is_model_generated,
     get_column_config_from_kwargs,
     get_column_display_order,
 )
@@ -29,7 +28,7 @@ from data_designer.config.datastore import DatastoreSettings, fetch_seed_dataset
 from data_designer.config.default_model_settings import get_default_model_configs
 from data_designer.config.errors import BuilderConfigurationError, InvalidColumnTypeError, InvalidConfigError
 from data_designer.config.models import ModelConfig, load_model_configs
-from data_designer.config.processors import ProcessorConfig, ProcessorType, get_processor_config_from_kwargs
+from data_designer.config.processors import ProcessorConfigT, ProcessorType, get_processor_config_from_kwargs
 from data_designer.config.sampler_constraints import (
     ColumnConstraintT,
     ColumnInequalityConstraint,
@@ -69,7 +68,7 @@ class BuilderConfig(ExportableConfigBase):
     """
 
     data_designer: DataDesignerConfig
-    datastore_settings: Optional[DatastoreSettings]
+    datastore_settings: DatastoreSettings | None
 
 
 class DataDesignerConfigBuilder:
@@ -79,7 +78,7 @@ class DataDesignerConfigBuilder:
     """
 
     @classmethod
-    def from_config(cls, config: Union[dict, str, Path, BuilderConfig]) -> Self:
+    def from_config(cls, config: dict | str | Path | BuilderConfig) -> Self:
         """Create a DataDesignerConfigBuilder from an existing configuration.
 
         Args:
@@ -130,7 +129,7 @@ class DataDesignerConfigBuilder:
 
         return builder
 
-    def __init__(self, model_configs: Optional[Union[list[ModelConfig], str, Path]] = None):
+    def __init__(self, model_configs: list[ModelConfig] | str | Path | None = None):
         """Initialize a new DataDesignerConfigBuilder instance.
 
         Args:
@@ -141,11 +140,11 @@ class DataDesignerConfigBuilder:
         """
         self._column_configs = {}
         self._model_configs = _load_model_configs(model_configs)
-        self._processor_configs: list[ProcessorConfig] = []
-        self._seed_config: Optional[SeedConfig] = None
+        self._processor_configs: list[ProcessorConfigT] = []
+        self._seed_config: SeedConfig | None = None
         self._constraints: list[ColumnConstraintT] = []
         self._profilers: list[ColumnProfilerConfigT] = []
-        self._datastore_settings: Optional[DatastoreSettings] = None
+        self._datastore_settings: DatastoreSettings | None = None
 
     @property
     def model_configs(self) -> list[ModelConfig]:
@@ -206,10 +205,10 @@ class DataDesignerConfigBuilder:
 
     def add_column(
         self,
-        column_config: Optional[ColumnConfigT] = None,
+        column_config: ColumnConfigT | None = None,
         *,
-        name: Optional[str] = None,
-        column_type: Optional[DataDesignerColumnType] = None,
+        name: str | None = None,
+        column_type: DataDesignerColumnType | None = None,
         **kwargs,
     ) -> Self:
         """Add a Data Designer column configuration to the current Data Designer configuration.
@@ -246,9 +245,9 @@ class DataDesignerConfigBuilder:
 
     def add_constraint(
         self,
-        constraint: Optional[ColumnConstraintT] = None,
+        constraint: ColumnConstraintT | None = None,
         *,
-        constraint_type: Optional[ConstraintType] = None,
+        constraint_type: ConstraintType | None = None,
         **kwargs,
     ) -> Self:
         """Add a constraint to the current Data Designer configuration.
@@ -298,9 +297,9 @@ class DataDesignerConfigBuilder:
 
     def add_processor(
         self,
-        processor_config: Optional[ProcessorConfig] = None,
+        processor_config: ProcessorConfigT | None = None,
         *,
-        processor_type: Optional[ProcessorType] = None,
+        processor_type: ProcessorType | None = None,
         **kwargs,
     ) -> Self:
         """Add a processor to the current Data Designer configuration.
@@ -447,12 +446,21 @@ class DataDesignerConfigBuilder:
         return [c for c in self._constraints if c.target_column == target_column]
 
     def get_llm_gen_columns(self) -> list[ColumnConfigT]:
-        """Get all LLM-generated column configurations.
+        """Get all model-generated column configurations.
 
         Returns:
-            A list of column configurations that use LLM generation.
+            A list of column configurations that use model generation.
         """
-        return [c for c in self._column_configs.values() if column_type_is_llm_generated(c.column_type)]
+        logger.warning("get_llm_gen_columns is deprecated. Use get_model_gen_columns instead.")
+        return self.get_model_gen_columns()
+
+    def get_model_gen_columns(self) -> list[ColumnConfigT]:
+        """Get all model-generated column configurations.
+
+        Returns:
+            A list of column configurations that use model generation.
+        """
+        return [c for c in self._column_configs.values() if column_type_is_model_generated(c.column_type)]
 
     def get_columns_of_type(self, column_type: DataDesignerColumnType) -> list[ColumnConfigT]:
         """Get all column configurations of the specified type.
@@ -478,7 +486,7 @@ class DataDesignerConfigBuilder:
         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[ProcessorConfig]]:
+    def get_processor_configs(self) -> dict[BuildStage, list[ProcessorConfigT]]:
         """Get processor configuration objects.
 
         Returns:
@@ -486,7 +494,7 @@ class DataDesignerConfigBuilder:
         """
         return self._processor_configs
 
-    def get_seed_config(self) -> Optional[SeedConfig]:
+    def get_seed_config(self) -> SeedConfig | None:
         """Get the seed config for the current Data Designer configuration.
 
         Returns:
@@ -494,7 +502,7 @@ class DataDesignerConfigBuilder:
         """
         return self._seed_config
 
-    def get_seed_datastore_settings(self) -> Optional[DatastoreSettings]:
+    def get_seed_datastore_settings(self) -> DatastoreSettings | None:
         """Get most recent datastore settings for the current Data Designer configuration.
 
         Returns:
@@ -513,7 +521,7 @@ class DataDesignerConfigBuilder:
         """
         return len(self.get_columns_of_type(column_type))
 
-    def set_seed_datastore_settings(self, datastore_settings: Optional[DatastoreSettings]) -> Self:
+    def set_seed_datastore_settings(self, datastore_settings: DatastoreSettings | None) -> Self:
         """Set the datastore settings for the seed dataset.
 
         Args:
@@ -554,7 +562,7 @@ class DataDesignerConfigBuilder:
         dataset_reference: SeedDatasetReference,
         *,
         sampling_strategy: SamplingStrategy = SamplingStrategy.ORDERED,
-        selection_strategy: Optional[Union[IndexRange, PartitionBlock]] = None,
+        selection_strategy: IndexRange | PartitionBlock | None = None,
     ) -> Self:
         """Add a seed dataset to the current Data Designer configuration.
 
@@ -582,7 +590,7 @@ class DataDesignerConfigBuilder:
             self._column_configs[column_name] = SeedDatasetColumnConfig(name=column_name)
         return self
 
-    def write_config(self, path: Union[str, Path], indent: Optional[int] = 2, **kwargs) -> None:
+    def write_config(self, path: str | Path, indent: int | None = 2, **kwargs) -> None:
         """Write the current configuration to a file.
 
         Args:
@@ -653,7 +661,7 @@ class DataDesignerConfigBuilder:
         return REPR_HTML_TEMPLATE.format(css=css, highlighted_html=highlighted_html)
 
 
-def _load_model_configs(model_configs: Optional[Union[list[ModelConfig], str, Path]] = None) -> list[ModelConfig]:
+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.
     """

data_designer/config/data_designer_config.py

@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from typing import Annotated, Optional
+from typing import Annotated
 
 from pydantic import Field
 
@@ -11,7 +11,7 @@ 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 ProcessorConfig
+from data_designer.config.processors import ProcessorConfigT
 from data_designer.config.sampler_constraints import ColumnConstraintT
 from data_designer.config.seed import SeedConfig
 
@@ -33,8 +33,8 @@ class DataDesignerConfig(ExportableConfigBase):
     """
 
     columns: list[Annotated[ColumnConfigT, Field(discriminator="column_type")]] = Field(min_length=1)
-    model_configs: Optional[list[ModelConfig]] = None
-    seed_config: Optional[SeedConfig] = None
-    constraints: Optional[list[ColumnConstraintT]] = None
-    profilers: Optional[list[ColumnProfilerConfigT]] = None
-    processors: Optional[list[ProcessorConfig]] = None
+    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/datastore.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 import pandas as pd
 import pyarrow.parquet as pq
@@ -28,10 +28,10 @@ class DatastoreSettings(BaseModel):
         ...,
         description="Datastore endpoint. Use 'https://huggingface.co' for the Hugging Face Hub.",
     )
-    token: Optional[str] = Field(default=None, description="If needed, token to use for authentication.")
+    token: str | None = Field(default=None, description="If needed, token to use for authentication.")
 
 
-def get_file_column_names(file_reference: Union[str, Path, HfFileSystem], file_type: str) -> list[str]:
+def get_file_column_names(file_reference: str | Path | HfFileSystem, file_type: str) -> list[str]:
     """Get column names from a dataset file.
 
     Args:
@@ -80,7 +80,7 @@ def fetch_seed_dataset_column_names(seed_dataset_reference: SeedDatasetReference
 def fetch_seed_dataset_column_names_from_datastore(
     repo_id: str,
     filename: str,
-    datastore_settings: Optional[Union[DatastoreSettings, dict]] = None,
+    datastore_settings: DatastoreSettings | dict | None = None,
 ) -> list[str]:
     file_type = filename.split(".")[-1]
     if f".{file_type}" not in VALID_DATASET_FILE_EXTENSIONS:
@@ -115,7 +115,7 @@ def resolve_datastore_settings(datastore_settings: DatastoreSettings | dict | No
 
 
 def upload_to_hf_hub(
-    dataset_path: Union[str, Path],
+    dataset_path: str | Path,
     filename: str,
     repo_id: str,
     datastore_settings: DatastoreSettings,
@@ -171,7 +171,7 @@ def _extract_single_file_path_from_glob_pattern_if_present(
     return matching_files[0]
 
 
-def _validate_dataset_path(dataset_path: Union[str, Path], allow_glob_pattern: bool = False) -> Path:
+def _validate_dataset_path(dataset_path: str | Path, allow_glob_pattern: bool = False) -> Path:
     if allow_glob_pattern and "*" in str(dataset_path):
         parts = str(dataset_path).split("*.")
         file_path = parts[0]

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():