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

data_designer/config/utils/visualization.py

@@ -1,482 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import json
-import os
-from collections import OrderedDict
-from enum import Enum
-from functools import cached_property
-from typing import TYPE_CHECKING, Any
-
-from rich.console import Console, Group
-from rich.padding import Padding
-from rich.panel import Panel
-from rich.pretty import Pretty
-from rich.rule import Rule
-from rich.syntax import Syntax
-from rich.table import Table
-from rich.text import Text
-
-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
-from data_designer.lazy_heavy_imports import np, pd
-
-if TYPE_CHECKING:
-    import numpy as np
-    import pandas as pd
-
-    from data_designer.config.config_builder import DataDesignerConfigBuilder
-    from data_designer.config.dataset_metadata import DatasetMetadata
-
-
-console = Console()
-
-
-def get_nvidia_api_key() -> str | None:
-    return os.getenv(NVIDIA_API_KEY_ENV_VAR_NAME)
-
-
-def get_openai_api_key() -> str | None:
-    return os.getenv(OPENAI_API_KEY_ENV_VAR_NAME)
-
-
-class ColorPalette(str, Enum):
-    NVIDIA_GREEN = "#76b900"
-    PURPLE = "#9525c6"
-    YELLOW = "#f9c500"
-    BLUE = "#0074df"
-    RED = "#e52020"
-    ORANGE = "#ef9100"
-    MAGENTA = "#d2308e"
-    TEAL = "#1dbba4"
-
-
-class WithRecordSamplerMixin:
-    _display_cycle_index: int = 0
-    dataset_metadata: DatasetMetadata | None
-
-    @cached_property
-    def _record_sampler_dataset(self) -> pd.DataFrame:
-        if hasattr(self, "dataset") and self.dataset is not None and isinstance(self.dataset, pd.DataFrame):
-            return self.dataset
-        elif (
-            hasattr(self, "load_dataset")
-            and callable(self.load_dataset)
-            and (dataset := self.load_dataset()) is not None
-            and isinstance(dataset, pd.DataFrame)
-        ):
-            return dataset
-        else:
-            raise DatasetSampleDisplayError("No valid dataset found in results object.")
-
-    def _has_processor_artifacts(self) -> bool:
-        return hasattr(self, "processor_artifacts") and self.processor_artifacts is not None
-
-    def display_sample_record(
-        self,
-        index: int | None = None,
-        *,
-        syntax_highlighting_theme: str = "dracula",
-        background_color: str | None = None,
-        processors_to_display: list[str] | None = None,
-        hide_seed_columns: bool = False,
-    ) -> None:
-        """Display a sample record from the Data Designer dataset preview.
-
-        Args:
-            index: Index of the record to display. If None, the next record will be displayed.
-                This is useful for running the cell in a notebook multiple times.
-            syntax_highlighting_theme: Theme to use for syntax highlighting. See the `Syntax`
-                documentation from `rich` for information about available themes.
-            background_color: Background color to use for the record. See the `Syntax`
-                documentation from `rich` for information about available background colors.
-            processors_to_display: List of processors to display the artifacts for. If None, all processors will be displayed.
-            hide_seed_columns: If True, seed columns will not be displayed separately.
-        """
-        i = index or self._display_cycle_index
-
-        try:
-            record = self._record_sampler_dataset.iloc[i]
-            num_records = len(self._record_sampler_dataset)
-        except IndexError:
-            raise DatasetSampleDisplayError(f"Index {i} is out of bounds for dataset of length {num_records}.")
-
-        processor_data_to_display = None
-        if self._has_processor_artifacts() and len(self.processor_artifacts) > 0:
-            if processors_to_display is None:
-                processors_to_display = list(self.processor_artifacts.keys())
-
-            if len(processors_to_display) > 0:
-                processor_data_to_display = {}
-                for processor in processors_to_display:
-                    if (
-                        isinstance(self.processor_artifacts[processor], list)
-                        and len(self.processor_artifacts[processor]) == num_records
-                    ):
-                        processor_data_to_display[processor] = self.processor_artifacts[processor][i]
-                    else:
-                        processor_data_to_display[processor] = self.processor_artifacts[processor]
-
-        seed_column_names = (
-            None if hide_seed_columns or self.dataset_metadata is None else self.dataset_metadata.seed_column_names
-        )
-
-        display_sample_record(
-            record=record,
-            processor_data_to_display=processor_data_to_display,
-            config_builder=self._config_builder,
-            background_color=background_color,
-            syntax_highlighting_theme=syntax_highlighting_theme,
-            record_index=i,
-            seed_column_names=seed_column_names,
-        )
-        if index is None:
-            self._display_cycle_index = (self._display_cycle_index + 1) % num_records
-
-
-def create_rich_histogram_table(
-    data: dict[str, int | float],
-    column_names: tuple[int, int],
-    name_style: str = ColorPalette.BLUE.value,
-    value_style: str = ColorPalette.TEAL.value,
-    title: str | None = None,
-    **kwargs,
-) -> Table:
-    table = Table(title=title, **kwargs)
-    table.add_column(column_names[0], justify="right", style=name_style)
-    table.add_column(column_names[1], justify="left", style=value_style)
-
-    max_count = max(data.values())
-    for name, value in data.items():
-        bar = "" if max_count <= 0 else "█" * int((value / max_count) * 20)
-        table.add_row(str(name), f"{bar} {value:.1f}")
-
-    return table
-
-
-def display_sample_record(
-    record: dict | pd.Series | pd.DataFrame,
-    config_builder: DataDesignerConfigBuilder,
-    processor_data_to_display: dict[str, list[str] | str] | None = None,
-    background_color: str | None = None,
-    syntax_highlighting_theme: str = "dracula",
-    record_index: int | None = None,
-    seed_column_names: list[str] | None = None,
-):
-    if isinstance(record, (dict, pd.Series)):
-        record = pd.DataFrame([record]).iloc[0]
-    elif isinstance(record, pd.DataFrame):
-        if record.shape[0] > 1:
-            raise DatasetSampleDisplayError(
-                f"The record must be a single record. You provided a DataFrame with {record.shape[0]} records."
-            )
-        record = record.iloc[0]
-    else:
-        raise DatasetSampleDisplayError(
-            "The record must be a single record in a dictionary, pandas Series, "
-            f"or pandas DataFrame. You provided: {type(record)}."
-        )
-
-    render_list = []
-    table_kws = dict(show_lines=True, expand=True)
-
-    # Display seed columns if seed_column_names is provided and not empty
-    if seed_column_names:
-        table = Table(title="Seed Columns", **table_kws)
-        table.add_column("Name")
-        table.add_column("Value")
-        for col_name in seed_column_names:
-            if col_name in record.index:
-                table.add_row(col_name, convert_to_row_element(record[col_name]))
-        render_list.append(pad_console_element(table))
-
-    non_code_columns = (
-        config_builder.get_columns_of_type(DataDesignerColumnType.SAMPLER)
-        + config_builder.get_columns_of_type(DataDesignerColumnType.EXPRESSION)
-        + config_builder.get_columns_of_type(DataDesignerColumnType.LLM_TEXT)
-        + config_builder.get_columns_of_type(DataDesignerColumnType.LLM_STRUCTURED)
-        + config_builder.get_columns_of_type(DataDesignerColumnType.EMBEDDING)
-    )
-    if len(non_code_columns) > 0:
-        table = Table(title="Generated Columns", **table_kws)
-        table.add_column("Name")
-        table.add_column("Value")
-        for col in non_code_columns:
-            if not col.drop:
-                if col.column_type == DataDesignerColumnType.EMBEDDING:
-                    record[col.name]["embeddings"] = [
-                        get_truncated_list_as_string(embd) for embd in record[col.name].get("embeddings")
-                    ]
-                table.add_row(col.name, convert_to_row_element(record[col.name]))
-        render_list.append(pad_console_element(table))
-
-    for col in config_builder.get_columns_of_type(DataDesignerColumnType.LLM_CODE):
-        panel = Panel(
-            Syntax(
-                record[col.name],
-                lexer=code_lang_to_syntax_lexer(col.code_lang),
-                theme=syntax_highlighting_theme,
-                word_wrap=True,
-                background_color=background_color,
-            ),
-            title=col.name,
-            expand=True,
-        )
-        render_list.append(pad_console_element(panel))
-
-    validation_columns = config_builder.get_columns_of_type(DataDesignerColumnType.VALIDATION)
-    if len(validation_columns) > 0:
-        table = Table(title="Validation", **table_kws)
-        table.add_column("Name")
-        table.add_column("Value", ratio=1)
-        for col in validation_columns:
-            if not col.drop:
-                # Add is_valid before other fields
-                if "is_valid" in record[col.name]:
-                    value_to_display = {"is_valid": record[col.name].get("is_valid")} | record[col.name]
-                else:  # if columns treated separately
-                    value_to_display = {}
-                    for col_name, validation_output in record[col.name].items():
-                        value_to_display[col_name] = {
-                            "is_valid": validation_output.get("is_valid", None)
-                        } | validation_output
-
-                table.add_row(col.name, convert_to_row_element(value_to_display))
-        render_list.append(pad_console_element(table, (1, 0, 1, 0)))
-
-    llm_judge_columns = config_builder.get_columns_of_type(DataDesignerColumnType.LLM_JUDGE)
-    if len(llm_judge_columns) > 0:
-        for col in llm_judge_columns:
-            if col.drop:
-                continue
-            table = Table(title=f"LLM-as-a-Judge: {col.name}", **table_kws)
-            row = []
-            judge = record[col.name]
-
-            for measure, results in judge.items():
-                table.add_column(measure)
-                row.append(f"score: {results['score']}\nreasoning: {results['reasoning']}")
-            table.add_row(*row)
-            render_list.append(pad_console_element(table, (1, 0, 1, 0)))
-
-    if processor_data_to_display and len(processor_data_to_display) > 0:
-        for processor_name, processor_data in processor_data_to_display.items():
-            table = Table(title=f"Processor Outputs: {processor_name}", **table_kws)
-            table.add_column("Name")
-            table.add_column("Value")
-            for col, value in processor_data.items():
-                table.add_row(col, convert_to_row_element(value))
-        render_list.append(pad_console_element(table, (1, 0, 1, 0)))
-
-    if record_index is not None:
-        index_label = Text(f"[index: {record_index}]", justify="center")
-        render_list.append(index_label)
-
-    console.print(Group(*render_list), markup=False)
-
-
-def get_truncated_list_as_string(long_list: list[Any], max_items: int = 2) -> str:
-    if max_items <= 0:
-        raise ValueError("max_items must be greater than 0")
-    if len(long_list) > max_items:
-        truncated_part = long_list[:max_items]
-        return f"[{', '.join(str(x) for x in truncated_part)}, ...]"
-    else:
-        return str(long_list)
-
-
-def display_sampler_table(
-    sampler_params: dict[SamplerType, ConfigBase],
-    title: str | None = None,
-) -> None:
-    table = Table(expand=True)
-    table.add_column("Type")
-    table.add_column("Parameter")
-    table.add_column("Data Type")
-    table.add_column("Required", justify="center")
-    table.add_column("Constraints")
-
-    for sampler_type, params in sampler_params.items():
-        num = 0
-        schema = params.model_json_schema()
-        for param_name, field_info in schema["properties"].items():
-            is_required = param_name in schema.get("required", [])
-            table.add_row(
-                sampler_type if num == 0 else "",
-                param_name,
-                _get_field_type(field_info),
-                "✓" if is_required else "",
-                _get_field_constraints(field_info, schema),
-            )
-            num += 1
-        table.add_section()
-
-    title = title or "NeMo Data Designer Samplers"
-
-    group = Group(Rule(title, end="\n\n"), table)
-    console.print(group)
-
-
-def display_model_configs_table(model_configs: list[ModelConfig]) -> None:
-    table_model_configs = Table(expand=True)
-    table_model_configs.add_column("Alias")
-    table_model_configs.add_column("Model")
-    table_model_configs.add_column("Provider")
-    table_model_configs.add_column("Inference Parameters")
-    for model_config in model_configs:
-        params_display = model_config.inference_parameters.format_for_display()
-
-        table_model_configs.add_row(
-            model_config.alias,
-            model_config.model,
-            model_config.provider,
-            params_display,
-        )
-    group_args: list = [Rule(title="Model Configs"), table_model_configs]
-    if len(model_configs) == 0:
-        subtitle = Text(
-            "‼️ No model configs found. Please provide at least one model config to the config builder",
-            style="dim",
-            justify="center",
-        )
-        group_args.insert(1, subtitle)
-    group = Group(*group_args)
-    console.print(group)
-
-
-def display_model_providers_table(model_providers: list[ModelProvider]) -> None:
-    table_model_providers = Table(expand=True)
-    table_model_providers.add_column("Name")
-    table_model_providers.add_column("Endpoint")
-    table_model_providers.add_column("API Key")
-    for model_provider in model_providers:
-        api_key = model_provider.api_key
-        if model_provider.api_key == OPENAI_API_KEY_ENV_VAR_NAME:
-            if get_openai_api_key() is not None:
-                api_key = mask_api_key(get_openai_api_key())
-            else:
-                api_key = f"* {OPENAI_API_KEY_ENV_VAR_NAME!r} not set in environment variables * "
-        elif model_provider.api_key == NVIDIA_API_KEY_ENV_VAR_NAME:
-            if get_nvidia_api_key() is not None:
-                api_key = mask_api_key(get_nvidia_api_key())
-            else:
-                api_key = f"* {NVIDIA_API_KEY_ENV_VAR_NAME!r} not set in environment variables *"
-        else:
-            api_key = mask_api_key(model_provider.api_key)
-        table_model_providers.add_row(model_provider.name, model_provider.endpoint, api_key)
-    group = Group(Rule(title="Model Providers"), table_model_providers)
-    console.print(group)
-
-
-def mask_api_key(api_key: str | None) -> str:
-    """Mask API keys for display.
-
-    Environment variable names (all uppercase) are kept visible.
-    Actual API keys are masked to show only the last 4 characters.
-
-    Args:
-        api_key: The API key to mask.
-
-    Returns:
-        Masked API key string or "(not set)" if None.
-    """
-    if not api_key:
-        return "(not set)"
-
-    # Keep environment variable names visible
-    if api_key.isupper():
-        return api_key
-
-    # Mask actual API keys
-    return "***" + api_key[-4:] if len(api_key) > 4 else "***"
-
-
-def convert_to_row_element(elem):
-    try:
-        elem = Pretty(json.loads(elem))
-    except (TypeError, json.JSONDecodeError):
-        pass
-    if isinstance(elem, (np.integer, np.floating, np.ndarray)):
-        elem = str(elem)
-    elif isinstance(elem, (list, dict)):
-        elem = Pretty(elem)
-    return elem
-
-
-def pad_console_element(elem, padding=(1, 0, 1, 0)):
-    return Padding(elem, padding)
-
-
-def _get_field_type(field: dict) -> str:
-    """Extract human-readable type information from a JSON Schema field."""
-
-    # single type
-    if "type" in field:
-        if field["type"] == "array":
-            return " | ".join([f"{f.strip()}[]" for f in _get_field_type(field["items"]).split("|")])
-        if field["type"] == "object":
-            return "dict"
-        return field["type"]
-
-    # union type
-    elif "anyOf" in field:
-        types = []
-        for f in field["anyOf"]:
-            if "$ref" in f:
-                types.append("enum")
-            elif f.get("type") == "array":
-                if "items" in f and "$ref" in f["items"]:
-                    types.append("enum[]")
-                else:
-                    types.append(f"{f['items']['type']}[]")
-            else:
-                types.append(f.get("type", ""))
-        return " | ".join(t for t in types if t)
-
-    return ""
-
-
-def _get_field_constraints(field: dict, schema: dict) -> str:
-    """Extract human-readable constraints from a JSON Schema field."""
-    constraints = []
-
-    # numeric constraints
-    if "minimum" in field:
-        constraints.append(f">= {field['minimum']}")
-    if "exclusiveMinimum" in field:
-        constraints.append(f"> {field['exclusiveMinimum']}")
-    if "maximum" in field:
-        constraints.append(f"<= {field['maximum']}")
-    if "exclusiveMaximum" in field:
-        constraints.append(f"< {field['exclusiveMaximum']}")
-
-    # string constraints
-    if "minLength" in field:
-        constraints.append(f"len > {field['minLength']}")
-    if "maxLength" in field:
-        constraints.append(f"len < {field['maxLength']}")
-
-    # array constraints
-    if "minItems" in field:
-        constraints.append(f"len > {field['minItems']}")
-    if "maxItems" in field:
-        constraints.append(f"len < {field['maxItems']}")
-
-    # enum constraints
-    if "enum" in _get_field_type(field) and "$defs" in schema:
-        enum_values = []
-        for defs in schema["$defs"].values():
-            if "enum" in defs:
-                enum_values.extend(defs["enum"])
-        if len(enum_values) > 0:
-            enum_values = OrderedDict.fromkeys(enum_values)
-            constraints.append(f"allowed: {', '.join(enum_values.keys())}")
-
-    return ", ".join(constraints)

