sdg-hub 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

sdg_hub/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.5.0'
-__version_tuple__ = version_tuple = (0, 5, 0)
+__version__ = version = '0.6.0'
+__version_tuple__ = version_tuple = (0, 6, 0)
 
 __commit_id__ = commit_id = None

sdg_hub/core/blocks/base.py

@@ -9,13 +9,14 @@ with unified constructor patterns, column handling, and common functionality.
 from abc import ABC, abstractmethod
 from typing import Any, Optional, Union
 
-# Third Party
-from datasets import Dataset
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 from rich.console import Console
 from rich.panel import Panel
 from rich.text import Text
 
+# Third Party
+import pandas as pd
+
 # Local
 from ..utils.error_handling import (
     EmptyDatasetError,
@@ -32,7 +33,7 @@ class BaseBlock(BaseModel, ABC):
     """Base class for all blocks, with standardized patterns and full Pydantic compatibility.
 
     This class defines a unified, configurable base for building composable data processing blocks
-    that operate over HuggingFace Datasets. It supports field-based initialization, validation,
+    that operate over pandas DataFrames. It supports field-based initialization, validation,
     and rich logging for inputs and outputs.
 
     Attributes
@@ -40,9 +41,9 @@ class BaseBlock(BaseModel, ABC):
     block_name : str
         Unique identifier for this block instance.
     input_cols : Union[List[str], Dict[str, Any]]
-        Input columns from the dataset (string, list of strings, or mapping).
+        Input columns from the DataFrame (string, list of strings, or mapping).
     output_cols : Union[List[str], Dict[str, Any]]
-        Output columns to write to the dataset (string, list of strings, or mapping).
+        Output columns to write to the DataFrame (string, list of strings, or mapping).
     """
 
     block_name: str = Field(
@@ -55,7 +56,7 @@ class BaseBlock(BaseModel, ABC):
         None, description="Output columns: str, list, or dict"
     )
 
-    # Allow extra config fields and complex types like Dataset
+    # Allow extra config fields and complex types like DataFrame
     model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
 
     # Normalize input columns before model construction
@@ -101,13 +102,13 @@ class BaseBlock(BaseModel, ABC):
             return dict(cols)
         raise ValueError(f"Invalid column specification: {cols} (type: {type(cols)})")
 
-    def _validate_columns(self, dataset: Dataset) -> None:
-        """Check that all required input columns are present in the dataset.
+    def _validate_columns(self, df: pd.DataFrame) -> None:
+        """Check that all required input columns are present in the DataFrame.
 
         Parameters
         ----------
-        dataset : Dataset
-            HuggingFace dataset to validate against.
+        df : pd.DataFrame
+            DataFrame to validate against.
 
         Raises
         ------
@@ -121,28 +122,29 @@ class BaseBlock(BaseModel, ABC):
             if isinstance(self.input_cols, dict)
             else self.input_cols
         )
+        available_columns = df.columns.tolist()
         missing_columns = [
-            col for col in columns_to_check if col not in dataset.column_names
+            col for col in columns_to_check if col not in available_columns
         ]
         if missing_columns:
             raise MissingColumnError(
                 block_name=self.block_name,
                 missing_columns=missing_columns,
-                available_columns=dataset.column_names,
+                available_columns=available_columns,
             )
 
-    def _validate_output_columns(self, dataset: Dataset) -> None:
+    def _validate_output_columns(self, df: pd.DataFrame) -> None:
         """Check that the output columns will not overwrite existing ones.
 
         Parameters
         ----------
-        dataset : Dataset
-            HuggingFace dataset to validate.
+        df : pd.DataFrame
+            DataFrame to validate.
 
         Raises
         ------
         OutputColumnCollisionError
-            If output columns already exist in the dataset.
+            If output columns already exist in the DataFrame.
         """
         if not self.output_cols:
             return
@@ -151,42 +153,43 @@ class BaseBlock(BaseModel, ABC):
             if isinstance(self.output_cols, dict)
             else self.output_cols
         )
-        collisions = [col for col in columns_to_check if col in dataset.column_names]
+        available_columns = df.columns.tolist()
+        collisions = [col for col in columns_to_check if col in available_columns]
         if collisions:
             raise OutputColumnCollisionError(
                 block_name=self.block_name,
                 collision_columns=collisions,
-                existing_columns=dataset.column_names,
+                existing_columns=available_columns,
             )
 
-    def _validate_dataset_not_empty(self, dataset: Dataset) -> None:
-        """Raise an error if the dataset is empty.
+    def _validate_dataframe_not_empty(self, df: pd.DataFrame) -> None:
+        """Raise an error if the DataFrame is empty.
 
         Parameters
         ----------
-        dataset : Dataset
+        df : pd.DataFrame
 
         Raises
         ------
         EmptyDatasetError
         """
-        if len(dataset) == 0:
+        if len(df) == 0:
             raise EmptyDatasetError(block_name=self.block_name)
 
-    def _validate_dataset(self, dataset: Dataset) -> None:
-        """Perform all default dataset validations."""
-        self._validate_dataset_not_empty(dataset)
-        self._validate_columns(dataset)
-        self._validate_output_columns(dataset)
+    def _validate_dataframe(self, df: pd.DataFrame) -> None:
+        """Perform all default DataFrame validations."""
+        self._validate_dataframe_not_empty(df)
+        self._validate_columns(df)
+        self._validate_output_columns(df)
 
-    def _validate_custom(self, dataset: Dataset) -> None:
+    def _validate_custom(self, df: pd.DataFrame) -> None:
         """Hook for subclasses to add extra validation logic."""
         pass
 
-    def _log_input_data(self, dataset: Dataset) -> None:
-        """Print a summary of the input dataset with Rich formatting."""
-        row_count = len(dataset)
-        columns = dataset.column_names
+    def _log_input_data(self, df: pd.DataFrame) -> None:
+        """Print a summary of the input DataFrame with Rich formatting."""
+        row_count = len(df)
+        columns = df.columns.tolist()
         content = Text()
         content.append("\U0001f4ca Processing Input Data\n", style="bold blue")
         content.append(f"Block Type: {self.__class__.__name__}\n", style="cyan")
@@ -207,13 +210,12 @@ class BaseBlock(BaseModel, ABC):
             Panel(content, title=f"[bold]{self.block_name}[/bold]", border_style="blue")
         )
 
-    def _log_output_data(self, input_dataset: Dataset, output_dataset: Dataset) -> None:
-        """Print a Rich panel summarizing output dataset differences."""
-        in_rows, out_rows = len(input_dataset), len(output_dataset)
-        in_cols, out_cols = (
-            set(input_dataset.column_names),
-            set(output_dataset.column_names),
-        )
+    def _log_output_data(self, input_df: pd.DataFrame, output_df: pd.DataFrame) -> None:
+        """Print a Rich panel summarizing output DataFrame differences."""
+        in_rows, out_rows = len(input_df), len(output_df)
+        in_cols = set(input_df.columns.tolist())
+        out_cols = set(output_df.columns.tolist())
+
         added_cols, removed_cols = out_cols - in_cols, in_cols - out_cols
         content = Text()
         content.append("\u2705 Processing Complete\n", style="bold green")
@@ -239,35 +241,35 @@ class BaseBlock(BaseModel, ABC):
         )
 
     @abstractmethod
-    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
+    def generate(self, samples: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
         """Subclass method to implement data generation logic.
 
         Parameters
         ----------
-        samples : Dataset
-            Input dataset to process.
+        samples : pd.DataFrame
+            Input DataFrame to process.
 
         Returns
         -------
-        Dataset
-            Transformed dataset with new columns or values.
+        pd.DataFrame
+            Transformed DataFrame with new columns or values.
         """
         pass
 
-    def __call__(self, samples: Dataset, **kwargs: Any) -> Dataset:
-        """Run the block on a dataset with full validation and logging.
+    def __call__(self, samples: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
+        """Run the block on a DataFrame with full validation and logging.
 
         Parameters
         ----------
-        samples : Dataset
-            Input dataset.
+        samples : pd.DataFrame
+            Input DataFrame.
         **kwargs : Any
             Runtime parameters to override block configuration
 
         Returns
         -------
-        Dataset
-            Output dataset after block processing.
+        pd.DataFrame
+            Output DataFrame after block processing.
         """
         # Handle runtime kwargs overrides
         if kwargs:
@@ -310,12 +312,12 @@ class BaseBlock(BaseModel, ABC):
 
             try:
                 self._log_input_data(samples)
-                self._validate_dataset(samples)
+                self._validate_dataframe(samples)
                 self._validate_custom(samples)
                 # Pass ALL kwargs to generate (including flow params)
-                output_dataset = self.generate(samples, **kwargs)
-                self._log_output_data(samples, output_dataset)
-                return output_dataset
+                output_df = self.generate(samples, **kwargs)
+                self._log_output_data(samples, output_df)
+                return output_df
             finally:
                 # Always restore original values for block fields
                 for key, value in original_values.items():
@@ -323,11 +325,11 @@ class BaseBlock(BaseModel, ABC):
         else:
             # Normal execution without overrides
             self._log_input_data(samples)
-            self._validate_dataset(samples)
+            self._validate_dataframe(samples)
             self._validate_custom(samples)
-            output_dataset = self.generate(samples)
-            self._log_output_data(samples, output_dataset)
-            return output_dataset
+            output_df = self.generate(samples)
+            self._log_output_data(samples, output_df)
+            return output_df
 
     def __repr__(self) -> str:
         """Compact string representation."""

sdg_hub/core/blocks/filtering/column_value_filter.py

@@ -9,10 +9,11 @@ using various operations with optional data type conversion.
 from typing import Any, Optional, Union
 import operator
 
-# Third Party
-from datasets import Dataset
 from pydantic import Field, field_validator
 
+# Third Party
+import pandas as pd
+
 # Local
 from ...utils.logger_config import setup_logger
 from ..base import BaseBlock
@@ -158,32 +159,44 @@ class ColumnValueFilterBlock(BaseBlock):
             sample[self.column_name] = None
         return sample
 
-    def generate(self, samples: Dataset, **_kwargs: Any) -> Dataset:
+    def generate(self, samples: pd.DataFrame, **_kwargs: Any) -> pd.DataFrame:
         """Generate filtered dataset based on specified conditions.
 
         Parameters
         ----------
-        samples : Dataset
+        samples : pd.DataFrame
             The input dataset to filter.
 
         Returns
         -------
-        Dataset
+        pd.DataFrame
             The filtered dataset.
         """
+        result = samples.copy()
+
+        # Convert dtype if specified
         if self._convert_dtype_func:
-            samples = samples.map(self._convert_dtype)
 
-        samples = samples.filter(
-            lambda x: x[self.column_name] is not None,
-        )
+            def safe_convert(x):
+                """Safely convert value, returning None on error."""
+                if pd.isna(x):
+                    return None
+                try:
+                    return self._convert_dtype_func(x)
+                except (ValueError, TypeError):
+                    return None
 
-        # Apply filter operation
-        samples = samples.filter(
-            lambda x: any(
-                self._operation_func(x[self.column_name], value)
-                for value in self.filter_value
-            )
+            result[self.column_name] = result[self.column_name].apply(safe_convert)
+
+        # Filter out None values
+        result = result[result[self.column_name].notna()]
+
+        # Apply filter operation using boolean indexing
+        # Create a mask that checks if any filter value matches
+        mask = result[self.column_name].apply(
+            lambda x: any(self._operation_func(x, value) for value in self.filter_value)
         )
 
-        return samples
+        result = result[mask]
+
+        return result

sdg_hub/core/blocks/llm/__init__.py

@@ -9,7 +9,6 @@ local models (vLLM, Ollama), and more.
 # Local
 from .error_handler import ErrorCategory, LLMErrorHandler
 from .llm_chat_block import LLMChatBlock
-from .llm_chat_with_parsing_retry_block import LLMChatWithParsingRetryBlock
 from .llm_parser_block import LLMParserBlock
 from .prompt_builder_block import PromptBuilderBlock
 from .text_parser_block import TextParserBlock
@@ -18,7 +17,6 @@ __all__ = [
     "LLMErrorHandler",
     "ErrorCategory",
     "LLMChatBlock",
-    "LLMChatWithParsingRetryBlock",
     "LLMParserBlock",
     "PromptBuilderBlock",
     "TextParserBlock",

sdg_hub/core/blocks/llm/llm_chat_block.py

@@ -5,12 +5,13 @@
 from typing import Any, Optional
 import asyncio
 
-# Third Party
-from datasets import Dataset
 from litellm import acompletion, completion
 from pydantic import ConfigDict, Field, field_validator
 import litellm
 
+# Third Party
+import pandas as pd
+
 from ...utils.error_handling import BlockValidationError
 from ...utils.logger_config import setup_logger
 
@@ -167,12 +168,12 @@ class LLMChatBlock(BaseBlock):
                 },
             )
 
-    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
+    def generate(self, samples: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
         """Generate responses from the LLM.
 
         Parameters
         ----------
-        samples : Dataset
+        samples : pd.DataFrame
             Input dataset containing the messages column.
         **kwargs : Any
             Runtime parameters that override initialization defaults.
@@ -180,7 +181,7 @@ class LLMChatBlock(BaseBlock):
 
         Returns
         -------
-        Dataset
+        pd.DataFrame
             Dataset with responses added to the output column.
 
         Raises
@@ -201,8 +202,8 @@ class LLMChatBlock(BaseBlock):
         # Build completion kwargs from ALL fields + runtime overrides
         completion_kwargs = self._build_completion_kwargs(**kwargs)
 
-        # Extract messages
-        messages_list = samples[self.input_cols[0]]
+        # Extract messages from pandas DataFrame
+        messages_list = samples[self.input_cols[0]].tolist()
 
         # Log generation start
         logger.info(
@@ -269,7 +270,9 @@ class LLMChatBlock(BaseBlock):
         )
 
         # Add responses as new column
-        return samples.add_column(self.output_cols[0], responses)
+        result = samples.copy()
+        result[self.output_cols[0]] = responses
+        return result
 
     def _build_completion_kwargs(self, **overrides) -> dict[str, Any]:
         """Build kwargs for LiteLLM completion call.
@@ -513,12 +516,14 @@ class LLMChatBlock(BaseBlock):
             )
             raise
 
-    def _validate_custom(self, dataset: Dataset) -> None:
+    def _validate_custom(self, dataset: pd.DataFrame) -> None:
         """Custom validation for LLMChatBlock message format.
 
+        Uses vectorized operations where possible for better performance.
+
         Parameters
         ----------
-        dataset : Dataset
+        dataset : pd.DataFrame
             The dataset to validate.
 
         Raises
@@ -526,28 +531,32 @@ class LLMChatBlock(BaseBlock):
         BlockValidationError
             If message format validation fails.
         """
+        messages_col = dataset[self.input_cols[0]]
+
+        # avoid using pd iterrows() when possible, it is notoriously slow: https://github.com/pandas-dev/pandas/issues/7683
+        # Vectorized check: all values must be lists
+        is_list = messages_col.apply(lambda x: isinstance(x, list))
+        if not is_list.all():
+            invalid_idx = is_list[~is_list].index[0]
+            invalid_value = messages_col.loc[invalid_idx]
+            raise BlockValidationError(
+                f"Messages column '{self.input_cols[0]}' must contain a list, "
+                f"got {type(invalid_value)} in row {invalid_idx}",
+                details=f"Block: {self.block_name}, Row: {invalid_idx}, Value: {invalid_value}",
+            )
 
-        def validate_sample(sample_with_index):
-            """Validate a single sample's message format."""
-            idx, sample = sample_with_index
-            messages = sample[self.input_cols[0]]
-
-            # Validate messages is a list
-            if not isinstance(messages, list):
-                raise BlockValidationError(
-                    f"Messages column '{self.input_cols[0]}' must contain a list, "
-                    f"got {type(messages)} in row {idx}",
-                    details=f"Block: {self.block_name}, Row: {idx}, Value: {messages}",
-                )
-
-            # Validate messages is not empty
-            if not messages:
-                raise BlockValidationError(
-                    f"Messages list is empty in row {idx}",
-                    details=f"Block: {self.block_name}, Row: {idx}",
-                )
+        # Vectorized check: no empty lists
+        is_empty = messages_col.apply(lambda x: len(x) == 0)
+        if is_empty.any():
+            invalid_idx = is_empty[is_empty].index[0]
+            raise BlockValidationError(
+                f"Messages list is empty in row {invalid_idx}",
+                details=f"Block: {self.block_name}, Row: {invalid_idx}",
+            )
 
-            # Validate each message format
+        # Validate nested message structure (requires iteration over messages column only)
+        def validate_message_structure(messages, idx):
+            """Validate structure of messages list."""
             for msg_idx, message in enumerate(messages):
                 if not isinstance(message, dict):
                     raise BlockValidationError(
@@ -555,7 +564,6 @@ class LLMChatBlock(BaseBlock):
                         details=f"Block: {self.block_name}, Row: {idx}, Message: {msg_idx}, Value: {message}",
                     )
 
-                # Validate required fields
                 if "role" not in message or message["role"] is None:
                     raise BlockValidationError(
                         f"Message {msg_idx} in row {idx} missing required 'role' field",
@@ -568,11 +576,9 @@ class LLMChatBlock(BaseBlock):
                         details=f"Block: {self.block_name}, Row: {idx}, Message: {msg_idx}, Available fields: {list(message.keys())}",
                     )
 
-            return True
-
-        # Validate all samples
-        indexed_samples = [(i, sample) for i, sample in enumerate(dataset)]
-        list(map(validate_sample, indexed_samples))
+        # Iterate only over the messages column (not the entire DataFrame)
+        for idx, messages in messages_col.items():
+            validate_message_structure(messages, idx)
 
     def __repr__(self) -> str:
         """String representation of the block."""

sdg_hub/core/blocks/llm/llm_parser_block.py

@@ -8,10 +8,11 @@ This module provides the LLMParserBlock for extracting specific fields
 # Standard
 from typing import Any
 
-# Third Party
-from datasets import Dataset
 from pydantic import Field, model_validator
 
+# Third Party
+import pandas as pd
+
 # Local
 from ...utils.logger_config import setup_logger
 from ..base import BaseBlock
@@ -26,6 +27,8 @@ logger = setup_logger(__name__)
     "Extracts specified fields from LLM response objects",
 )
 class LLMParserBlock(BaseBlock):
