data-designer 0.1.3__py3-none-any.whl → 0.1.4__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

@@ -7,9 +7,9 @@ from typing import Optional
 
 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):

data_designer/config/processors.py

@@ -7,8 +7,8 @@ from typing import 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
 
 SUPPORTED_STAGES = [BuildStage.POST_BATCH]
 

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,11 @@ 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
+from data_designer.config.validator_params import ValidatorType
 
 
 class ViolationType(str, Enum):

data_designer/config/utils/visualization.py

@@ -3,11 +3,11 @@
 
 from __future__ import annotations
 
+import json
+import os
 from collections import OrderedDict
 from enum import Enum
 from functools import cached_property
-import json
-import os
 from typing import TYPE_CHECKING, Optional, Union
 
 import numpy as np
@@ -21,16 +21,16 @@ from rich.syntax import Syntax
 from rich.table import Table
 from rich.text import Text
 
-from ..base import ConfigBase
-from ..column_types import DataDesignerColumnType
-from ..models import ModelConfig, ModelProvider
-from ..sampler_params import SamplerType
-from .code_lang import code_lang_to_syntax_lexer
-from .constants import NVIDIA_API_KEY_ENV_VAR_NAME, OPENAI_API_KEY_ENV_VAR_NAME
-from .errors import DatasetSampleDisplayError
+from data_designer.config.base import ConfigBase
+from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.config.models import ModelConfig, ModelProvider
+from data_designer.config.sampler_params import SamplerType
+from data_designer.config.utils.code_lang import code_lang_to_syntax_lexer
+from data_designer.config.utils.constants import NVIDIA_API_KEY_ENV_VAR_NAME, OPENAI_API_KEY_ENV_VAR_NAME
+from data_designer.config.utils.errors import DatasetSampleDisplayError
 
 if TYPE_CHECKING:
-    from ..config_builder import DataDesignerConfigBuilder
+    from data_designer.config.config_builder import DataDesignerConfigBuilder
 
 
 console = Console()

data_designer/config/validator_params.py

@@ -7,8 +7,8 @@ from typing import Any, Optional, Union
 from pydantic import Field, field_serializer, model_validator
 from typing_extensions import Self, TypeAlias
 
-from .base import ConfigBase
-from .utils.code_lang import SQL_DIALECTS, CodeLang
+from data_designer.config.base import ConfigBase
+from data_designer.config.utils.code_lang import SQL_DIALECTS, CodeLang
 
 SUPPORTED_CODE_LANGUAGES = {CodeLang.PYTHON, *SQL_DIALECTS}
 

data_designer/engine/analysis/column_profilers/base.py

@@ -3,8 +3,8 @@
 
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
 import logging
+from abc import ABC, abstractmethod
 
 import pandas as pd
 import pyarrow as pa

data_designer/engine/analysis/dataset_profiler.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 collections.abc import Sequence
 from functools import cached_property
-import logging
 
 import pandas as pd
 from pydantic import Field, field_validator

data_designer/engine/analysis/utils/judge_score_processing.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from collections import defaultdict
 import logging
+from collections import defaultdict
 from typing import Any, Optional, Union
 
 import pandas as pd

data_designer/engine/column_generators/generators/samplers.py

@@ -1,9 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from functools import partial
 import logging
 import random
+from functools import partial
 from typing import Callable
 
 import pandas as pd

data_designer/engine/dataset_builders/artifact_storage.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 datetime
-from functools import cached_property
 import json
 import logging
-from pathlib import Path
 import shutil
+from datetime import datetime
+from functools import cached_property
+from pathlib import Path
 from typing import Union
 
 import pandas as pd

data_designer/engine/dataset_builders/column_wise_builder.py

@@ -4,8 +4,8 @@
 import functools
 import json
 import logging
-from pathlib import Path
 import time
+from pathlib import Path
 from typing import Callable
 
 import pandas as pd

data_designer/engine/dataset_builders/utils/concurrency.py

@@ -3,10 +3,10 @@
 
 from __future__ import annotations
 
-from concurrent.futures import Future, ThreadPoolExecutor
 import contextvars
 import json
 import logging
+from concurrent.futures import Future, ThreadPoolExecutor
 from threading import Lock, Semaphore
 from typing import Any, Optional, Protocol
 

data_designer/engine/dataset_builders/utils/dataset_batch_manager.py

@@ -2,8 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-from pathlib import Path
 import shutil
+from pathlib import Path
 from typing import Callable, Container, Iterator
 
 import pandas as pd

data_designer/engine/errors.py

@@ -3,7 +3,7 @@
 
 from pydantic import BaseModel, Field
 
-from ..errors import DataDesignerError
+from data_designer.errors import DataDesignerError
 
 
 class DataDesignerRuntimeError(DataDesignerError): ...

data_designer/engine/models/errors.py

@@ -3,9 +3,9 @@
 
 from __future__ import annotations
 
+import logging
 from collections.abc import Callable
 from functools import wraps
-import logging
 from typing import Any
 
 from litellm.exceptions import (

data_designer/engine/models/facade.py

@@ -3,9 +3,9 @@
 
 from __future__ import annotations
 
+import logging
 from collections.abc import Callable
 from copy import deepcopy
-import logging
 from typing import Any
 
 from litellm.types.router import DeploymentTypedDict, LiteLLM_Params

data_designer/engine/models/parsers/parser.py

@@ -4,12 +4,12 @@
 from functools import reduce
 from typing import Optional
 
+import marko
 from lxml import etree
 from lxml.etree import _Element
-import marko
 
-from data_designer.engine.models.parsers.postprocessors import merge_text_blocks
 import data_designer.engine.models.parsers.tag_parsers as tp
+from data_designer.engine.models.parsers.postprocessors import merge_text_blocks
 from data_designer.engine.models.parsers.types import (
     LLMStructuredResponse,
     PostProcessor,

data_designer/engine/models/recipes/response_recipes.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from collections.abc import Callable
 import json
+from collections.abc import Callable
 
 from pydantic import BaseModel
 

data_designer/engine/processing/ginja/environment.py

@@ -1,9 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import re
 from collections.abc import Callable
 from functools import partial, wraps
-import re
 from typing import Any
 
 from jinja2 import meta

data_designer/engine/processing/gsonschema/validators.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from copy import deepcopy
 import logging
+from copy import deepcopy
 from typing import Any, overload
 
 from jsonschema import Draft202012Validator, ValidationError, validators