data_designer/config/validator_params.py

@@ -1,94 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from enum import Enum
-from typing import Any
-
-from pydantic import Field, field_serializer, model_validator
-from typing_extensions import Self, TypeAlias
-
-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}
-
-
-class ValidatorType(str, Enum):
-    CODE = "code"
-    LOCAL_CALLABLE = "local_callable"
-    REMOTE = "remote"
-
-
-class CodeValidatorParams(ConfigBase):
-    """Configuration for code validation. Supports Python and SQL code validation.
-
-    Attributes:
-        code_lang: The language of the code to validate. Supported values include: `python`,
-            `sql:sqlite`, `sql:postgres`, `sql:mysql`, `sql:tsql`, `sql:bigquery`, `sql:ansi`.
-    """
-
-    code_lang: CodeLang = Field(description="The language of the code to validate")
-
-    @model_validator(mode="after")
-    def validate_code_lang(self) -> Self:
-        if self.code_lang not in SUPPORTED_CODE_LANGUAGES:
-            raise ValueError(
-                f"Unsupported code language, supported languages are: {[lang.value for lang in SUPPORTED_CODE_LANGUAGES]}"
-            )
-        return self
-
-
-class LocalCallableValidatorParams(ConfigBase):
-    """Configuration for local callable validation. Expects a function to be passed that validates the data.
-
-    Attributes:
-        validation_function: Function (`Callable[[pd.DataFrame], pd.DataFrame]`) to validate the
-            data. Output must contain a column `is_valid` of type `bool`.
-        output_schema: The JSON schema for the local callable validator's output. If not provided,
-            the output will not be validated.
-    """
-
-    validation_function: Any = Field(
-        description="Function (Callable[[pd.DataFrame], pd.DataFrame]) to validate the data"
-    )
-    output_schema: dict[str, Any] | None = Field(
-        default=None, description="Expected schema for local callable validator's output"
-    )
-
-    @field_serializer("validation_function")
-    def serialize_validation_function(self, v: Any) -> Any:
-        return v.__name__
-
-    @model_validator(mode="after")
-    def validate_validation_function(self) -> Self:
-        if not callable(self.validation_function):
-            raise ValueError("Validation function must be a callable")
-        return self
-
-
-class RemoteValidatorParams(ConfigBase):
-    """Configuration for remote validation. Sends data to a remote endpoint for validation.
-
-    Attributes:
-        endpoint_url: The URL of the remote endpoint.
-        output_schema: The JSON schema for the remote validator's output. If not provided,
-            the output will not be validated.
-        timeout: The timeout for the HTTP request in seconds. Defaults to 30.0.
-        max_retries: The maximum number of retry attempts. Defaults to 3.
-        retry_backoff: The backoff factor for the retry delay in seconds. Defaults to 2.0.
-        max_parallel_requests: The maximum number of parallel requests to make. Defaults to 4.
-    """
-
-    endpoint_url: str = Field(description="URL of the remote endpoint")
-    output_schema: dict[str, Any] | None = Field(
-        default=None, description="Expected schema for remote validator's output"
-    )
-    timeout: float = Field(default=30.0, gt=0, description="The timeout for the HTTP request")
-    max_retries: int = Field(default=3, ge=0, description="The maximum number of retry attempts")
-    retry_backoff: float = Field(default=2.0, gt=1, description="The backoff factor for the retry delay")
-    max_parallel_requests: int = Field(default=4, ge=1, description="The maximum number of parallel requests to make")
-
-
-ValidatorParamsT: TypeAlias = CodeValidatorParams | LocalCallableValidatorParams | RemoteValidatorParams

