data-designer-engine 0.4.0__py3-none-any.whl

data_designer/engine/__init__.py

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

data_designer/engine/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.4.0'
+__version_tuple__ = version_tuple = (0, 4, 0)
+
+__commit_id__ = commit_id = None

data_designer/engine/analysis/column_profilers/base.py

@@ -0,0 +1,49 @@
+# 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}'")

data_designer/engine/analysis/column_profilers/judge_score_profiler.py

@@ -0,0 +1,153 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+import random
+from typing import TYPE_CHECKING
+
+from data_designer.config.analysis.column_profilers import (
+    JudgeScoreProfilerConfig,
+    JudgeScoreProfilerResults,
+    JudgeScoreSummary,
+)
+from data_designer.config.analysis.column_statistics import (
+    ColumnDistributionType,
+    MissingValue,
+)
+from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame, ColumnProfiler
+from data_designer.engine.analysis.utils.judge_score_processing import (
+    extract_judge_score_distributions,
+    sample_scores_and_reasoning,
+)
+from data_designer.engine.models.recipes.response_recipes import TextResponseRecipe
+
+if TYPE_CHECKING:
+    from data_designer.config.analysis.column_profilers import JudgeScoreSample
+    from data_designer.config.analysis.column_statistics import (
+        CategoricalDistribution,
+        CategoricalHistogramData,
+        NumericalDistribution,
+    )
+    from data_designer.engine.models.facade import ModelFacade
+
+logger = logging.getLogger(__name__)
+
+
+class JudgeScoreProfiler(ColumnProfiler[JudgeScoreProfilerConfig]):
+    @staticmethod
+    def get_applicable_column_types() -> list[DataDesignerColumnType]:
+        return [DataDesignerColumnType.LLM_JUDGE]
+
+    def get_model(self, model_alias: str) -> ModelFacade:
+        return self.resource_provider.model_registry.get_model(model_alias=model_alias)
+
+    def profile(self, column_config_with_df: ColumnConfigWithDataFrame) -> JudgeScoreProfilerResults:
+        column_config, df = column_config_with_df.as_tuple()
+
+        logger.info(
+            f"{column_config.get_column_emoji()} Analyzing LLM-as-judge scores for column: '{column_config.name}'"
+        )
+
+        score_summaries = {}
+        score_distributions = extract_judge_score_distributions(column_config, df)
+
+        if self.config.summary_score_sample_size is None or isinstance(score_distributions, MissingValue):
+            return JudgeScoreProfilerResults(
+                summaries={},
+                column_name=column_config.name,
+                score_distributions=score_distributions,
+            )
+
+        for score in column_config.scores:
+            score_name = score.name
+            logger.info(f"{random.choice(['👩‍⚖️', '👨‍⚖️'])} Summarizing LLM-as-judge score: '{score_name}'")
+            score_sample = sample_scores_and_reasoning(
+                scores=score_distributions.scores[score_name],
+                reasoning=score_distributions.reasoning[score_name],
+                num_samples=self.config.summary_score_sample_size,
+            )
+
+            score_summaries[score_name] = self._summarize_score_sample(
+                name=score_name,
+                sample=score_sample,
+                histogram=score_distributions.histograms[score_name],
+                distribution=score_distributions.distributions[score_name],
+                distribution_type=score_distributions.distribution_types[score_name],
+            )
+
+        return JudgeScoreProfilerResults(
+            column_name=column_config.name,
+            summaries=score_summaries,
+            score_distributions=score_distributions,
+        )
+
+    def _summarize_score_sample(
+        self,
+        name: str,
+        sample: list[JudgeScoreSample],
+        histogram: CategoricalHistogramData,
+        distribution: CategoricalDistribution | NumericalDistribution | MissingValue,
+        distribution_type: ColumnDistributionType,
+    ) -> JudgeScoreSummary:
+        if isinstance(distribution, MissingValue) or not sample:
+            return JudgeScoreSummary(
+                score_name=name,
+                summary="No judge score information available to summarize.",
+                score_samples=sample,
+            )
+
+        category_info = []
+        total_count = sum(histogram.counts)
+        for cat, count in zip(histogram.categories, histogram.counts):
+            percentage = (count / total_count) * 100
+            category_info.append(f"{cat}: {count} records ({percentage:.1f}%)")
+
+        distribution_context = f"Score distribution - {', '.join(category_info)}, "
+        if distribution_type == ColumnDistributionType.CATEGORICAL:
+            distribution_context += f"Most common value: {distribution.most_common_value}. "
+        if distribution_type == ColumnDistributionType.NUMERICAL:
+            distribution_context += f"Mean score: {distribution.mean:.2f}. "
+
+        logger.info(f"  |-- number of score samples: {len(sample)}")
+        logger.info(f"  |-- {distribution_context.lower()}")
+
+        combined_reasoning = "\n".join([r.reasoning for r in sample])
+        prompt = (
+            f"Based on the following evaluator reasoning for the '{name}' criterion, "
+            "provide a concise summary that captures both the strengths and areas for improvement mentioned. "
+            "Be specific about what worked well and what needs improvement.\n\n"
+            f"Overall distribution of scores: {distribution_context}"
+            f"\nA sample of reasoning:\n{combined_reasoning}\n\n"
+            "Do not include any titles like `Summary` or `Summary:`. "
+            "Do not wrap the summary in quotation marks. "
+            "YOU WILL PRODUCE LESS THAN 75 WORDS in a readable sentence format. "
+            "No need to use bullets or headers. Write naturally."
+        )
+
+        system_prompt = (
+            "You are an expert at distilling complex feedback into concise summaries. "
+            "Focus on specificity and balance, incorporating both the distribution context and individual reasoning examples."
+        )
+
+        try:
+            model = self.get_model(self.config.model_alias)
+            recipe = TextResponseRecipe()
+            summary, _ = model.generate(
+                prompt=recipe.apply_recipe_to_user_prompt(prompt),
+                system_prompt=recipe.apply_recipe_to_system_prompt(system_prompt),
+                parser=recipe.parse,
+            )
+            return JudgeScoreSummary(
+                score_name=name,
+                summary=summary.strip(),
+                score_samples=sample,
+            )
+        except Exception as e:
+            return JudgeScoreSummary(
+                score_name=name,
+                summary=f"Score summarization failed: {e}",
+                score_samples=sample,
+            )

