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

sdg_hub/flow_runner.py

@@ -0,0 +1,437 @@
+"""Script for running data generation flows with configurable parameters."""
+
+# Standard
+from importlib import resources
+from typing import Optional
+import os
+import sys
+import traceback
+
+# Third Party
+from datasets import load_dataset
+from openai import OpenAI
+import click
+import yaml
+
+# First Party
+from sdg_hub.flow import Flow
+from sdg_hub.logger_config import setup_logger
+from sdg_hub.sdg import SDG
+from sdg_hub.utils.error_handling import (
+    APIConnectionError,
+    DataGenerationError,
+    DataSaveError,
+    DatasetLoadError,
+    FlowConfigurationError,
+    FlowRunnerError,
+)
+from sdg_hub.utils.path_resolution import resolve_path
+
+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: Optional[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.
+    dataset_start_index : int, optional
+        Start index for dataset slicing, by default 0.
+    dataset_end_index : Optional[int], optional
+        End index for dataset slicing, by default None.
+
+    Returns
+    -------
+    None
+
+    Raises
+    ------
+    DatasetLoadError
+        If the dataset cannot be loaded or processed.
+    FlowConfigurationError
+        If the flow configuration is invalid or cannot be loaded.
+    APIConnectionError
+        If connection to the API endpoint fails.
+    DataGenerationError
+        If data generation fails during processing.
+    DataSaveError
+        If saving the generated data fails.
+    """
+    logger.info(f"Generation configuration: {locals()}\n\n")
+
+    try:
+        # Load and validate dataset
+        try:
+            ds = load_dataset("json", data_files=ds_path, split="train")
+            logger.info(
+                f"Successfully loaded dataset from {ds_path} with {len(ds)} rows"
+            )
+        except Exception as e:
+            raise DatasetLoadError(
+                f"Failed to load dataset from '{ds_path}'. "
+                f"Please check if the file exists and is a valid JSON file.",
+                details=str(e),
+            ) from e
+
+        # Apply dataset slicing if specified
+        try:
+            if dataset_start_index is not None and dataset_end_index is not None:
+                if dataset_start_index >= len(ds) or dataset_end_index > len(ds):
+                    raise DatasetLoadError(
+                        f"Dataset slice indices ({dataset_start_index}, {dataset_end_index}) "
+                        f"are out of bounds for dataset with {len(ds)} rows"
+                    )
+                if dataset_start_index >= dataset_end_index:
+                    raise DatasetLoadError(
+                        f"Start index ({dataset_start_index}) must be less than end index ({dataset_end_index})"
+                    )
+                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:
+                if len(ds) < 30:
+                    logger.warning(
+                        f"Debug mode requested 30 samples but dataset only has {len(ds)} rows"
+                    )
+                ds = ds.shuffle(seed=42).select(range(min(30, len(ds))))
+                logger.info(
+                    f"Debug mode enabled. Using {len(ds)} samples from the dataset."
+                )
+        except DatasetLoadError:
+            raise
+        except Exception as e:
+            raise DatasetLoadError(
+                "Failed to process dataset slicing or debug mode.", details=str(e)
+            ) from e
+
+        # Validate API configuration
+        openai_api_key = os.environ.get("OPENAI_API_KEY")
+        if not openai_api_key or openai_api_key == "EMPTY":
+            logger.warning("OPENAI_API_KEY not set or is 'EMPTY'. API calls may fail.")
+
+        openai_api_base = endpoint
+        if not openai_api_base:
+            raise APIConnectionError("API endpoint cannot be empty")
+
+        # Initialize OpenAI client
+        try:
+            client = OpenAI(
+                api_key=openai_api_key or "EMPTY",
+                base_url=openai_api_base,
+            )
+            # test connection with a model list
+            models = client.models.list()
+            logger.info(f"Initialized OpenAI client with endpoint: {openai_api_base}")
+            logger.info(f"Available models: {[model.id for model in models.data]}")
+        except Exception as e:
+            raise APIConnectionError(
+                f"Failed to initialize OpenAI client with endpoint '{openai_api_base}'. "
+                f"Please check if the endpoint is valid and accessible.",
+                details=str(e),
+            ) from e
+
+        # Load and validate flow configuration
+        try:
+            base_path = str(resources.files(__package__))
+            flow_path = resolve_path(flow_path, [".", base_path])
+            if not os.path.exists(flow_path):
+                raise FlowConfigurationError(
+                    f"Flow configuration file not found: {flow_path}"
+                )
+
+            # Validate flow file is readable YAML
+            try:
+                with open(flow_path, "r", encoding="utf-8") as f:
+                    flow_config = yaml.safe_load(f)
+                if not flow_config:
+                    raise FlowConfigurationError(
+                        f"Flow configuration file is empty: {flow_path}"
+                    )
+                logger.info(f"Successfully loaded flow configuration from {flow_path}")
+            except yaml.YAMLError as e:
+                raise FlowConfigurationError(
+                    f"Flow configuration file '{flow_path}' contains invalid YAML.",
+                    details=str(e),
+                ) from e
+            except Exception as e:
+                raise FlowConfigurationError(
+                    f"Failed to read flow configuration file '{flow_path}'.",
+                    details=str(e),
+                ) from e
+
+            flow = Flow(client).get_flow_from_file(flow_path)
+            logger.info("Successfully initialized flow from configuration")
+        except FlowConfigurationError:
+            raise
+        except Exception as e:
+            raise FlowConfigurationError(
+                f"Failed to create flow from configuration file '{flow_path}'. "
+                f"Please check the flow configuration format and block definitions.",
+                details=str(e),
+            ) from e
+
+        # Initialize SDG and generate data
+        try:
+            sdg = SDG(
+                flows=[flow],
+                num_workers=num_workers,
+                batch_size=batch_size,
+                save_freq=save_freq,
+            )
+            logger.info(
+                f"Initialized SDG with {num_workers} workers, batch size {batch_size}"
+            )
+
+            # Ensure checkpoint directory exists if specified
+            if checkpoint_dir and not os.path.exists(checkpoint_dir):
+                os.makedirs(checkpoint_dir, exist_ok=True)
+                logger.info(f"Created checkpoint directory: {checkpoint_dir}")
+
+            generated_data = sdg.generate(ds, checkpoint_dir=checkpoint_dir)
+
+            if generated_data is None or len(generated_data) == 0:
+                raise DataGenerationError(
+                    "Data generation completed but no data was generated. "
+                    "This may indicate issues with the flow configuration or input data."
+                )
+
+            logger.info(f"Successfully generated {len(generated_data)} rows of data")
+
+        except Exception as e:
+            if isinstance(e, DataGenerationError):
+                raise
+            raise DataGenerationError(
+                "Data generation failed during processing. This could be due to:"
+                "\n- API connection issues with the endpoint"
+                "\n- Invalid flow configuration or block parameters"
+                "\n- Insufficient system resources (try reducing batch_size or num_workers)"
+                "\n- Input data format incompatibility",
+                details=f"Endpoint: {openai_api_base}, Error: {e}",
+            ) from e
+
+        # Save generated data
+        try:
+            # Adjust save path for dataset slicing
+            final_save_path = save_path
+            if dataset_end_index is not None and dataset_start_index is not None:
+                final_save_path = save_path.replace(
+                    ".jsonl", f"_{dataset_start_index}_{dataset_end_index}.jsonl"
+                )
+
+            # Ensure save directory exists
+            save_dir = os.path.dirname(final_save_path)
+            if save_dir and not os.path.exists(save_dir):
+                os.makedirs(save_dir, exist_ok=True)
+                logger.info(f"Created save directory: {save_dir}")
+
+            generated_data.to_json(final_save_path, orient="records", lines=True)
+            logger.info(f"Data successfully saved to {final_save_path}")
+
+        except Exception as e:
+            raise DataSaveError(
+                f"Failed to save generated data to '{final_save_path}'. "
+                f"Please check write permissions and disk space.",
+                details=str(e),
+            ) from e
+
+    except (
+        DatasetLoadError,
+        FlowConfigurationError,
+        APIConnectionError,
+        DataGenerationError,
+        DataSaveError,
+    ):
+        # Re-raise our custom exceptions with their detailed messages
+        raise
+    except Exception as e:
+        # Catch any unexpected errors
+        logger.error(f"Unexpected error during flow execution: {e}")
+        logger.error(f"Traceback: {traceback.format_exc()}")
+        raise FlowRunnerError(
+            "An unexpected error occurred during flow execution. "
+            "Please check the logs for more details.",
+            details=str(e),
+        ) from e
+
+
+@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: Optional[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.
+    dataset_start_index : int
+        Start index for dataset slicing.
+    dataset_end_index : Optional[int]
+        End index for dataset slicing.
+
+    Returns
+    -------
+    None
+    """
+    try:
+        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,
+        )
+    except (
+        DatasetLoadError,
+        FlowConfigurationError,
+        APIConnectionError,
+        DataGenerationError,
+        DataSaveError,
+        FlowRunnerError,
+    ) as e:
+        logger.error(f"Flow execution failed: {e}")
+        click.echo(f"Error: {e}", err=True)
+        sys.exit(1)
+    except KeyboardInterrupt:
+        logger.info("Flow execution interrupted by user")
+        click.echo("Flow execution interrupted by user", err=True)
+        sys.exit(130)  # Standard exit code for SIGINT
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}")
+        logger.error(f"Traceback: {traceback.format_exc()}")
+        click.echo(
+            f"Unexpected error occurred. Please check the logs for details. Error: {e}",
+            err=True,
+        )
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    # pylint: disable=no-value-for-parameter
+    main()

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