data_designer/engine/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0

data_designer/engine/analysis/column_profilers/base.py

@@ -1,49 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import logging
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
-
-from pydantic import BaseModel, model_validator
-from typing_extensions import Self
-
-from data_designer.config.base import ConfigBase
-from data_designer.config.column_configs import SingleColumnConfig
-from data_designer.config.column_types import DataDesignerColumnType
-from data_designer.engine.configurable_task import ConfigurableTask, TaskConfigT
-from data_designer.lazy_heavy_imports import pd
-
-if TYPE_CHECKING:
-    import pandas as pd
-
-logger = logging.getLogger(__name__)
-
-
-class ColumnConfigWithDataFrame(ConfigBase):
-    column_config: SingleColumnConfig
-    df: pd.DataFrame
-
-    @model_validator(mode="after")
-    def validate_column_exists(self) -> Self:
-        if self.column_config.name not in self.df.columns:
-            raise ValueError(f"Column {self.column_config.name!r} not found in DataFrame")
-        return self
-
-    def as_tuple(self) -> tuple[SingleColumnConfig, pd.DataFrame]:
-        return (self.column_config, self.df)
-
-
-class ColumnProfiler(ConfigurableTask[TaskConfigT], ABC):
-    @staticmethod
-    @abstractmethod
-    def get_applicable_column_types() -> list[DataDesignerColumnType]:
-        """Returns a list of column types that this profiler can be applied to during dataset profiling."""
-
-    @abstractmethod
-    def profile(self, column_config_with_df: ColumnConfigWithDataFrame) -> BaseModel: ...
-
-    def _initialize(self) -> None:
-        logger.info(f"💫 Initializing column profiler: '{self.name}'")