data_designer/engine/analysis/column_profilers/registry.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from data_designer.config.analysis.column_profilers import ColumnProfilerType
+from data_designer.config.base import ConfigBase
+from data_designer.engine.analysis.column_profilers.base import ColumnProfiler
+from data_designer.engine.analysis.column_profilers.judge_score_profiler import (
+    JudgeScoreProfiler,
+    JudgeScoreProfilerConfig,
+)
+from data_designer.engine.registry.base import TaskRegistry
+
+
+class ColumnProfilerRegistry(TaskRegistry[ColumnProfilerType, ColumnProfiler, ConfigBase]): ...
+
+
+def create_default_column_profiler_registry() -> ColumnProfilerRegistry:
+    registry = ColumnProfilerRegistry()
+    registry.register(ColumnProfilerType.JUDGE_SCORE, JudgeScoreProfiler, JudgeScoreProfilerConfig, False)
+    return registry

data_designer/engine/analysis/column_statistics.py

@@ -0,0 +1,145 @@
+# 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 typing import TYPE_CHECKING, Any, TypeAlias
+
+from pydantic import BaseModel
+from typing_extensions import Self
+
+from data_designer.config.analysis.column_statistics import (
+    DEFAULT_COLUMN_STATISTICS_MAP,
+    ColumnStatisticsT,
+    GeneralColumnStatistics,
+)
+from data_designer.config.column_types import ColumnConfigT, DataDesignerColumnType
+from data_designer.config.sampler_params import SamplerType, is_numerical_sampler_type
+from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame
+from data_designer.engine.analysis.utils.column_statistics_calculations import (
+    ColumnDistributionType,
+    calculate_column_distribution,
+    calculate_general_column_info,
+    calculate_token_stats,
+    calculate_validation_column_info,
+)
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+
+class GeneralColumnStatisticsCalculator(BaseModel):
+    column_config_with_df: ColumnConfigWithDataFrame
+
+    @property
+    def column_config(self) -> ColumnConfigT:
+        return self.column_config_with_df.column_config
+
+    @property
+    def df(self) -> pd.DataFrame:
+        return self.column_config_with_df.df
+
+    @property
+    def column_statistics_type(self) -> type[ColumnStatisticsT]:
+        return DEFAULT_COLUMN_STATISTICS_MAP.get(self.column_config.column_type, GeneralColumnStatistics)
+
+    def calculate(self) -> Self:
+        """Calculate all the column statistics fields for the given column configuration and dataset profiler.
+
+        This method dynamically collects all class methods prefixed with 'calculate_' and invokes them to
+        compute various column statistics, aggregating their results into a single statistics object.
+        """
+        calculate_methods = [
+            name for name in dir(self) if name.startswith("calculate_") and callable(getattr(self, name))
+        ]
+        return self.column_statistics_type(
+            column_name=self.column_config.name,
+            **{k: v for name in calculate_methods for k, v in getattr(self, name)().items()},
+        )
+
+    def calculate_general_column_info(self) -> dict[str, Any]:
+        return calculate_general_column_info(self.column_config.name, self.df)
+
+    def __repr__(self) -> str:
+        params = []
+        for field, value in self.model_dump(mode="json").items():
+            params.append(f"    {field}: {value}")
+        params_str = "\n".join(params)
+        return f"{self.__class__.__name__}(\n{params_str}\n)"
+
+
+class LLMTextColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
+    def calculate_token_stats(self) -> dict[str, Any]:
+        return calculate_token_stats(self.column_config, self.df)
+
+
+class LLMCodeColumnStatisticsCalculator(LLMTextColumnStatisticsCalculator): ...
+
+
+class LLMStructuredColumnStatisticsCalculator(LLMTextColumnStatisticsCalculator): ...
+
+
+class LLMJudgedColumnStatisticsCalculator(LLMTextColumnStatisticsCalculator): ...
+
+
+class SamplerColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
+    def calculate_sampler_distribution(self) -> dict[str, Any]:
+        make_dist, dist_type = False, ColumnDistributionType.OTHER
+        if self.column_config.sampler_type in [SamplerType.CATEGORY, SamplerType.SUBCATEGORY]:
+            make_dist, dist_type = True, ColumnDistributionType.CATEGORICAL
+        elif is_numerical_sampler_type(self.column_config.sampler_type):
+            make_dist, dist_type = True, ColumnDistributionType.NUMERICAL
+        return (
+            {
+                "sampler_type": SamplerType(self.column_config.sampler_type),
+                **calculate_column_distribution(self.column_config.name, self.df, dist_type),
+            }
+            if make_dist
+            else {
+                "sampler_type": SamplerType(self.column_config.sampler_type),
+                "distribution_type": dist_type,
+                "distribution": None,
+            }
+        )
+
+
+class SeedDatasetColumnStatisticsCalculator(GeneralColumnStatisticsCalculator): ...
+
+
+class ValidationColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
+    def calculate_validation_column_info(self) -> dict[str, Any]:
+        return calculate_validation_column_info(self.column_config.name, self.df)
+
+
+class ExpressionColumnStatisticsCalculator(GeneralColumnStatisticsCalculator): ...
+
+
+ColumnStatisticsCalculatorT: TypeAlias = (
+    ExpressionColumnStatisticsCalculator
+    | ValidationColumnStatisticsCalculator
+    | GeneralColumnStatisticsCalculator
+    | LLMCodeColumnStatisticsCalculator
+    | LLMJudgedColumnStatisticsCalculator
+    | LLMStructuredColumnStatisticsCalculator
+    | LLMTextColumnStatisticsCalculator
+    | SamplerColumnStatisticsCalculator
+    | SeedDatasetColumnStatisticsCalculator
+)
+DEFAULT_COLUMN_STATISTICS_CALCULATOR_MAP = {
+    DataDesignerColumnType.EXPRESSION: ExpressionColumnStatisticsCalculator,
+    DataDesignerColumnType.VALIDATION: ValidationColumnStatisticsCalculator,
+    DataDesignerColumnType.LLM_CODE: LLMCodeColumnStatisticsCalculator,
+    DataDesignerColumnType.LLM_JUDGE: LLMJudgedColumnStatisticsCalculator,
+    DataDesignerColumnType.LLM_STRUCTURED: LLMStructuredColumnStatisticsCalculator,
+    DataDesignerColumnType.LLM_TEXT: LLMTextColumnStatisticsCalculator,
+    DataDesignerColumnType.SAMPLER: SamplerColumnStatisticsCalculator,
+    DataDesignerColumnType.SEED_DATASET: SeedDatasetColumnStatisticsCalculator,
+}
+
+
+def get_column_statistics_calculator(column_type: DataDesignerColumnType) -> ColumnStatisticsCalculatorT:
+    return DEFAULT_COLUMN_STATISTICS_CALCULATOR_MAP.get(column_type, GeneralColumnStatisticsCalculator)