@@ -12,7 +12,9 @@
     output_cols:
       - summary_detailed
   gen_kwargs:
-    max_tokens: 2048
+    max_tokens: 4096
+    temperature: 0.7
+    n: 50
 
 - block_type: LLMBlock
   block_config:
@@ -22,7 +24,8 @@
     output_cols:
       - summary_atomic_facts
   gen_kwargs:
-    max_tokens: 2048
+    max_tokens: 4096
+    temperature: 0.7
 
 - block_type: LLMBlock
   block_config:
@@ -32,7 +35,8 @@
     output_cols:
       - summary_extractive
   gen_kwargs:
-    max_tokens: 2048
+    max_tokens: 4096
+    temperature: 0.7
 
 - block_type: FlattenColumnsBlock
   block_config:
@@ -55,18 +59,26 @@
 - 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
+
+- 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
 
 - block_type: LLMBlock

sdg_hub/flows/generation/skills/{agentic_improve_skill.yaml → improve_responses.yaml}

@@ -2,34 +2,34 @@
   block_config:
     block_name: router
     config_path: configs/skills/router.yaml
-    model_id: skill-classifier-v3-clm
+    model_id: meta-llama/Llama-3.3-70B-Instruct
     output_cols:
       - route
   gen_kwargs:
     temperature: 0
-    max_tokens: 1
+    max_tokens: 5
     extra_body:
-      allowed_token_ids:  
-        - 32001
-        - 32002
-        - 32003
-        - 32004
-        - 32005
-        - 32006
-        - 32007
-        - 32008
+      guided_choice:
+        - "coding"
+        - "extraction"
+        - "humanities"
+        - "math"
+        - "reasoning"
+        - "roleplay"
+        - "STEM"
+        - "writing"
 - block_type: SamplePopulatorBlock
   block_config:
     block_name: icl_populator
     config_paths: 
-      - configs/skills/_A_.yaml
-      - configs/skills/_B_.yaml
-      - configs/skills/_C_.yaml
-      - configs/skills/_D_.yaml
-      - configs/skills/_E_.yaml
-      - configs/skills/_F_.yaml
-      - configs/skills/_G_.yaml
-      - configs/skills/_H_.yaml
+      - configs/skills/icl_examples/coding.yaml
+      - configs/skills/icl_examples/extraction.yaml
+      - configs/skills/icl_examples/humanities.yaml
+      - configs/skills/icl_examples/math.yaml
+      - configs/skills/icl_examples/reasoning.yaml
+      - configs/skills/icl_examples/roleplay.yaml
+      - configs/skills/icl_examples/STEM.yaml
+      - configs/skills/icl_examples/writing.yaml
     column_name: route
     batch_kwargs:
       num_procs: 8
