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

data_designer/engine/analysis/utils/column_statistics_calculations.py

@@ -20,10 +20,8 @@ from data_designer.config.analysis.column_statistics import (
 )
 from data_designer.config.column_configs import (
     LLMTextColumnConfig,
-    SingleColumnConfig,
-    ValidationColumnConfig,
 )
-from data_designer.engine.column_generators.generators.llm_generators import (
+from data_designer.engine.column_generators.utils.prompt_renderer import (
     PromptType,
     RecordBasedPromptRenderer,
     create_response_recipe,
@@ -39,41 +37,54 @@ logger = logging.getLogger(__name__)
 
 
 def calculate_column_distribution(
-    column_config: SingleColumnConfig, df: pd.DataFrame, distribution_type: ColumnDistributionType
+    column_name: str, df: pd.DataFrame, distribution_type: ColumnDistributionType
 ) -> dict[str, CategoricalDistribution | NumericalDistribution | MissingValue | None]:
     distribution_type = ColumnDistributionType(distribution_type)
     try:
         if distribution_type == ColumnDistributionType.CATEGORICAL:
             return {
                 "distribution_type": ColumnDistributionType.CATEGORICAL,
-                "distribution": CategoricalDistribution.from_series(df[column_config.name]),
+                "distribution": CategoricalDistribution.from_series(df[column_name]),
             }
 
         if distribution_type == ColumnDistributionType.NUMERICAL:
             return {
                 "distribution_type": ColumnDistributionType.NUMERICAL,
-                "distribution": NumericalDistribution.from_series(df[column_config.name]),
+                "distribution": NumericalDistribution.from_series(df[column_name]),
             }
     except Exception as e:
-        logger.warning(f"{WARNING_PREFIX} failed to calculate column distribution for '{column_config.name}' {e}")
+        logger.warning(f"{WARNING_PREFIX} failed to calculate column distribution for '{column_name}' {e}")
         return {
             "distribution_type": ColumnDistributionType.UNKNOWN,
             "distribution": MissingValue.CALCULATION_FAILED,
         }
 
 
-def calculate_general_column_info(column_config: SingleColumnConfig, df: pd.DataFrame) -> dict[str, Any]:
+def calculate_general_column_info(column_name: str, df: pd.DataFrame) -> dict[str, Any]:
     try:
-        _df = pd.DataFrame(df[column_config.name].apply(ensure_hashable))
+        _df = pd.DataFrame(df[column_name].apply(ensure_hashable))
+
+        if has_pyarrow_backend(df):
+            pyarrow_dtype = str(df[column_name].dtype.pyarrow_dtype)
+            simple_dtype = convert_pyarrow_dtype_to_simple_dtype(df[column_name].dtype.pyarrow_dtype)
+        else:
+            # We do not log a warning at the column-level because it would be too noisy.
+            # However, there is a logged warning at the dataset-profiler level.
+            try:
+                simple_dtype = get_column_data_type_from_first_non_null_value(column_name, df)
+            except Exception:
+                simple_dtype = MissingValue.CALCULATION_FAILED
+            pyarrow_dtype = "n/a"
+
         return {
-            "pyarrow_dtype": str(df[column_config.name].dtype.pyarrow_dtype),
-            "simple_dtype": convert_pyarrow_dtype_to_simple_dtype(df[column_config.name].dtype.pyarrow_dtype),
-            "num_records": len(_df[column_config.name]),
-            "num_null": _df[column_config.name].isnull().sum(),
-            "num_unique": _df[column_config.name].nunique(),
+            "pyarrow_dtype": pyarrow_dtype,
+            "simple_dtype": simple_dtype,
+            "num_records": len(_df[column_name]),
+            "num_null": _df[column_name].isnull().sum(),
+            "num_unique": _df[column_name].nunique(),
         }
     except Exception as e:
-        logger.warning(f"{WARNING_PREFIX} failed to calculate general column info for '{column_config.name}': {e}")
+        logger.warning(f"{WARNING_PREFIX} failed to calculate general column info for '{column_name}': {e}")
         return {
             "pyarrow_dtype": MissingValue.CALCULATION_FAILED,
             "simple_dtype": MissingValue.CALCULATION_FAILED,
@@ -83,7 +94,7 @@ def calculate_general_column_info(column_config: SingleColumnConfig, df: pd.Data
         }
 
 
-def calculate_prompt_token_stats(
+def calculate_input_token_stats(
     column_config: LLMTextColumnConfig, df: pd.DataFrame
 ) -> dict[str, float | MissingValue]:
     try:
@@ -100,22 +111,20 @@ def calculate_prompt_token_stats(
             concatenated_prompt = str(system_prompt + "\n\n" + prompt)
             num_tokens.append(len(TOKENIZER.encode(concatenated_prompt, disallowed_special=())))
     except Exception as e:
-        logger.warning(
-            f"{WARNING_PREFIX} failed to calculate prompt token stats for column {column_config.name!r}: {e}"
-        )
+        logger.warning(f"{WARNING_PREFIX} failed to calculate input token stats for column {column_config.name!r}: {e}")
         return {
-            "prompt_tokens_mean": MissingValue.CALCULATION_FAILED,
-            "prompt_tokens_median": MissingValue.CALCULATION_FAILED,
-            "prompt_tokens_stddev": MissingValue.CALCULATION_FAILED,
+            "input_tokens_mean": MissingValue.CALCULATION_FAILED,
+            "input_tokens_median": MissingValue.CALCULATION_FAILED,
+            "input_tokens_stddev": MissingValue.CALCULATION_FAILED,
         }
     return {
-        "prompt_tokens_mean": np.mean(num_tokens),
-        "prompt_tokens_median": np.median(num_tokens),
-        "prompt_tokens_stddev": np.std(num_tokens),
+        "input_tokens_mean": np.mean(num_tokens),
+        "input_tokens_median": np.median(num_tokens),
+        "input_tokens_stddev": np.std(num_tokens),
     }
 
 
-def calculate_completion_token_stats(
+def calculate_output_token_stats(
     column_config: LLMTextColumnConfig, df: pd.DataFrame
 ) -> dict[str, float | MissingValue]:
     try:
@@ -123,34 +132,32 @@ def calculate_completion_token_stats(
             lambda value: len(TOKENIZER.encode(str(value), disallowed_special=()))
         )
         return {
-            "completion_tokens_mean": tokens_per_record.mean(),
-            "completion_tokens_median": tokens_per_record.median(),
-            "completion_tokens_stddev": tokens_per_record.std(),
+            "output_tokens_mean": tokens_per_record.mean(),
+            "output_tokens_median": tokens_per_record.median(),
+            "output_tokens_stddev": tokens_per_record.std(),
         }
     except Exception as e:
-        logger.warning(
-            f"{WARNING_PREFIX} failed to calculate completion token stats for column {column_config.name}: {e}"
-        )
+        logger.warning(f"{WARNING_PREFIX} failed to calculate output token stats for column {column_config.name}: {e}")
         return {
-            "completion_tokens_mean": MissingValue.CALCULATION_FAILED,
-            "completion_tokens_median": MissingValue.CALCULATION_FAILED,
-            "completion_tokens_stddev": MissingValue.CALCULATION_FAILED,
+            "output_tokens_mean": MissingValue.CALCULATION_FAILED,
+            "output_tokens_median": MissingValue.CALCULATION_FAILED,
+            "output_tokens_stddev": MissingValue.CALCULATION_FAILED,
         }
 
 
 def calculate_token_stats(column_config: LLMTextColumnConfig, df: pd.DataFrame) -> dict[str, float | MissingValue]:
     return {
-        **calculate_prompt_token_stats(column_config, df),
-        **calculate_completion_token_stats(column_config, df),
+        **calculate_input_token_stats(column_config, df),
+        **calculate_output_token_stats(column_config, df),
     }
 
 
-def calculate_validation_column_info(column_config: ValidationColumnConfig, df: pd.DataFrame) -> dict[str, Any]:
+def calculate_validation_column_info(column_name: str, df: pd.DataFrame) -> dict[str, Any]:
     try:
-        return {"num_valid_records": df[column_config.name].apply(lambda x: ensure_boolean(x["is_valid"])).sum()}
+        return {"num_valid_records": df[column_name].apply(lambda x: ensure_boolean(x["is_valid"])).sum()}
     except Exception as e:
         logger.warning(
-            f"{WARNING_PREFIX} failed to calculate code validation column info for column {column_config.name}: {e}"
+            f"{WARNING_PREFIX} failed to calculate code validation column info for column {column_name}: {e}"
         )
         return {"num_valid_records": MissingValue.CALCULATION_FAILED}
 
@@ -160,22 +167,33 @@ def convert_pyarrow_dtype_to_simple_dtype(pyarrow_dtype: pa.DataType) -> str:
         return f"list[{convert_pyarrow_dtype_to_simple_dtype(pyarrow_dtype.value_type)}]"
     if isinstance(pyarrow_dtype, pa.StructType):
         return "dict"
-    pyarrow_dtype_str = str(pyarrow_dtype)
-    if "int" in pyarrow_dtype_str:
+    return convert_to_simple_dtype(str(pyarrow_dtype))
+
+
+def convert_to_simple_dtype(dtype: str) -> str:
+    if "int" in dtype:
         return "int"
-    if "double" in pyarrow_dtype_str:
+    if "double" in dtype:
         return "float"
-    if "float" in pyarrow_dtype_str:
+    if "float" in dtype:
         return "float"
-    if "string" in pyarrow_dtype_str:
+    if "str" in dtype:
         return "string"
-    if "timestamp" in pyarrow_dtype_str:
+    if "timestamp" in dtype:
         return "timestamp"
-    if "time" in pyarrow_dtype_str:
+    if "time" in dtype:
         return "time"
-    if "date" in pyarrow_dtype_str:
+    if "date" in dtype:
         return "date"
-    return pyarrow_dtype_str
+    return dtype
+
+
+def get_column_data_type_from_first_non_null_value(column_name: str, df: pd.DataFrame) -> str:
+    df_no_nulls = df[column_name].dropna()
+    if len(df_no_nulls) == 0:
+        return MissingValue.CALCULATION_FAILED
+    dtype = type(df_no_nulls.iloc[0]).__name__
+    return convert_to_simple_dtype(dtype)
 
 
 def ensure_hashable(x: Any) -> str:
@@ -207,3 +225,7 @@ def ensure_boolean(v: bool | str | int | None) -> bool:
     if v is None:
         return False
     raise ValueError(f"Invalid boolean value: {v}")
+
+
+def has_pyarrow_backend(df: pd.DataFrame) -> bool:
+    return all(isinstance(dtype, pd.ArrowDtype) for dtype in df.dtypes)

data_designer/engine/analysis/utils/judge_score_processing.py

@@ -3,7 +3,7 @@
 
 import logging
 from collections import defaultdict
-from typing import Any, Optional, Union
+from typing import Any
 
 import pandas as pd
 
@@ -21,7 +21,7 @@ logger = logging.getLogger(__name__)
 
 def extract_judge_score_distributions(
     column_config: LLMJudgeColumnConfig, df: pd.DataFrame
-) -> Union[JudgeScoreDistributions, MissingValue]:
+) -> JudgeScoreDistributions | MissingValue:
     scores = defaultdict(list)
     reasoning = defaultdict(list)
 
@@ -32,7 +32,7 @@ def extract_judge_score_distributions(
 
     for score in column_config.scores:
         is_numerical = True
-        name = score.name.lower()
+        name = score.name
         for results in df[column_config.name]:
             try:
                 score = results[name].get("score", None)
@@ -79,10 +79,10 @@ def extract_judge_score_distributions(
 
 
 def sample_scores_and_reasoning(
-    scores: list[Union[int, str]],
+    scores: list[int | str],
     reasoning: list[str],
     num_samples: int,
-    random_seed: Optional[int] = None,
+    random_seed: int | None = None,
 ) -> list[JudgeScoreSample]:
     if len(scores) != len(reasoning):
         raise ValueError("scores and reasoning must have the same length")

data_designer/engine/column_generators/generators/base.py

@@ -1,13 +1,20 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import functools
+import logging
 from abc import ABC, abstractmethod
 from typing import overload
 
 import pandas as pd
 
+from data_designer.config.column_types import COLUMN_TYPE_EMOJI_MAP
+from data_designer.config.models import BaseInferenceParams, ModelConfig
 from data_designer.config.utils.type_helpers import StrEnum
 from data_designer.engine.configurable_task import ConfigurableTask, ConfigurableTaskMetadata, DataT, TaskConfigT
+from data_designer.engine.models.facade import ModelFacade
+
+logger = logging.getLogger(__name__)
 
 
 class GenerationStrategy(StrEnum):
@@ -59,3 +66,30 @@ class FromScratchColumnGenerator(ColumnGenerator[TaskConfigT], ABC):
 
     @abstractmethod
     def generate_from_scratch(self, num_records: int) -> pd.DataFrame: ...
+
+
+class WithModelGeneration:
+    @functools.cached_property
+    def model(self) -> ModelFacade:
+        return self.resource_provider.model_registry.get_model(model_alias=self.config.model_alias)
+
+    @functools.cached_property
+    def model_config(self) -> ModelConfig:
+        return self.resource_provider.model_registry.get_model_config(model_alias=self.config.model_alias)
+
+    @functools.cached_property
+    def inference_parameters(self) -> BaseInferenceParams:
+        return self.model_config.inference_parameters
+
+    def log_pre_generation(self) -> None:
+        emoji = COLUMN_TYPE_EMOJI_MAP[self.config.column_type]
+        logger.info(f"{emoji} Preparing {self.config.column_type} column generation")
+        logger.info(f"  |-- column name: {self.config.name!r}")
+        logger.info(f"  |-- model config:\n{self.model_config.model_dump_json(indent=4)}")
+        if self.model_config.provider is None:
+            logger.info(f"  |-- default model provider: {self._get_provider_name()!r}")
+
+    def _get_provider_name(self) -> str:
+        model_alias = self.model_config.alias
+        provider = self.resource_provider.model_registry.get_model_provider(model_alias=model_alias)
+        return provider.name

data_designer/engine/column_generators/generators/embedding.py

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+from pydantic import BaseModel, computed_field
+
+from data_designer.config.column_configs import EmbeddingColumnConfig
+from data_designer.engine.column_generators.generators.base import (
+    ColumnGenerator,
+    GenerationStrategy,
+    GeneratorMetadata,
+    WithModelGeneration,
+)
+from data_designer.engine.processing.utils import deserialize_json_values, parse_list_string
+from data_designer.engine.resources.resource_provider import ResourceType
+
+
+class EmbeddingGenerationResult(BaseModel):
+    embeddings: list[list[float]]
+
+    @computed_field
+    def num_embeddings(self) -> int:
+        return len(self.embeddings)
+
+    @computed_field
+    def dimension(self) -> int:
+        return len(self.embeddings[0]) if len(self.embeddings) > 0 else 0
+
+
+class EmbeddingCellGenerator(WithModelGeneration, ColumnGenerator[EmbeddingColumnConfig]):
+    @staticmethod
+    def metadata() -> GeneratorMetadata:
+        return GeneratorMetadata(
+            name="embedding_cell_generator",
+            description="Generate embeddings for a text column.",
+            generation_strategy=GenerationStrategy.CELL_BY_CELL,
+            required_resources=[ResourceType.MODEL_REGISTRY],
+        )
+
+    def generate(self, data: dict) -> dict:
+        deserialized_record = deserialize_json_values(data)
+        input_texts = parse_list_string(deserialized_record[self.config.target_column])
+        embeddings = self.model.generate_text_embeddings(input_texts=input_texts)
+        data[self.config.name] = EmbeddingGenerationResult(embeddings=embeddings).model_dump(mode="json")
+        return data

data_designer/engine/column_generators/generators/{llm_generators.py → llm_completion.py}

@@ -10,43 +10,41 @@ from data_designer.config.column_configs import (
     LLMStructuredColumnConfig,
     LLMTextColumnConfig,
 )
-from data_designer.config.column_types import COLUMN_TYPE_EMOJI_MAP
-from data_designer.config.models import InferenceParameters, ModelConfig
 from data_designer.config.utils.constants import REASONING_TRACE_COLUMN_POSTFIX
 from data_designer.engine.column_generators.generators.base import (
     ColumnGenerator,
     GenerationStrategy,
     GeneratorMetadata,
+    WithModelGeneration,
 )
 from data_designer.engine.column_generators.utils.prompt_renderer import (
     PromptType,
     RecordBasedPromptRenderer,
     create_response_recipe,
 )
-from data_designer.engine.models.facade import ModelFacade
 from data_designer.engine.models.recipes.base import ResponseRecipe
 from data_designer.engine.processing.utils import deserialize_json_values
 from data_designer.engine.resources.resource_provider import ResourceType
 
-DEFAULT_MAX_CONVERSATION_RESTARTS = 5
-DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS = 0
+logger = logging.getLogger(__name__)
 
 
-logger = logging.getLogger(__name__)
+DEFAULT_MAX_CONVERSATION_RESTARTS = 5
+DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS = 0
 
 
-class WithLLMGeneration:
+class WithChatCompletionGeneration(WithModelGeneration):
     @functools.cached_property
-    def model(self) -> ModelFacade:
-        return self.resource_provider.model_registry.get_model(model_alias=self.config.model_alias)
+    def response_recipe(self) -> ResponseRecipe:
+        return create_response_recipe(self.config, self.model_config)
 
-    @functools.cached_property
-    def model_config(self) -> ModelConfig:
-        return self.resource_provider.model_registry.get_model_config(model_alias=self.config.model_alias)
+    @property
+    def max_conversation_correction_steps(self) -> int:
+        return DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS
 
-    @functools.cached_property
-    def inference_parameters(self) -> InferenceParameters:
-        return self.model_config.inference_parameters
+    @property
+    def max_conversation_restarts(self) -> int:
+        return DEFAULT_MAX_CONVERSATION_RESTARTS
 
     @functools.cached_property
     def prompt_renderer(self) -> RecordBasedPromptRenderer:
@@ -59,18 +57,6 @@ class WithLLMGeneration:
             },
         )
 
-    @functools.cached_property
-    def response_recipe(self) -> ResponseRecipe:
-        return create_response_recipe(self.config, self.model_config)
-
-    @property
-    def max_conversation_correction_steps(self) -> int:
-        return DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS
-
-    @property
-    def max_conversation_restarts(self) -> int:
-        return DEFAULT_MAX_CONVERSATION_RESTARTS
-
     def generate(self, data: dict) -> dict:
         deserialized_record = deserialize_json_values(data)
 
@@ -96,7 +82,6 @@ class WithLLMGeneration:
             max_correction_steps=self.max_conversation_correction_steps,
             max_conversation_restarts=self.max_conversation_restarts,
             purpose=f"running generation for column '{self.config.name}'",
-            **self.inference_parameters.generate_kwargs,
         )
 
         data[self.config.name] = deserialize_json_values(self.response_recipe.serialize_output(response))
@@ -106,21 +91,8 @@ class WithLLMGeneration:
 
         return data
 
-    def log_pre_generation(self) -> None:
-        emoji = COLUMN_TYPE_EMOJI_MAP[self.config.column_type]
-        logger.info(f"{emoji} Preparing {self.config.column_type} column generation")
-        logger.info(f"  |-- column name: {self.config.name!r}")
-        logger.info(f"  |-- model config:\n{self.model_config.model_dump_json(indent=4)}")
-        if self.model_config.provider is None:
-            logger.info(f"  |-- default model provider: {self._get_provider_name()!r}")
-
-    def _get_provider_name(self) -> str:
-        model_alias = self.model_config.alias
-        provider = self.resource_provider.model_registry.get_model_provider(model_alias=model_alias)
-        return provider.name
-
 
-class LLMTextCellGenerator(WithLLMGeneration, ColumnGenerator[LLMTextColumnConfig]):
+class LLMTextCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMTextColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
@@ -131,7 +103,7 @@ class LLMTextCellGenerator(WithLLMGeneration, ColumnGenerator[LLMTextColumnConfi
         )
 
 
-class LLMCodeCellGenerator(WithLLMGeneration, ColumnGenerator[LLMCodeColumnConfig]):
+class LLMCodeCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMCodeColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
@@ -142,7 +114,7 @@ class LLMCodeCellGenerator(WithLLMGeneration, ColumnGenerator[LLMCodeColumnConfi
         )
 
 
-class LLMStructuredCellGenerator(WithLLMGeneration, ColumnGenerator[LLMStructuredColumnConfig]):
+class LLMStructuredCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMStructuredColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
@@ -153,7 +125,7 @@ class LLMStructuredCellGenerator(WithLLMGeneration, ColumnGenerator[LLMStructure
         )
 
 
-class LLMJudgeCellGenerator(WithLLMGeneration, ColumnGenerator[LLMJudgeColumnConfig]):
+class LLMJudgeCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMJudgeColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
@@ -163,10 +135,6 @@ class LLMJudgeCellGenerator(WithLLMGeneration, ColumnGenerator[LLMJudgeColumnCon
             required_resources=[ResourceType.MODEL_REGISTRY],
         )
 
-    @property
-    def max_conversation_correction_steps(self) -> int:
-        return DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS
-
     @property
     def max_conversation_restarts(self) -> int:
         return 2 * DEFAULT_MAX_CONVERSATION_RESTARTS

data_designer/engine/column_generators/registry.py

@@ -3,6 +3,7 @@
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.column_configs import (
+    EmbeddingColumnConfig,
     ExpressionColumnConfig,
     LLMCodeColumnConfig,
     LLMJudgeColumnConfig,
@@ -12,8 +13,9 @@ from data_designer.config.column_configs import (
 )
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from data_designer.engine.column_generators.generators.embedding import EmbeddingCellGenerator
 from data_designer.engine.column_generators.generators.expression import ExpressionColumnGenerator
-from data_designer.engine.column_generators.generators.llm_generators import (
+from data_designer.engine.column_generators.generators.llm_completion import (
     LLMCodeCellGenerator,
     LLMJudgeCellGenerator,
     LLMStructuredCellGenerator,
@@ -40,11 +42,11 @@ def create_default_column_generator_registry(with_plugins: bool = True) -> Colum
     registry.register(DataDesignerColumnType.LLM_CODE, LLMCodeCellGenerator, LLMCodeColumnConfig)
     registry.register(DataDesignerColumnType.LLM_JUDGE, LLMJudgeCellGenerator, LLMJudgeColumnConfig)
     registry.register(DataDesignerColumnType.EXPRESSION, ExpressionColumnGenerator, ExpressionColumnConfig)
+    registry.register(DataDesignerColumnType.EMBEDDING, EmbeddingCellGenerator, EmbeddingColumnConfig)
     registry.register(DataDesignerColumnType.SAMPLER, SamplerColumnGenerator, SamplerMultiColumnConfig)
     registry.register(DataDesignerColumnType.SEED_DATASET, SeedDatasetColumnGenerator, SeedDatasetMultiColumnConfig)
     registry.register(DataDesignerColumnType.VALIDATION, ValidationColumnGenerator, ValidationColumnConfig)
     registry.register(DataDesignerColumnType.LLM_STRUCTURED, LLMStructuredCellGenerator, LLMStructuredColumnConfig)
-
     if with_plugins:
         for plugin in PluginRegistry().get_plugins(PluginType.COLUMN_GENERATOR):
             registry.register(

data_designer/engine/column_generators/utils/judge_score_factory.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import Enum
-from typing import Type
 
 from pydantic import BaseModel, ConfigDict, Field, create_model
 
@@ -19,7 +18,7 @@ class BaseJudgeResponse(BaseModel):
     reasoning: str = Field(..., description="Reasoning for the assigned score.")
 
 
-def _stringify_scoring(options: dict, enum_type: Type[Enum]) -> str:
+def _stringify_scoring(options: dict, enum_type: type[Enum]) -> str:
     """Convert score descriptions into a single text block."""
     list_block = "\n".join(
         [SCORING_FORMAT.format(score=score, description=description) for score, description in options.items()]
@@ -27,7 +26,7 @@ def _stringify_scoring(options: dict, enum_type: Type[Enum]) -> str:
     return SCORE_FIELD_DESCRIPTION_FORMAT.format(enum_name=enum_type.__name__, scoring=list_block)
 
 
-def create_judge_response_model(score: Score) -> Type[BaseJudgeResponse]:
+def create_judge_response_model(score: Score) -> type[BaseJudgeResponse]:
     """Create a JudgeResponse data type."""
     enum_members = {}
     for option in score.options.keys():
@@ -46,12 +45,12 @@ def create_judge_response_model(score: Score) -> Type[BaseJudgeResponse]:
 
 
 def create_judge_structured_output_model(
-    judge_responses: list[Type[BaseJudgeResponse]],
-) -> Type[BaseModel]:
+    judge_responses: list[type[BaseJudgeResponse]],
+) -> type[BaseModel]:
     """Create a JudgeStructuredOutput class dynamically."""
     return create_model(
         "JudgeStructuredOutput",
         __doc__=f"Response schema for scores with the following names: {[response.__name__ for response in judge_responses]}.",
         __base__=BaseModel,
-        **{response.__name__.lower(): (response, ...) for response in judge_responses},
+        **{response.__name__: (response, ...) for response in judge_responses},
     )

data_designer/engine/configurable_task.py

@@ -3,7 +3,7 @@
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Generic, Type, TypeVar, get_origin
+from typing import Generic, TypeVar, get_origin
 
 import pandas as pd
 
@@ -30,7 +30,7 @@ class ConfigurableTask(ABC, Generic[TaskConfigT]):
         self._initialize()
 
     @classmethod
-    def get_config_type(cls) -> Type[TaskConfigT]:
+    def get_config_type(cls) -> type[TaskConfigT]:
         for base in cls.__orig_bases__:
             if hasattr(base, "__args__") and len(base.__args__) == 1:
                 arg = base.__args__[0]

data_designer/engine/dataset_builders/artifact_storage.py

@@ -7,7 +7,6 @@ import shutil
 from datetime import datetime
 from functools import cached_property
 from pathlib import Path
-from typing import Union
 
 import pandas as pd
 from pydantic import BaseModel, field_validator, model_validator
@@ -77,7 +76,7 @@ class ArtifactStorage(BaseModel):
         return self.base_dataset_path / self.processors_outputs_folder_name
 
     @field_validator("artifact_path")
-    def validate_artifact_path(cls, v: Union[Path, str]) -> Path:
+    def validate_artifact_path(cls, v: Path | str) -> Path:
         v = Path(v)
         if not v.is_dir():
             raise ArtifactStorageError("Artifact path must exist and be a directory")

data_designer/engine/dataset_builders/column_wise_builder.py

@@ -10,15 +10,18 @@ from typing import Callable
 
 import pandas as pd
 
-from data_designer.config.column_types import ColumnConfigT, column_type_is_llm_generated
+from data_designer.config.column_types import ColumnConfigT, column_type_is_model_generated
 from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
     ProcessorConfig,
     ProcessorType,
 )
-from data_designer.engine.column_generators.generators.base import ColumnGenerator, GenerationStrategy
-from data_designer.engine.column_generators.generators.llm_generators import WithLLMGeneration
+from data_designer.engine.column_generators.generators.base import (
+    ColumnGenerator,
+    GenerationStrategy,
+    WithModelGeneration,
+)
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError, DatasetProcessingError
 from data_designer.engine.dataset_builders.multi_column_configs import (
@@ -72,7 +75,7 @@ class ColumnWiseDatasetBuilder:
 
     @functools.cached_property
     def llm_generated_column_configs(self) -> list[ColumnConfigT]:
-        return [config for config in self.single_column_configs if column_type_is_llm_generated(config.column_type)]
+        return [config for config in self.single_column_configs if column_type_is_model_generated(config.column_type)]
 
     def build(
         self,
@@ -169,10 +172,8 @@ class ColumnWiseDatasetBuilder:
 
     def _run_cell_by_cell_generator(self, generator: ColumnGenerator) -> None:
         max_workers = MAX_CONCURRENCY_PER_NON_LLM_GENERATOR
-        if isinstance(generator, WithLLMGeneration):
+        if isinstance(generator, WithModelGeneration):
             max_workers = generator.inference_parameters.max_parallel_requests
-        elif hasattr(generator.config, "max_parallel_requests"):
-            max_workers = generator.config.max_parallel_requests
         self._fan_out_with_threads(generator, max_workers=max_workers)
 
     def _run_full_column_generator(self, generator: ColumnGenerator) -> None:
@@ -180,12 +181,12 @@ class ColumnWiseDatasetBuilder:
         self.batch_manager.update_records(df.to_dict(orient="records"))
 
     def _run_model_health_check_if_needed(self) -> bool:
-        if any(column_type_is_llm_generated(config.column_type) for config in self.single_column_configs):
+        if any(column_type_is_model_generated(config.column_type) for config in self.single_column_configs):
             self._resource_provider.model_registry.run_health_check(
-                set(config.model_alias for config in self.llm_generated_column_configs)
+                list(set(config.model_alias for config in self.llm_generated_column_configs))
             )
 
-    def _fan_out_with_threads(self, generator: WithLLMGeneration, max_workers: int) -> None:
+    def _fan_out_with_threads(self, generator: WithModelGeneration, max_workers: int) -> None:
         if generator.generation_strategy != GenerationStrategy.CELL_BY_CELL:
             raise DatasetGenerationError(
                 f"Generator {generator.metadata().name} is not a {GenerationStrategy.CELL_BY_CELL} "