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

data_designer/engine/analysis/column_profilers/judge_score_profiler.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 
 import logging
 import random
-from typing import Union
 
 from data_designer.config.analysis.column_profilers import (
     JudgeScoreProfilerConfig,
@@ -69,7 +68,7 @@ class JudgeScoreProfiler(ColumnProfiler[JudgeScoreProfilerConfig]):
             )
 
         for score in column_config.scores:
-            score_name = score.name.lower()
+            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],
@@ -96,7 +95,7 @@ class JudgeScoreProfiler(ColumnProfiler[JudgeScoreProfilerConfig]):
         name: str,
         sample: list[JudgeScoreSample],
         histogram: CategoricalHistogramData,
-        distribution: Union[CategoricalDistribution, NumericalDistribution, MissingValue],
+        distribution: CategoricalDistribution | NumericalDistribution | MissingValue,
         distribution_type: ColumnDistributionType,
     ) -> JudgeScoreSummary:
         if isinstance(distribution, MissingValue) or not sample:

data_designer/engine/analysis/column_statistics.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Type, TypeAlias, Union
+from typing import Any, TypeAlias
 
 import pandas as pd
 from pydantic import BaseModel
@@ -41,7 +41,7 @@ class GeneralColumnStatisticsCalculator(BaseModel):
         return self.column_config_with_df.df
 
     @property
-    def column_statistics_type(self) -> Type[ColumnStatisticsT]:
+    def column_statistics_type(self) -> type[ColumnStatisticsT]:
         return DEFAULT_COLUMN_STATISTICS_MAP.get(self.column_config.column_type, GeneralColumnStatistics)
 
     def calculate(self) -> Self:
@@ -59,7 +59,7 @@ class GeneralColumnStatisticsCalculator(BaseModel):
         )
 
     def calculate_general_column_info(self) -> dict[str, Any]:
-        return calculate_general_column_info(self.column_config, self.df)
+        return calculate_general_column_info(self.column_config.name, self.df)
 
     def __repr__(self) -> str:
         params = []
@@ -93,7 +93,7 @@ class SamplerColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
         return (
             {
                 "sampler_type": SamplerType(self.column_config.sampler_type),
-                **calculate_column_distribution(self.column_config, self.df, dist_type),
+                **calculate_column_distribution(self.column_config.name, self.df, dist_type),
             }
             if make_dist
             else {
@@ -109,23 +109,23 @@ class SeedDatasetColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
 
 class ValidationColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
     def calculate_validation_column_info(self) -> dict[str, Any]:
-        return calculate_validation_column_info(self.column_config, self.df)
+        return calculate_validation_column_info(self.column_config.name, self.df)
 
 
 class ExpressionColumnStatisticsCalculator(GeneralColumnStatisticsCalculator): ...
 
 
-ColumnStatisticsCalculatorT: TypeAlias = Union[
-    ExpressionColumnStatisticsCalculator,
-    ValidationColumnStatisticsCalculator,
-    GeneralColumnStatisticsCalculator,
-    LLMCodeColumnStatisticsCalculator,
-    LLMJudgedColumnStatisticsCalculator,
-    LLMStructuredColumnStatisticsCalculator,
-    LLMTextColumnStatisticsCalculator,
-    SamplerColumnStatisticsCalculator,
-    SeedDatasetColumnStatisticsCalculator,
-]
+ColumnStatisticsCalculatorT: TypeAlias = (
+    ExpressionColumnStatisticsCalculator
+    | ValidationColumnStatisticsCalculator
+    | GeneralColumnStatisticsCalculator
+    | LLMCodeColumnStatisticsCalculator
+    | LLMJudgedColumnStatisticsCalculator
+    | LLMStructuredColumnStatisticsCalculator
+    | LLMTextColumnStatisticsCalculator
+    | SamplerColumnStatisticsCalculator
+    | SeedDatasetColumnStatisticsCalculator
+)
 DEFAULT_COLUMN_STATISTICS_CALCULATOR_MAP = {
     DataDesignerColumnType.EXPRESSION: ExpressionColumnStatisticsCalculator,
     DataDesignerColumnType.VALIDATION: ValidationColumnStatisticsCalculator,

data_designer/engine/analysis/dataset_profiler.py

@@ -6,6 +6,7 @@ from collections.abc import Sequence
 from functools import cached_property
 
 import pandas as pd
+import pyarrow as pa
 from pydantic import Field, field_validator
 
 from data_designer.config.analysis.column_profilers import ColumnProfilerConfigT
@@ -19,10 +20,8 @@ from data_designer.config.column_types import (
 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.dataset_builders.multi_column_configs import (
-    DatasetBuilderColumnConfigT,
-    MultiColumnConfig,
-)
+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
 
@@ -68,6 +67,7 @@ class DataDesignerDatasetProfiler:
         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:
@@ -100,6 +100,27 @@ class DataDesignerDatasetProfiler:
             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

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