@@ -37,8 +37,7 @@
   block_config:
     block_name: analyzer
     config_path: configs/skills/analyzer.yaml
-    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
-    model_prompt: <s> [INST] {prompt} [/INST]
+    model_id: meta-llama/Llama-3.3-70B-Instruct
     output_cols:
       - analysis
       - rubric
@@ -46,24 +45,21 @@
   block_config:
     block_name: critic
     config_path: configs/skills/critic.yaml
-    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
-    model_prompt: <s> [INST] {prompt} [/INST]
+    model_id: meta-llama/Llama-3.3-70B-Instruct
     output_cols: 
       - critique
 - block_type: LLMBlock
   block_config:
     block_name: planner
     config_path: configs/skills/planner.yaml
-    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
-    model_prompt: <s> [INST] {prompt} [/INST]
+    model_id: meta-llama/Llama-3.3-70B-Instruct
     output_cols: 
       - plan
 - block_type: LLMBlock
   block_config:
     block_name: revised_responder
     config_path: configs/skills/revised_responder.yaml
-    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
-    model_prompt: <s> [INST] {prompt} [/INST]
+    model_id: meta-llama/Llama-3.3-70B-Instruct
     output_cols: 
       - revised_response
   drop_columns:
@@ -78,8 +74,7 @@
   block_config:
     block_name: judge
     config_path: configs/skills/judge.yaml
-    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
-    model_prompt: <s> [INST] {prompt} [/INST]
+    model_id: meta-llama/Llama-3.3-70B-Instruct
     output_cols:
       - judgement
       - verdict
@@ -100,9 +95,9 @@
        Assistant A: "response"
        Assistant B: "revised_response"
     choice_col: verdict
-    output_col: chosen_reponse
+    output_col: chosen_response
     batch_kwargs:
       num_procs: 8
   drop_columns:
     - judgemnent
-    - verdict
+    - verdict

sdg_hub/flows/generation/skills/synth_skills.yaml

@@ -2,7 +2,7 @@
   block_config:
     block_name: gen_questions
     config_path: configs/skills/freeform_questions.yaml
-    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    model_id: meta-llama/Llama-3.3-70B-Instruct
     output_cols:
       - question
     batch_kwargs:
@@ -13,7 +13,7 @@
   block_config:
     block_name: eval_questions
     config_path: configs/skills/evaluate_freeform_questions.yaml
-    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    model_id: meta-llama/Llama-3.3-70B-Instruct
     output_cols:
       - evaluation
       - score
@@ -34,14 +34,14 @@
   block_config:
     block_name: gen_responses
     config_path: configs/skills/freeform_responses.yaml
-    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    model_id: meta-llama/Llama-3.3-70B-Instruct
     output_cols:
       - response
 - block_type: LLMBlock
   block_config:
     block_name: evaluate_qa_pair
     config_path: configs/skills/evaluate_freeform_pair.yaml
-    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    model_id: meta-llama/Llama-3.3-70B-Instruct
     output_cols:
       - evaluation
       - score