sdg-hub 0.5.1__py3-none-any.whl → 0.6.1__py3-none-any.whl

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

sdg_hub/core/blocks/llm/text_parser_block.py

@@ -7,17 +7,15 @@ start/end tags, custom regex patterns, and cleanup operations.
 
 # Standard
 from typing import Any, Optional
-from weakref import finalize
-import json
 import re
 
-# Third Party
-from datasets import Dataset, load_dataset
 from pydantic import Field, field_validator, model_validator
 
+# Third Party
+import pandas as pd
+
 # Local
 from ...utils.logger_config import setup_logger
-from ...utils.temp_manager import cleanup_path, create_temp_dir, create_temp_file
 from ..base import BaseBlock
 from ..registry import BlockRegistry
 
@@ -122,12 +120,12 @@ class TextParserBlock(BaseBlock):
 
         return self
 
-    def _validate_custom(self, dataset: Dataset) -> None:
+    def _validate_custom(self, dataset: pd.DataFrame) -> None:
         """Validate TextParserBlock specific requirements.
 
         Parameters
         ----------
-        dataset : Dataset
+        dataset : pd.DataFrame
             The dataset to validate.
 
         Raises
@@ -316,60 +314,15 @@ class TextParserBlock(BaseBlock):
             )
             return []
 
-    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
+    def generate(self, samples: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
         logger.debug(f"Parsing outputs for {len(samples)} samples")
         if len(samples) == 0:
             logger.warning("No samples to parse, returning empty dataset")
-            return Dataset.from_list([])
-
-        tmp_jsonl_path = kwargs.get("_flow_tmp_jsonl_path")
-        cleanup_locally = False
-
-        if tmp_jsonl_path is None:
-            tmp_jsonl_path = str(
-                create_temp_file(
-                    prefix=f"{self.block_name}_text_parser", suffix=".jsonl"
-                )
-            )
-            cleanup_locally = True
-
-        rows_written = 0
-        batch = []
-        with open(tmp_jsonl_path, "w") as f:
-            for sample in samples:
-                out = self._generate(sample)
-                for row in out:
-                    batch.append(json.dumps(row) + "\n")
-                    rows_written += 1
-                    if len(batch) >= 5:
-                        f.writelines(batch)
-                        batch.clear()
-            if batch:
-                f.writelines(batch)
-
-        if rows_written == 0:
-            if cleanup_locally:
-                cleanup_path(tmp_jsonl_path)
-            return Dataset.from_list([])
-
-        hf_cache_dir = None
-        try:
-            hf_cache_dir = create_temp_dir(
-                prefix=f"{self.block_name}_text_parser_hf_cache"
-            )
-            ret = load_dataset(
-                "json",
-                data_files=tmp_jsonl_path,
-                split="train",
-                keep_in_memory=False,
-                cache_dir=str(hf_cache_dir),
-            )
-            finalize(ret, cleanup_path, hf_cache_dir)
-            return ret
-        except Exception:
-            if hf_cache_dir is not None:
-                cleanup_path(hf_cache_dir)
-            raise
-        finally:
-            if cleanup_locally:
-                cleanup_path(tmp_jsonl_path)
+            return pd.DataFrame()
+
+        # Convert DataFrame to list of dicts to avoid iterrows and improve performance
+        samples_list = samples.to_dict("records")
+        new_data: list[dict] = []
+        for sample in samples_list:
+            new_data.extend(self._generate(sample))
+        return pd.DataFrame(new_data)

sdg_hub/core/blocks/transform/duplicate_columns.py

@@ -8,10 +8,11 @@ according to a mapping specification.
 # Standard
 from typing import Any
 
-# Third Party
-from datasets import Dataset
 from pydantic import field_validator
 
+# Third Party
+import pandas as pd
+
 # Local
 from ...utils.logger_config import setup_logger
 from ..base import BaseBlock
@@ -62,27 +63,27 @@ class DuplicateColumnsBlock(BaseBlock):
         if self.output_cols is None:
             self.output_cols = list(self.input_cols.values())
 
-    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
+    def generate(self, samples: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
         """Generate a dataset with duplicated columns.
 
         Parameters
         ----------
-        samples : Dataset
+        samples : pd.DataFrame
             Input dataset to duplicate columns from.
 
         Returns
         -------
-        Dataset
+        pd.DataFrame
             Dataset with additional duplicated columns.
         """
         # Create a copy to avoid modifying the original
-        result = samples
+        result = samples.copy()
 
         # Duplicate each column as specified in the mapping
         for source_col, target_col in self.input_cols.items():
-            if source_col not in result.column_names:
+            if source_col not in result.columns.tolist():
                 raise ValueError(f"Source column '{source_col}' not found in dataset")
 
-            result = result.add_column(target_col, result[source_col])
+            result[target_col] = result[source_col]
 
         return result

sdg_hub/core/blocks/transform/index_based_mapper.py

@@ -8,10 +8,11 @@ to another based on a choice column's value.
 # Standard
 from typing import Any
 
-# Third Party
-from datasets import Dataset
 from pydantic import Field, field_validator, model_validator
 
+# Third Party
+import pandas as pd
+
 # Local
 from ...utils.error_handling import MissingColumnError
 from ...utils.logger_config import setup_logger
@@ -103,12 +104,12 @@ class IndexBasedMapperBlock(BaseBlock):
         # Create mapping from choice_col to output_col for easy access
         self.choice_to_output_map = dict(zip(self.choice_cols, self.output_cols))
 
-    def _validate_custom(self, samples: Dataset) -> None:
+    def _validate_custom(self, samples: pd.DataFrame) -> None:
         """Validate that required columns exist in the dataset.
 
         Parameters
         ----------
-        samples : Dataset
+        samples : pd.DataFrame
             Input dataset to validate.
 
         Raises
@@ -120,29 +121,29 @@ class IndexBasedMapperBlock(BaseBlock):
         """
         # Check that all choice_cols exist
         missing_choice_cols = [
-            col for col in self.choice_cols if col not in samples.column_names
+            col for col in self.choice_cols if col not in samples.columns.tolist()
         ]
         if missing_choice_cols:
             raise MissingColumnError(
                 block_name=self.block_name,
                 missing_columns=missing_choice_cols,
-                available_columns=samples.column_names,
+                available_columns=samples.columns.tolist(),
             )
 
         # Check that all mapped columns exist
         mapped_cols = list(self.choice_map.values())
-        missing_cols = list(set(mapped_cols) - set(samples.column_names))
+        missing_cols = list(set(mapped_cols) - set(samples.columns.tolist()))
         if missing_cols:
             raise MissingColumnError(
                 block_name=self.block_name,
                 missing_columns=missing_cols,
-                available_columns=samples.column_names,
+                available_columns=samples.columns.tolist(),
             )
 
         # Check that all choice values in all choice columns have corresponding mappings
         all_unique_choices = set()
         for choice_col in self.choice_cols:
-            all_unique_choices.update(samples[choice_col])
+            all_unique_choices.update(samples[choice_col].unique())
 
         mapped_choices = set(self.choice_map.keys())
         unmapped_choices = all_unique_choices - mapped_choices
@@ -174,23 +175,23 @@ class IndexBasedMapperBlock(BaseBlock):
             sample[output_col] = sample[source_col]
         return sample
 
-    def generate(self, samples: Dataset, **kwargs) -> Dataset:
+    def generate(self, samples: pd.DataFrame, **kwargs) -> pd.DataFrame:
         """Generate a new dataset with selected values.
 
         Parameters
         ----------
-        samples : Dataset
+        samples : pd.DataFrame
             Input dataset to process.
 
         Returns
         -------
-        Dataset
+        pd.DataFrame
             Dataset with selected values stored in output column.
         """
         # Log the operation
         all_unique_choices = set()
         for choice_col in self.choice_cols:
-            all_unique_choices.update(samples[choice_col])
+            all_unique_choices.update(samples[choice_col].unique())
         mapped_choices = set(self.choice_map.keys())
 
         logger.info(
@@ -205,8 +206,21 @@ class IndexBasedMapperBlock(BaseBlock):
             },
         )
 
-        # Apply the mapping
-        result = samples.map(self._generate)
+        # Create a copy to avoid modifying the input
+        result = samples.copy()
+
+        # Handle empty DataFrame case
+        if len(result) == 0:
+            # Add empty output columns
+            for output_col in self.output_cols:
+                result[output_col] = []
+        else:
+            # Apply the mapping for each choice column and output column pair
+            for choice_col, output_col in self.choice_to_output_map.items():
+                # Map the choice values to source columns, then get values from those columns
+                result[output_col] = result.apply(
+                    lambda row: row[self.choice_map[row[choice_col]]], axis=1
+                )
 
         # Log completion
         logger.info(

sdg_hub/core/blocks/transform/json_structure_block.py

@@ -9,10 +9,11 @@ containing a structured JSON object with specified field names.
 from typing import Any, Dict
 import json
 
-# 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
@@ -90,17 +91,17 @@ class JSONStructureBlock(BaseBlock):
 
         raise ValueError("input_cols must be a list of column names")
 
-    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
+    def generate(self, samples: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
         """Generate a dataset with JSON structured output.
 
         Parameters
         ----------
-        samples : Dataset
+        samples : pd.DataFrame
             Input dataset to process.
 
         Returns
         -------
-        Dataset
+        pd.DataFrame
             Dataset with JSON structured output in the specified column.
         """
         if not self.output_cols:
@@ -109,17 +110,17 @@ class JSONStructureBlock(BaseBlock):
         output_col = self.output_cols[0]
         field_mapping = self._get_field_mapping()
 
-        def _create_json_structure(sample):
+        def _create_json_structure(row):
             """Create JSON structure from input columns."""
             json_obj = {}
 
             # Build the JSON object using the field mapping
             for json_field, col_name in field_mapping.items():
-                if col_name not in sample:
-                    logger.warning(f"Input column '{col_name}' not found in sample")
+                if col_name not in row.index:
+                    logger.warning(f"Input column '{col_name}' not found in row")
                     json_obj[json_field] = None
                 else:
-                    value = sample[col_name]
+                    value = row[col_name]
                     if self.ensure_json_serializable:
                         value = self._make_json_serializable(value)
                     json_obj[json_field] = value
@@ -130,13 +131,15 @@ class JSONStructureBlock(BaseBlock):
                     json_string = json.dumps(json_obj, indent=2, ensure_ascii=False)
                 else:
                     json_string = json.dumps(json_obj, ensure_ascii=False)
-                sample[output_col] = json_string
+                return json_string
             except (TypeError, ValueError) as e:
                 logger.error(f"Failed to serialize JSON object: {e}")
-                sample[output_col] = "{}"
+                return "{}"
 
-            return sample
+        # Create a copy to avoid modifying the input
+        result = samples.copy()
 
         # Apply the JSON structuring to all samples
-        result = samples.map(_create_json_structure)
+        result[output_col] = result.apply(_create_json_structure, axis=1)
+
         return result

sdg_hub/core/blocks/transform/melt_columns.py

@@ -8,10 +8,11 @@ by melting specified columns into rows.
 # Standard
 from typing import Any
 
-# Third Party
-from datasets import Dataset
 from pydantic import field_validator
 
+# Third Party
+import pandas as pd
+
 # Local
 from ...utils.error_handling import MissingColumnError
 from ...utils.logger_config import setup_logger
@@ -79,12 +80,12 @@ class MeltColumnsBlock(BaseBlock):
             self.input_cols if isinstance(self.input_cols, list) else [self.input_cols]
         )
 
-    def _validate_custom(self, samples: Dataset) -> None:
+    def _validate_custom(self, samples: pd.DataFrame) -> None:
         """Validate that required columns exist in the dataset.
 
         Parameters
         ----------
-        samples : Dataset
+        samples : pd.DataFrame
             Input dataset to validate.
 
         Raises
@@ -93,34 +94,34 @@ class MeltColumnsBlock(BaseBlock):
             If required columns are missing from the dataset.
         """
         # Check that all var_cols exist in the dataset
-        missing_cols = list(set(self.var_cols) - set(samples.column_names))
+        missing_cols = list(set(self.var_cols) - set(samples.columns.tolist()))
         if missing_cols:
             raise MissingColumnError(
                 block_name=self.block_name,
                 missing_columns=missing_cols,
-                available_columns=samples.column_names,
+                available_columns=samples.columns.tolist(),
             )
 
-    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
+    def generate(self, samples: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
         """Generate a flattened dataset in long format.
 
         Parameters
         ----------
-        samples : Dataset
+        samples : pd.DataFrame
             Input dataset to flatten.
 
         Returns
         -------
-        Dataset
+        pd.DataFrame
             Flattened dataset in long format with new variable and value columns.
         """
         # Use the original simple logic - just adapted to use derived attributes
-        df = samples.to_pandas()
-        id_cols = [col for col in samples.column_names if col not in self.var_cols]
+        df = samples
+        id_cols = [col for col in samples.columns.tolist() if col not in self.var_cols]
         flatten_df = df.melt(
             id_vars=id_cols,
             value_vars=self.var_cols,
             value_name=self.value_name,
             var_name=self.var_name,
         )
-        return Dataset.from_pandas(flatten_df)
+        return flatten_df

sdg_hub/core/blocks/transform/rename_columns.py

@@ -8,10 +8,11 @@ to a mapping specification.
 # Standard
 from typing import Any
 
-# Third Party
-from datasets import Dataset
 from pydantic import field_validator
 
+# Third Party
+import pandas as pd
+
 # Local
 from ...utils.logger_config import setup_logger
 from ..base import BaseBlock
@@ -52,28 +53,38 @@ class RenameColumnsBlock(BaseBlock):
             )
         return v
 
-    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
+    def generate(self, samples: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
         """Generate a dataset with renamed columns.
 
         Parameters
         ----------
