sdg-hub 0.1.0a4__py3-none-any.whl → 0.1.1__py3-none-any.whl

sdg_hub/flow.py

@@ -1,58 +1,145 @@
+"""
+Flow module for managing data generation pipelines.
+
+This module provides the core Flow class that handles both configuration loading and execution
+of data generation blocks. The Flow class serves as the main interface for defining and running
+data generation pipelines, supporting both direct usage with SDG and backward compatibility
+through the deprecated Pipeline class.
+
+Example:
+    >>> flow = Flow(llm_client)
+    >>> flow = flow.get_flow_from_file("path/to/flow.yaml")
+    >>> dataset = flow.generate(input_dataset)
+
+Note:
+    This module is part of the SDG Hub package and is designed to work in conjunction
+    with the SDG class for distributed data generation.
+"""
+
 # SPDX-License-Identifier: Apache-2.0
 # Standard
 from abc import ABC
 from importlib import resources
-from typing import Optional
+from typing import Any, Callable, Dict, List, Optional
 import operator
 import os
 
 # Third Party
+from datasets import Dataset
+from datasets.data_files import EmptyDatasetError
+from jinja2 import Environment, meta
+from rich.console import Console
+from rich.table import Table
 import yaml
 
 # Local
+from .blocks import *  # needed to register blocks
+from .logger_config import setup_logger
+from .prompts import *  # needed to register prompts
 from .registry import BlockRegistry, PromptRegistry