data_designer/engine/analysis/dataset_profiler.py

@@ -0,0 +1,149 @@
+# 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 collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING
+
+from pydantic import Field, field_validator
+
+from data_designer.config.analysis.column_profilers import ColumnProfilerConfigT
+from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
+from data_designer.config.base import ConfigBase
+from data_designer.config.column_configs import SingleColumnConfig
+from data_designer.config.column_types import ColumnConfigT
+from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame, ColumnProfiler
+from data_designer.engine.analysis.column_statistics import get_column_statistics_calculator
+from data_designer.engine.analysis.errors import DatasetProfilerConfigurationError
+from data_designer.engine.analysis.utils.column_statistics_calculations import has_pyarrow_backend
+from data_designer.engine.dataset_builders.multi_column_configs import DatasetBuilderColumnConfigT, MultiColumnConfig
+from data_designer.engine.registry.data_designer_registry import DataDesignerRegistry
+from data_designer.engine.resources.resource_provider import ResourceProvider
+from data_designer.lazy_heavy_imports import pa, pd
+
+if TYPE_CHECKING:
+    import pandas as pd
+    import pyarrow as pa
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetProfilerConfig(ConfigBase):
+    column_configs: Sequence[DatasetBuilderColumnConfigT] = Field(..., min_length=1)
+    column_profiler_configs: Sequence[ColumnProfilerConfigT] | None = None
+
+    @field_validator("column_configs")
+    def flatten_and_validate_column_configs(cls, v: list[DatasetBuilderColumnConfigT]) -> list[ColumnConfigT]:
+        column_configs = []
+        for config in v:
+            if isinstance(config, SingleColumnConfig) and not config.drop:
+                column_configs.append(config)
+            elif isinstance(config, MultiColumnConfig):
+                column_configs.extend([c for c in config.columns if not c.drop])
+        if len(column_configs) == 0:
+            raise DatasetProfilerConfigurationError("All columns were dropped!")
+        return column_configs
+
+
+class DataDesignerDatasetProfiler:
+    def __init__(self, config: DatasetProfilerConfig, resource_provider: ResourceProvider):
+        self.config = config
+        self.resource_provider = resource_provider
+        self._validate_column_profiler_configs()
+
+    @cached_property
+    def column_names_from_configs(self) -> list[str]:
+        return [c.name for c in self.config.column_configs]
+
+    @cached_property
+    def registry(self) -> DataDesignerRegistry:
+        return DataDesignerRegistry()
+
+    def profile_dataset(
+        self,
+        target_num_records: int,
+        dataset: pd.DataFrame,
+    ) -> DatasetProfilerResults:
+        logger.info("📐 Measuring dataset column statistics:")
+
+        self._validate_schema_consistency(list(dataset.columns))
+        dataset = self._convert_to_pyarrow_backend_if_needed(dataset)
+
+        column_statistics = []
+        for c in self.config.column_configs:
+            logger.info(f"  |-- {c.get_column_emoji()} column: '{c.name}'")
+            column_statistics.append(
+                get_column_statistics_calculator(c.column_type)(
+                    column_config_with_df=ColumnConfigWithDataFrame(column_config=c, df=dataset)
+                ).calculate()
+            )
+
+        column_profiles = []
+        for profiler_config in self.config.column_profiler_configs or []:
+            profiler = self._create_column_profiler(profiler_config)
+            applicable_column_types = profiler.get_applicable_column_types()
+            for c in self.config.column_configs:
+                if c.column_type in applicable_column_types:
+                    params = ColumnConfigWithDataFrame(column_config=c, df=dataset)
+                    column_profiles.append(profiler.profile(params))
+            if len(column_profiles) == 0:
+                logger.warning(
+                    f"⚠️ No applicable column types found for the '{profiler.name}' profiler. "
+                    f"This profiler is applicable to the following column types: {applicable_column_types}"
+                )
+
+        return DatasetProfilerResults(
+            num_records=len(dataset),
+            target_num_records=target_num_records,
+            side_effect_column_names=list(set(dataset.columns) - set(self.column_names_from_configs)),
+            column_statistics=column_statistics,
+            column_profiles=column_profiles if column_profiles else None,
+        )
+
+    def _convert_to_pyarrow_backend_if_needed(self, dataset: pd.DataFrame) -> pd.DataFrame:
+        if not has_pyarrow_backend(dataset):
+            try:
+                dataset = pa.Table.from_pandas(dataset).to_pandas(types_mapper=pd.ArrowDtype)
+            except Exception as e:
+                # For ArrowTypeError, the second arg contains the more informative message
+                if isinstance(e, pa.lib.ArrowTypeError) and len(e.args) > 1:
+                    error_msg = str(e.args[1])
+                else:
+                    error_msg = str(e)
+                for col in dataset.columns:
+                    # Make sure column names are clear in the error message
+                    error_msg = error_msg.replace(col, f"'{col}'")
+                logger.warning("⚠️ Unable to convert the dataset to a PyArrow backend")
+                logger.warning(f"  |-- Conversion Error Message: {error_msg}")
+                logger.warning("  |-- This is often due to at least one column having mixed data types")
+                logger.warning(
+                    "  |-- Note: Reported data types will be inferred from the first non-null value of each column"
+                )
+        return dataset
+
+    def _create_column_profiler(self, profiler_config: ColumnProfilerConfigT) -> ColumnProfiler:
+        return self.registry.column_profilers.get_for_config_type(type(profiler_config))(
+            config=profiler_config, resource_provider=self.resource_provider
+        )
+
+    def _validate_column_profiler_configs(self) -> None:
+        if self.config.column_profiler_configs:
+            if self.resource_provider.model_registry is None:
+                raise DatasetProfilerConfigurationError("Model registry is required for column profiler configs")
+            self._validate_model_configs()
+
+    def _validate_model_configs(self) -> None:
+        aliases = [alias for alias in self.resource_provider.model_registry.model_configs.keys()]
+        for column_config in self.config.column_configs:
+            if hasattr(column_config, "model_alias") and column_config.model_alias not in aliases:
+                raise DatasetProfilerConfigurationError(
+                    f"Model config '{column_config.model_alias}' not found in model configs"
+                )
+
+    def _validate_schema_consistency(self, dataset_column_names: list[str]) -> None:
+        for column_name in self.column_names_from_configs:
+            if column_name not in dataset_column_names:
+                raise DatasetProfilerConfigurationError(f"Column '{column_name}' not found in dataset")

data_designer/engine/analysis/errors.py

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