-        samples : Dataset
+        samples : pd.DataFrame
             Input dataset to rename columns in.
 
         Returns
         -------
-        Dataset
+        pd.DataFrame
             Dataset with renamed columns.
 
         Raises
         ------
         ValueError
-            If attempting to rename to a column name that already exists.
+            If attempting to rename to a column name that already exists,
+            or if the original column names don't exist in the dataset.
         """
+        # Check that all original column names exist in the dataset
+        existing_cols = set(samples.columns.tolist())
+        original_cols = set(self.input_cols.keys())
+
+        missing_cols = original_cols - existing_cols
+        if missing_cols:
+            raise ValueError(
+                f"Original column names {sorted(missing_cols)} not in the dataset"
+            )
+
         # Check for column name collisions
         # Strict validation: no target column name can be an existing column name
         # This prevents chained/circular renames which can be confusing
-        existing_cols = set(samples.column_names)
         target_cols = set(self.input_cols.values())
 
         collision = target_cols & existing_cols
@@ -84,5 +95,5 @@ class RenameColumnsBlock(BaseBlock):
                 "Chained renames are not supported."
             )
 
-        # Rename columns using HuggingFace datasets method
-        return samples.rename_columns(self.input_cols)
+        # Rename columns using pandas method
+        return samples.rename(columns=self.input_cols)

sdg_hub/core/blocks/transform/text_concat.py

@@ -8,10 +8,11 @@ using a specified separator.
 # Standard
 from typing import Any
 
-# 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
@@ -65,17 +66,17 @@ class TextConcatBlock(BaseBlock):
             raise ValueError("TextConcatBlock requires exactly one output column")
         return v
 
-    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
+    def generate(self, samples: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
         """Generate a dataset with combined columns.
 
         Parameters
         ----------
-        samples : Dataset
+        samples : pd.DataFrame
             Input dataset to process.
 
         Returns
         -------
-        Dataset
+        pd.DataFrame
             Dataset with combined values stored in output column.
         """
         if not self.output_cols:
@@ -83,20 +84,18 @@ class TextConcatBlock(BaseBlock):
 
         output_col = self.output_cols[0]
 
-        def _combine_columns(sample):
-            """Combine values from input columns."""
-            # Check that all input columns exist
-            for col in self.input_cols:
-                if col not in sample:
-                    raise ValueError(f"Input column '{col}' not found in sample")
-
-            # Combine values using separator
-            combined_value = self.separator.join(
-                [str(sample[col]) for col in self.input_cols]
-            )
-            sample[output_col] = combined_value
-            return sample
-
-        # Apply the combination to all samples
-        result = samples.map(_combine_columns)
+        # Validate that all input columns exist in the dataset
+        for col in self.input_cols:
+            if col not in samples.columns:
+                raise ValueError(f"Input column '{col}' not found in sample")
+
+        # Create a copy to avoid modifying the input
+        result = samples.copy()
+
+        # Combine columns using vectorized string operations
+        # Convert all input columns to strings and concatenate with separator
+        result[output_col] = (
+            result[self.input_cols].astype(str).agg(self.separator.join, axis=1)
+        )
+
         return result

sdg_hub/core/blocks/transform/uniform_col_val_setter.py

@@ -8,11 +8,12 @@ mode, min, max, mean, or median.
 # Standard
 from typing import Any, Literal
 
-# Third Party
-from datasets import Dataset
 from pydantic import field_validator
 import numpy as np
 
+# Third Party
+import pandas as pd
+
 # Local
 from ...utils.logger_config import setup_logger
 from ..base import BaseBlock
@@ -66,8 +67,8 @@ class UniformColumnValueSetter(BaseBlock):
         self.output_cols = []
         self.col_name = self.input_cols[0]
 
-    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
-        df = samples.to_pandas()
+    def generate(self, samples: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
+        df = samples.copy()
 
         if df.empty:
             raise ValueError("Cannot compute reduction for empty dataset")
@@ -98,4 +99,4 @@ class UniformColumnValueSetter(BaseBlock):
         )
 
         df[self.col_name] = value
-        return Dataset.from_pandas(df)
+        return df