+    _flow_requires_jsonl_tmp: bool = True
+
     """Block for extracting fields from LLM response objects.
 
     This block extracts specified fields from chat completion response objects.
@@ -102,12 +105,12 @@ class LLMParserBlock(BaseBlock):
 
         return self
 
-    def _validate_custom(self, dataset: Dataset) -> None:
+    def _validate_custom(self, dataset: pd.DataFrame) -> None:
         """Validate LLMParserBlock specific requirements.
 
         Parameters
         ----------
-        dataset : Dataset
+        dataset : pd.DataFrame
             The dataset to validate.
 
         Raises
@@ -308,13 +311,16 @@ class LLMParserBlock(BaseBlock):
         extracted = self._extract_fields_from_response(raw_output)
         return [{**sample, **extracted}]
 
-    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
+    def generate(self, samples: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
         logger.debug(f"Extracting fields from {len(samples)} samples")
         if len(samples) == 0:
             logger.warning("No samples to process, returning empty dataset")
-            return Dataset.from_list([])
+            return pd.DataFrame()
 
         new_data = []
+        samples = samples.to_dict("records")  # Avoid Iterrows() when possible
+
         for sample in samples:
             new_data.extend(self._generate(sample))
-        return Dataset.from_list(new_data)
+
+        return pd.DataFrame(new_data)

sdg_hub/core/blocks/llm/prompt_builder_block.py

@@ -8,10 +8,11 @@ including conversion to OpenAI Messages format and template rendering.
 # Standard
 from typing import Any, Literal, Optional
 
-# Third Party
-from datasets import Dataset
 from jinja2 import Template, meta
 from pydantic import BaseModel, Field, field_validator
+
+# Third Party
+import pandas as pd
 import yaml
 
 # Local
@@ -279,12 +280,14 @@ class PromptBuilderBlock(BaseBlock):
         message_templates = self.prompt_template_config.get_message_templates()
         self.prompt_renderer = PromptRenderer(message_templates)
 
-    def _validate_custom(self, dataset: Dataset) -> None:
+    def _validate_custom(self, dataset: pd.DataFrame) -> None:
         if len(dataset) > 0:
             # Get required variables from all message templates
             required_vars = self.prompt_renderer.get_required_variables()
 
-            sample = dataset[0]
+            # Get first row as dict
+            sample = dataset.iloc[0].to_dict()
+
             template_vars = self.prompt_renderer.resolve_template_vars(
                 sample, self.input_cols
             )
@@ -344,25 +347,27 @@ class PromptBuilderBlock(BaseBlock):
 
         return sample
 
-    def generate(self, samples: Dataset, **_kwargs: Any) -> Dataset:
-        """Generate formatted output for all samples using dataset map.
+    def generate(self, samples: pd.DataFrame, **_kwargs: Any) -> pd.DataFrame:
+        """Generate formatted output for all samples.
 
         Parameters
         ----------
-        samples : Dataset
+        samples : pd.DataFrame
             Input dataset containing samples to be formatted.
         **kwargs : Dict[str, Any]
             Additional keyword arguments (unused in this block).
 
         Returns
         -------
-        Dataset
+        pd.DataFrame
             Dataset with the formatted output added to the specified column.
         """
         logger.debug(f"Formatting prompts for {len(samples)} samples")
 
-        # Use dataset map for efficient processing
-        formatted_dataset = samples.map(self._generate)
+        # Convert DataFrame to list of dicts, process each, and convert back
+        samples_list = samples.to_dict("records")
+        formatted_samples = [self._generate(sample) for sample in samples_list]
+        formatted_dataset = pd.DataFrame(formatted_samples)
 
         logger.debug(f"Successfully formatted {len(formatted_dataset)} samples")
         return formatted_dataset