-from . import prompts
-from . import blocks
+from .utils.config_validation import validate_prompt_config_schema
+from .utils.validation_result import ValidationResult
+
+logger = setup_logger(__name__)
 
 
-OPERATOR_MAP = {
+OPERATOR_MAP: Dict[str, Callable] = {
     "operator.eq": operator.eq,
     "operator.ge": operator.ge,
+    "operator.le": operator.le,
+    "operator.gt": operator.gt,
+    "operator.lt": operator.lt,
+    "operator.ne": operator.ne,
     "operator.contains": operator.contains,
 }
 
-CONVERT_DTYPE_MAP = {
+CONVERT_DTYPE_MAP: Dict[str, Callable] = {
     "float": float,
     "int": int,
 }
 
 
 class Flow(ABC):
+    """A class representing a data generation flow.
+
+    This class handles both configuration loading and execution of data generation
+    blocks. It can be used directly with SDG or through the deprecated Pipeline class.
+    """
+
     def __init__(
         self,
-        llm_client,
+        llm_client: Any,
         num_samples_to_generate: Optional[int] = None,
+        log_level: Optional[str] = None,
     ) -> None:
+        """
+        Initialize the Flow class.
+
+        Parameters
+        ----------
+        llm_client : Any
+            The LLM client to use for generation.
+        num_samples_to_generate : Optional[int], optional
+            Number of samples to generate, by default None
+        log_level : Optional[str], optional
+            Logging verbosity level, by default None
+
+        Attributes
+        ----------
+        llm_client : Any
+            The LLM client instance.
+        base_path : str
+            Base path for resource files.
+        registered_blocks : Dict[str, Any]
+            Registry of available blocks.
+        chained_blocks : Optional[List[Dict[str, Any]]]
+            List of block configurations.
+        num_samples_to_generate : Optional[int]
+            Number of samples to generate.
+
+        """
         self.llm_client = llm_client
-        self.num_samples_to_generate = num_samples_to_generate
         self.base_path = str(resources.files(__package__))
         self.registered_blocks = BlockRegistry.get_registry()
+        self.chained_blocks = None  # Will be set by get_flow_from_file
+        self.num_samples_to_generate = num_samples_to_generate
 
-    def _getFilePath(self, dirs, filename):
-        """
-        Find a named configuration file.
+        # Logging verbosity level
+        self.log_level = log_level or os.getenv("SDG_HUB_LOG_LEVEL", "normal").lower()
+        self.console = Console() if self.log_level in ["verbose", "debug"] else None
 
-        Files are checked in the following order
-            - absulute path is always used
-            - checked relative to the directories in "dirs"
-            - relative the the current directory
+    def _log_block_info(
+        self, index: int, total: int, name: str, ds: Dataset, stage: str
+    ) -> None:
+        if self.log_level in ["verbose", "debug"] and self.console:
+            table = Table(
+                title=f"{stage} Block {index + 1}/{total}: {name}", show_header=True
+            )
+            table.add_column("Metric", style="cyan", no_wrap=True)
+            table.add_column("Value", style="magenta")
+            table.add_row("Rows", str(len(ds)))
+            table.add_row("Columns", ", ".join(ds.column_names))
+            self.console.print(table)
+
+    def _getFilePath(self, dirs: List[str], filename: str) -> str:
+        """Find a named configuration file.
 
-        Args:
-            dirs (list): Directories in which to search for "config_path"
-            config_path (str): The path to the configuration file.
+        Files are checked in the following order:
+            1. Absolute path is always used
+            2. Checked relative to the directories in "dirs"
+            3. Relative to the current directory
 
-        Returns:
-            Selected file path
+        Parameters
+        ----------
+        dirs : List[str]
+            Directories in which to search for the file.
+        filename : str
+            The path to the configuration file.
+
+        Returns
+        -------
+        str
+            Selected file path.
         """
         if os.path.isabs(filename):
             return filename
@@ -64,7 +151,175 @@ class Flow(ABC):
         # assume the path is relative to the current directory
         return filename
 
-    def get_flow_from_file(self, yaml_path: str) -> list:
+    def _drop_duplicates(self, dataset: Dataset, cols: List[str]) -> Dataset:
+        """Drop duplicates from the dataset based on the columns provided.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The input dataset.
+        cols : List[str]
+            Columns to consider for duplicate detection.
+
+        Returns
+        -------
+        Dataset
+            Dataset with duplicates removed.
+        """
+        df = dataset.to_pandas()
+        df = df.drop_duplicates(subset=cols).reset_index(drop=True)
+        return Dataset.from_pandas(df)
+
+    def generate(self, dataset: Dataset) -> Dataset:
+        """Generate the dataset by running the pipeline steps.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The input dataset to process.
+
+        Returns
+        -------
+        Dataset
+            The processed dataset.
+
+        Raises
+        ------
+        ValueError
+            If Flow has not been initialized with blocks.
+        EmptyDatasetError
+            If a block produces an empty dataset.
+        """
+        if self.chained_blocks is None:
+            raise ValueError(
+                "Flow has not been initialized with blocks. "
+                "Call get_flow_from_file() first. "
+                "Or pass a list of blocks to the Flow constructor."
+            )
+
+        for i, block_prop in enumerate(self.chained_blocks):
+            block_type = block_prop["block_type"]
+            block_config = block_prop["block_config"]
+            drop_columns = block_prop.get("drop_columns", [])
+            gen_kwargs = block_prop.get("gen_kwargs", {})
+            drop_duplicates_cols = block_prop.get("drop_duplicates", False)
+            block = block_type(**block_config)
+
+            name = block_config.get("block_name", f"block_{i}")
+
+            # Logging: always show basic progress unless in quiet mode
+            if self.log_level in ["normal", "verbose", "debug"]:
+                logger.info(
+                    f"🔄 Running block {i + 1}/{len(self.chained_blocks)}: {name}"
+                )
+
+            # Log dataset shape before block (verbose/debug)
+            self._log_block_info(i, len(self.chained_blocks), name, dataset, "Input")
+
+            if self.log_level == "debug":
+                logger.debug(f"Input dataset (truncated): {dataset}")
+
+            dataset = block.generate(dataset, **gen_kwargs)
+
+            if len(dataset) == 0:
+                raise EmptyDatasetError(
+                    f"Pipeline stopped: "
+                    f"Empty dataset after running block: "
+                    f"{block_config['block_name']}"
+                )
+
+            drop_columns_in_ds = [e for e in drop_columns if e in dataset.column_names]
+            if drop_columns:
+                dataset = dataset.remove_columns(drop_columns_in_ds)
+
+            if drop_duplicates_cols:
+                dataset = self._drop_duplicates(dataset, cols=drop_duplicates_cols)
+
+            # Log dataset shape after block (verbose/debug)
+            self._log_block_info(i, len(self.chained_blocks), name, dataset, "Output")
+
+            if self.log_level == "debug":
+                logger.debug(f"Output dataset (truncated): {dataset}")
+
+        return dataset
+
+    def validate_config_files(self) -> "ValidationResult":
+        """
+        Validate all configuration file paths referenced in the flow blocks.
+
+        This method checks that all config files specified via `config_path` or `config_paths`
+        in each block:
+            - Exist on the filesystem
+            - Are readable by the current process
+            - Are valid YAML files (optional format check)
+
+        Returns
+        -------
+        ValidationResult
+            An object indicating whether all config files passed validation, along with a list
+            of error messages for any missing, unreadable, or invalid YAML files.
+
+        Notes
+        -----
+        This method is automatically called at the end of `get_flow_from_file()` to ensure
+        early detection of misconfigured blocks.
+        """
+        errors = []
+
+        def check_file(path: str, context: str):
+            if not os.path.isfile(path):
+                errors.append(f"[{context}] File does not exist: {path}")
+            else:
+                try:
+                    with open(path, "r", encoding="utf-8") as f:
+                        config_data = yaml.safe_load(f)
+                        _, validation_errors = validate_prompt_config_schema(config_data, path)
+
+                        if validation_errors:
+                            errors.extend(validation_errors)
+
+                except PermissionError:
+                    errors.append(f"[{context}] File is not readable: {path}")
+                except yaml.YAMLError as e:
+                    errors.append(f"[{context}] YAML load failed: {path} ({e})")
+
+        for i, block in enumerate(self.chained_blocks or []):
+            block_name = block["block_config"].get("block_name", f"block_{i}")
+
+            config_path = block["block_config"].get("config_path")
+            if config_path:
+                check_file(config_path, f"{block_name}.config_path")
+
+            config_paths = block["block_config"].get("config_paths")
+            if isinstance(config_paths, list):
+                for idx, path in enumerate(config_paths):
+                    check_file(path, f"{block_name}.config_paths[{idx}]")
+            elif isinstance(config_paths, dict):
+                for key, path in config_paths.items():
+                    check_file(path, f"{block_name}.config_paths['{key}']")
+
+        return ValidationResult(valid=(len(errors) == 0), errors=errors)
+
+    def get_flow_from_file(self, yaml_path: str) -> "Flow":
+        """Load and initialize flow configuration from a YAML file.
+
+        Parameters
+        ----------
+        yaml_path : str
+            Path to the YAML configuration file.
+
+        Returns
+        -------
+        Flow
+            Self with initialized chained_blocks.
+
+        Raises
+        ------
+        FileNotFoundError
+            If the YAML file cannot be found.
+        KeyError
+            If a required block or prompt is not found in the registry.
+        """
         yaml_path_relative_to_base = os.path.join(self.base_path, yaml_path)
         if os.path.isfile(yaml_path_relative_to_base):
             yaml_path = yaml_path_relative_to_base
@@ -141,4 +396,79 @@ class Flow(ABC):
                     block["block_config"]["convert_dtype"]
                 ]
 
-        return flow
+        # Store the chained blocks and return self
+        self.chained_blocks = flow
+
+        # Validate config files
+        result = self.validate_config_files()
+        if not result.valid:
+            raise ValueError("Invalid config files:\n\n".join(result.errors))
+
+        return self
+
+    def validate_flow(self, dataset: Dataset) -> "ValidationResult":
+        """
+        Validate that all required dataset columns are present before executing the flow.
+
+        This includes:
+        - Columns referenced in Jinja templates for LLM blocks
+        - Columns required by specific utility blocks (e.g. filter_column, choice_col, etc.)
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The input dataset to validate against.
+
+        Returns
+        -------
+        ValidationResult
+            Whether the dataset has all required columns, and which ones are missing.
+        """
+        errors = []
+        all_columns = set(dataset.column_names)
+
+        for i, block in enumerate(self.chained_blocks or []):
+            name = block["block_config"].get("block_name", f"block_{i}")
+            block_type = block["block_type"]
+            config = block["block_config"]
+
+            # LLM Block: parse Jinja vars
+            cls_name = block_type.__name__ if isinstance(block_type, type) else block_type.__class__.__name__
+            logger.info(f"Validating block: {name} ({cls_name})")
+            if "LLM" in cls_name:
+                config_path = config.get("config_path")
+                if config_path and os.path.isfile(config_path):
+                    with open(config_path, "r", encoding="utf-8") as f:
+                        content = f.read()
+                        env = Environment()
+                        ast = env.parse(content)
+                        vars_found = meta.find_undeclared_variables(ast)
+                        for var in vars_found:
+                            if var not in all_columns:
+                                errors.append(f"[{name}] Missing column for prompt var: '{var}'")
+
+            # FilterByValueBlock
+            if "FilterByValueBlock" in str(block_type):
+                col = config.get("filter_column")
+                if col and col not in all_columns:
+                    errors.append(f"[{name}] Missing filter_column: '{col}'")
+
+            # SelectorBlock
+            if "SelectorBlock" in str(block_type):
+                col = config.get("choice_col")
+                if col and col not in all_columns:
+                    errors.append(f"[{name}] Missing choice_col: '{col}'")
+
+                choice_map = config.get("choice_map", {})
+                for col in choice_map.values():
+                    if col not in all_columns:
+                        errors.append(f"[{name}] choice_map references missing column: '{col}'")
+
+            # CombineColumnsBlock
+            if "CombineColumnsBlock" in str(block_type):
+                cols = config.get("columns", [])
+                for col in cols:
+                    if col not in all_columns:
+                        errors.append(f"[{name}] CombineColumnsBlock requires column: '{col}'")
+
+        return ValidationResult(valid=(len(errors) == 0), errors=errors)

sdg_hub/flow_runner.py

@@ -0,0 +1,216 @@
+"""Script for running data generation flows with configurable parameters."""
+
+# Standard
+import os
+
+# Third Party
+from datasets import load_dataset
+from openai import OpenAI
+import click
+
+# First Party
+from sdg_hub.flow import Flow
+from sdg_hub.logger_config import setup_logger
+from sdg_hub.sdg import SDG
+
+
+logger = setup_logger(__name__)
+
+
+def run_flow(
+    ds_path: str,
+    save_path: str,
+    endpoint: str,
+    flow_path: str,
+    checkpoint_dir: str,
+    batch_size: int = 8,
+    num_workers: int = 32,
+    save_freq: int = 2,
+    debug: bool = False,
+    dataset_start_index: int = 0,
+    dataset_end_index: int = None,
+) -> None:
+    """Process the dataset using the specified configuration.
+
+    Parameters
+    ----------
+    ds_path : str
+        Path to the dataset file.
+    save_path : str
+        Path where the output will be saved.
+    endpoint : str
+        API endpoint for data processing.
+    flow_path : str
+        Path to the flow configuration file.
+    checkpoint_dir : str
+        Directory path for saving checkpoints.
+    batch_size : int, optional
+        Batch size for processing, by default 8.
+    num_workers : int, optional
+        Number of worker processes to use, by default 32.
+    save_freq : int, optional
+        Frequency (in batches) at which to save checkpoints, by default 2.
+    debug : bool, optional
+        If True, enables debug mode with a smaller dataset subset, by default False.
+
+    Returns
+    -------
+    None
+
+    Raises
+    ------
+    FileNotFoundError
+        If the flow configuration file is not found.
+    """
+    logger.info(f"Generation configuration: {locals()}\n\n")
+    ds = load_dataset("json", data_files=ds_path, split="train")
+    if dataset_start_index is not None and dataset_end_index is not None:
+        ds = ds.select(range(dataset_start_index, dataset_end_index))
+        logger.info(f"Dataset sliced from {dataset_start_index} to {dataset_end_index}")
+    if debug:
+        ds = ds.shuffle(seed=42).select(range(30))
+        logger.info("Debug mode enabled. Using a subset of the dataset.")
+
+    openai_api_key = os.environ.get("OPENAI_API_KEY", "EMPTY")
+    openai_api_base = endpoint
+
+    client = OpenAI(
+        api_key=openai_api_key,
+        base_url=openai_api_base,
+    )
+
+    if not os.path.exists(flow_path):
+        raise FileNotFoundError(f"Flow file not found: {flow_path}")
+
+    flow = Flow(client).get_flow_from_file(flow_path)
+    sdg = SDG(
+        flows=[flow],
+        num_workers=num_workers,
+        batch_size=batch_size,
+        save_freq=save_freq,
+    )
+
+    generated_data = sdg.generate(ds, checkpoint_dir=checkpoint_dir)
+    if dataset_end_index is not None and dataset_start_index is not None:
+        save_path = save_path.replace(".jsonl", f"_{dataset_start_index}_{dataset_end_index}.jsonl")
+    generated_data.to_json(save_path, orient="records", lines=True)
+    logger.info(f"Data saved to {save_path}")
+
+
+@click.command()
+@click.option(
+    "--ds_path",
+    type=click.Path(exists=True),
+    required=True,
+    help="Path to the dataset.",
+)
+@click.option(
+    "--bs",
+    type=int,
+    default=8,
+    show_default=True,
+    help="Batch size for processing.",
+)
+@click.option(
+    "--num_workers",
+    type=int,
+    default=32,
+    show_default=True,
+    help="Number of worker processes to use.",
+)
+@click.option(
+    "--save_path",
+    type=click.Path(),
+    required=True,
+    help="Path to save the output.",
+)
+@click.option(
+    "--endpoint",
+    type=str,
+    required=True,
+    help="API endpoint for data processing.",
+)
+@click.option(
+    "--flow",
+    type=click.Path(exists=True),
+    required=True,
+    help="Flow configuration for the process.",
+)
+@click.option(
+    "--checkpoint_dir",
+    type=click.Path(),
+    required=True,
+    help="Path to save checkpoints.",
+)
+@click.option(
+    "--save_freq",
+    type=int,
+    default=2,
+    show_default=True,
+    help="Frequency to save checkpoints.",
+)
+@click.option(
+    "--debug",
+    is_flag=True,
+    help="Enable debug mode with a smaller dataset subset.",
+)
+@click.option("--dataset_start_index", type=int, default=0, help="Start index of the dataset.")
+@click.option("--dataset_end_index", type=int, default=None, help="End index of the dataset.")
+def main(
+    ds_path: str,
+    bs: int,
+    num_workers: int,
+    save_path: str,
+    endpoint: str,
+    flow: str,
+    checkpoint_dir: str,
+    save_freq: int,
+    debug: bool,
+    dataset_start_index: int,
+    dataset_end_index: int,
+) -> None:
+    """CLI entry point for running data generation flows.
+
+    Parameters
+    ----------
+    ds_path : str
+        Path to the dataset file.
+    bs : int
+        Batch size for processing.
+    num_workers : int
+        Number of worker processes to use.
+    save_path : str
+        Path where the output will be saved.
+    endpoint : str
+        API endpoint for data processing.
+    flow : str
+        Path to the flow configuration file.
+    checkpoint_dir : str
+        Directory path for saving checkpoints.
+    save_freq : int
+        Frequency (in batches) at which to save checkpoints.
+    debug : bool
+        If True, enables debug mode with a smaller dataset subset.
+
+    Returns
+    -------
+    None
+    """
+    run_flow(
+        ds_path=ds_path,
+        batch_size=bs,
+        num_workers=num_workers,
+        save_path=save_path,
+        endpoint=endpoint,
+        flow_path=flow,
+        checkpoint_dir=checkpoint_dir,
+        save_freq=save_freq,
+        debug=debug,
+        dataset_start_index=dataset_start_index,
+        dataset_end_index=dataset_end_index,
+    )
+
+
+if __name__ == "__main__":
+    # pylint: disable=no-value-for-parameter
+    main()

sdg_hub/flows/generation/knowledge/synth_knowledge1.5.yaml

@@ -12,7 +12,10 @@
     output_cols:
       - summary_detailed
   gen_kwargs:
-    max_tokens: 2048
+    max_tokens: 4096
+    temperature: 0.7
+    seed: 7452
+    n: 50
 
 - block_type: LLMBlock
   block_config:
@@ -22,7 +25,9 @@
     output_cols:
       - summary_atomic_facts
   gen_kwargs:
-    max_tokens: 2048
+    max_tokens: 4096
+    temperature: 0.7
+    seed: 7452
 
 - block_type: LLMBlock
   block_config:
@@ -32,7 +37,9 @@
     output_cols:
       - summary_extractive
   gen_kwargs:
-    max_tokens: 2048
+    max_tokens: 4096
+    temperature: 0.7
+    seed: 7452
 
 - block_type: FlattenColumnsBlock
   block_config:
@@ -55,19 +62,29 @@
 - block_type: LLMBlock
   block_config:
     block_name: knowledge generation
-    config_path: configs/knowledge/generate_questions_responses.yaml
+    config_path: configs/knowledge/generate_questions.yaml
     model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
     output_cols:
       - question
-      - response
     parser_kwargs:
       parser_name: custom
-      parsing_pattern: "\\[(?:Question|QUESTION)\\]\\s*(.*?)\\s*\\[(?:Answer|ANSWER)\\]\\s*(.*?)\\s*(?=\\[(?:Question|QUESTION)\\]|$)"
-      parser_cleanup_tags:
-        - "[END]"
+      parsing_pattern: "\\[(?:Question|QUESTION)\\]\\s*(.*?)\\s*(?=\\[(?:Question|QUESTION)\\]|$)"
+  gen_kwargs:
+    temperature: 0.7
+    max_tokens: 100
+    seed: 7452
+
+- block_type: LLMBlock
+  block_config:
+    block_name: knowledge generation
+    config_path: configs/knowledge/generate_responses.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - response
   gen_kwargs:
-    temperature: 0.0
+    temperature: 0.7
     max_tokens: 2048
+    seed: 7452
 
 - block_type: LLMBlock
   block_config: