sdg-hub 0.4.1__tar.gz → 0.4.2__tar.gz

{sdg_hub-0.4.1 → sdg_hub-0.4.2}/.github/workflows/integration-test.yml

@@ -20,6 +20,7 @@ on:
     branches:
       - "main"
       - "release-**"
+    types: [opened, synchronize, reopened, labeled]
     paths:
       # Only trigger on changes to relevant flows and examples (EXTEND THIS):
       - 'src/sdg_hub/flows/qa_generation/document_grounded_qa/enhanced_multi_summary_qa/**'
@@ -47,11 +48,13 @@ jobs:
     # Require manual approval before running (via GitHub Environment)
     environment: integration-tests
     # Skip fork PRs (they can't access environment secrets anyway)
+    # Also check for 'run-integration-tests' label on labeled events
     if: |
       github.event_name == 'workflow_dispatch' ||
       github.event_name == 'push' ||
       (github.event_name == 'pull_request' &&
-       github.event.pull_request.head.repo.full_name == github.repository)
+       github.event.pull_request.head.repo.full_name == github.repository &&
+       (github.event.action != 'labeled' || contains(github.event.pull_request.labels.*.name, 'run-integration-tests')))
     strategy:
       matrix:
         python:

{sdg_hub-0.4.1/src/sdg_hub.egg-info → sdg_hub-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sdg_hub
-Version: 0.4.1
+Version: 0.4.2
 Summary: Synthetic Data Generation
 Author-email: Red Hat AI Innovation <abhandwa@redhat.com>
 License: Apache-2.0

{sdg_hub-0.4.1 → sdg_hub-0.4.2}/docs/concepts.md

@@ -148,9 +148,22 @@ Every block validates data at runtime:
 ## 🚀 Best Practices
 
 ### 1. Start Small
-- Use `dry_run()` to test with small samples
+- Use `dry_run()` to test with small samples before processing full datasets
+- Add `enable_time_estimation=True` to predict execution time for the complete dataset
 - Validate your pipeline before scaling up
 
+```python
+# Test AND estimate in one call
+result = flow.dry_run(dataset, sample_size=5, enable_time_estimation=True, max_concurrency=100)
+
+# Access dry run results
+print(f"Tested with {result['sample_size']} samples")
+print(f"Output columns: {result['final_dataset']['columns']}")
+
+# Time estimation is automatically displayed in a Rich table format
+# No need to access it programmatically - the table shows all estimation details
+```
+
 ### 2. Layer Validation
 - Use basic block composition (PromptBuilder → LLMChat → Parser → Filter) to assess quality
 - Implement filtering to maintain data standards

{sdg_hub-0.4.1 → sdg_hub-0.4.2}/docs/flows/discovery.md

@@ -67,7 +67,44 @@ for flow_name in all_flows:
 
 ### Getting Flow Information
 
-#TODO: Add flow info example
+Access detailed flow metadata and configuration:
+
+```python
+from sdg_hub.core.flow import FlowRegistry, Flow
+
+# Get metadata for a specific flow
+flow_name = "Advanced Document Grounded Question-Answer Generation Flow for Knowledge Tuning"
+metadata = FlowRegistry.get_flow_metadata(flow_name)
+
+if metadata:
+    print(f"Flow: {metadata.name}")
+    print(f"Version: {metadata.version}")
+    print(f"Author: {metadata.author}")
+    print(f"Description: {metadata.description}")
+    print(f"Tags: {', '.join(metadata.tags)}")
+    print(f"Recommended model: {metadata.recommended_models.get('default', 'Not specified')}")
+
+# Load flow and get detailed information
+flow_path = FlowRegistry.get_flow_path(flow_name)
+flow = Flow.from_yaml(flow_path)
+
+# Get comprehensive flow info
+info = flow.get_info()
+print(f"Total blocks: {info['total_blocks']}")
+print(f"Block sequence: {', '.join(info['block_names'])}")
+
+# Get dataset requirements
+requirements = flow.get_dataset_requirements()
+if requirements:
+    print(f"Required columns: {requirements.required_columns}")
+    print(f"Description: {requirements.description}")
+    print(f"Min samples: {requirements.min_samples}")
+
+# Get model recommendations
+recommendations = flow.get_model_recommendations()
+print(f"Default model: {recommendations.get('default')}")
+print(f"Compatible models: {recommendations.get('compatible', [])}")
+```
 
 ### Getting Flow Paths
 

{sdg_hub-0.4.1 → sdg_hub-0.4.2}/docs/flows/overview.md

@@ -292,6 +292,41 @@ print(f"Output columns: {dry_result['final_dataset']['columns']}")
 print(f"Sample output: {dry_result['sample_output']}")
 ```
 
+### Time Estimation
+
+Predict execution time for your full dataset before running:
+
+```python
+# Get dry run results AND time estimation in one call
+result = flow.dry_run(
+    dataset, 
+    sample_size=5, 
+    enable_time_estimation=True,
+    max_concurrency=100
+)
+
+# Time estimation is automatically displayed in a Rich table format
+# The table shows estimated time, total API requests, and per-block breakdowns
+print(f"Dry run completed with {result['sample_size']} samples")
+print(f"Output columns: {result['final_dataset']['columns']}")
+```
+
+**How It Works:**
+
+The estimation uses 2 dry runs to accurately predict execution time:
+- Extracts startup overhead (one-time costs)
+- Calculates per-sample throughput (variable costs)
+- Uses linear regression to separate fixed from variable costs
+
+**Accuracy:**
+- Includes a 20% conservative buffer to account for API variability
+- Typical accuracy: within 15-40% of actual runtime depending on workload
+- Better to finish early than run over time!
+
+**When to Use:**
+- Before processing with your full dataset
+- To identify bottleneck blocks and optimize your pipeline
+
 ### Runtime Parameters
 
 Runtime parameters allow you to customize block behavior at execution time without modifying flow YAML files. You can override global parameters for all blocks or configure specific blocks individually.

{sdg_hub-0.4.1 → sdg_hub-0.4.2}/docs/quick-start.md

@@ -62,11 +62,14 @@ dataset = Dataset.from_dict({
     'icl_response_3': ['Java provides platform independence and strong object-oriented features.']
 })
 
-# Test with a small sample first (recommended!)
-print("🧪 Running dry run...")
-dry_result = flow.dry_run(dataset, sample_size=1)
+# Test with a small sample AND get time estimate (recommended!)
+print("🧪 Running dry run with time estimation...")
+dry_result = flow.dry_run(dataset, sample_size=5, enable_time_estimation=True, max_concurrency=100)
 print(f"✅ Dry run completed in {dry_result['execution_time_seconds']:.2f}s")
 print(f"📊 Output columns: {list(dry_result['final_dataset']['columns'])}")
+
+# Time estimation is automatically displayed in a Rich table format
+# The table shows estimated time, total API calls, and per-block breakdowns
 ```
 
 ## 📊 Step 3: Generate Synthetic Data

{sdg_hub-0.4.1 → sdg_hub-0.4.2}/examples/knowledge_tuning/enhanced_summary_knowledge_tuning/.env.example

@@ -16,6 +16,7 @@ VLLM_MODEL=hosted_vllm/meta-llama/Llama-3.3-70B-Instruct
 VLLM_API_BASE=http://localhost:8000/v1
 VLLM_API_KEY=EMPTY
 ENABLE_REASONING=false
+MAX_CONCURRENCY=50
 # =============================================================================
 # OPENAI CONFIGURATION
 # =============================================================================
@@ -38,7 +39,9 @@ MAAS_API_KEY=your-maas-api-key-here
 # =============================================================================
 # DATA CONFIGURATION
 # =============================================================================
-SEED_DATA_PATH=seed_data_val.jsonl
+SEED_DATA_PATH=seed_data.jsonl
+# Set this for subsampling the seed data. Useful for debugging or running validation
+SEED_DATA_SUBSAMPLE=24
 OUTPUT_DATA_FOLDER=output_data
 RUN_ON_VALIDATION_SET=true
 NUMBER_OF_SUMMARIES=50

sdg_hub-0.4.2/examples/knowledge_tuning/enhanced_summary_knowledge_tuning/document_pre_processing.ipynb

@@ -0,0 +1,214 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "83f458de",
+   "metadata": {},
+   "source": [
+    "# Document Pre-processing for Knowledge Tuning\n",
+    "\n",
+    "## Overview\n",
+    "\n",
+    "This notebook demonstrates a complete document preprocessing pipeline designed specifically for **knowledge tuning** with sdg-hub. \n",
+    "\n",
+    "## What This Notebook Does\n",
+    "\n",
+    "This preprocessing pipeline transforms raw documents (PDFs, Word docs, etc.) into seed data for data generation:\n",
+    "\n",
+    "1. **Document Parsing**: Converts raw documents to structured markdown format\n",
+    "2. **Chunking**: Splits documents into manageable chunks while preserving structure and context\n",
+    "3. **Seed Data Creation**: Formats chunks with in-context learning (ICL) templates for effective knowledge tuning\n",
+    "\n",
+    "## Prerequisites\n",
+    "\n",
+    "- We will use the existing InstructLab document parser (`docparser_v2.py`) and Document parsing configuration (`docling_v2_config.yaml`)\n",
+    "- Raw pdf documents in the `document_collection/` directory\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "daa22c74",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Step 1: Document Processing Pipeline\n",
+    "# Define the directory containing raw documents to be processed\n",
+    "data_dir = 'document_collection/'\n",
+    "\n",
+    "# Run the document parser to convert documents to markdown\n",
+    "# - input-dir: Directory containing source documents\n",
+    "# - output-dir: Directory where processed markdown files will be saved\n",
+    "# - c: Configuration file specifying parsing parameters\n",
+    "!python ../instructlab/docparser_v2.py --input-dir {data_dir} --output-dir {data_dir} -c ../instructlab/docling_v2_config.yaml"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "295749b5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Step 2: Install Required Dependencies\n",
+    "# Install packages needed for document processing and text chunking\n",
+    "\n",
+    "%pip install docling markdown-it-py\n",
+    "%pip install --upgrade transformers"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dd8a4a2a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Step 3: Load Processed Document\n",
+    "import glob\n",
+    "\n",
+    "# In our example above docling step produces markdown of all the pdf files in the document_collection\n",
+    "with open(glob.glob(f'{data_dir}/*.md')[0], 'r') as f:\n",
+    "    text = f.read()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7614dc73",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Step 4: Text Chunking and Dataset Creation\n",
+    "\n",
+    "from markdown_it import MarkdownIt  \n",
+    "from typing import List\n",
+    "import datasets  \n",
+    "\n",
+    "\n",
+    "def chunk_markdown(\n",
+    "    text: str,\n",
+    "    max_tokens: int = 200,\n",
+    "    overlap: int = 50\n",
+    ") -> List[str]:\n",
+    "    \"\"\"\n",
+    "    Splits Markdown text into chunks at block-level elements\n",
+    "    (headings, paragraphs, lists, tables, code, blockquotes).\n",
+    "    Adds overlap (in words) between all consecutive chunks.\n",
+    "    \n",
+    "    Args:\n",
+    "        text: The markdown text to be chunked\n",
+    "        max_tokens: Maximum number of words per chunk\n",
+    "        overlap: Number of overlapping words between consecutive chunks\n",
+    "    \n",
+    "    Returns:\n",
+    "        List of text chunks with specified overlap\n",
+    "    \"\"\"\n",
+    "\n",
+    "    # Initialize markdown parser to understand document structure\n",
+    "    md = MarkdownIt()\n",
+    "    tokens = md.parse(text)\n",
+    "\n",
+    "    # Group tokens into block-level segments to preserve markdown structure\n",
+    "    # This ensures we don't split in the middle of headings, lists, etc.\n",
+    "    blocks = []\n",
+    "    buf = []\n",
+    "    for tok in tokens:\n",
+    "        if tok.block and tok.type.endswith(\"_open\"):\n",
+    "            buf = []\n",
+    "        elif tok.block and tok.type.endswith(\"_close\"):\n",
+    "            if buf:\n",
+    "                blocks.append(\"\\n\".join(buf).strip())\n",
+    "                buf = []\n",
+    "        elif tok.content:\n",
+    "            buf.append(tok.content)\n",
+    "    if buf:\n",
+    "        blocks.append(\"\\n\".join(buf).strip())\n",
+    "\n",
+    "    # Split blocks into chunks with overlap to maintain context continuity\n",
+    "    chunks = []\n",
+    "    current_words = []\n",
+    "    for block in blocks:\n",
+    "        words = block.split()\n",
+    "        for w in words:\n",
+    "            current_words.append(w)\n",
+    "            if len(current_words) >= max_tokens:\n",
+    "                # Emit a complete chunk\n",
+    "                chunks.append(\" \".join(current_words))\n",
+    "                # Prepare next buffer with overlap from the end of this chunk\n",
+    "                # This ensures context continuity between chunks\n",
+    "                current_words = current_words[-overlap:] if overlap > 0 else []\n",
+    "\n",
+    "    # Add any remaining words as the final chunk\n",
+    "    if current_words:\n",
+    "        chunks.append(\" \".join(current_words))\n",
+    "\n",
+    "    return chunks\n",
+    "\n",
+    "\n",
+    "chunks = chunk_markdown(text, max_tokens=5000, overlap=1000)\n",
+    "\n",
+    "\n",
+    "# Prepare seed data for the SDG-Hub knowledge pipeline.\n",
+    "# \n",
+    "# The seed data requires the following fields:\n",
+    "#   - document_outline: A concise title or summary that accurately represents the entire document.\n",
+    "#     For documents covering multiple themes, consider providing multiple outlines (one per section).\n",
+    "#   - icl_document: A representative sample extract from the document. This may include tables, code snippets, definitions, etc.\n",
+    "#   - icl_query_1, icl_query_2, icl_query_3: Three questions based on the icl_document sample.\n",
+    "#   - domain: The domain or subject area of the document.\n",
+    "#\n",
+    "# The code below creates a HuggingFace Dataset from the document chunks,\n",
+    "# then maps the required ICL fields to each entry, and finally saves the result as a JSONL file.\n",
+    "\n",
+    "seed_data = datasets.Dataset.from_dict({'document': chunks})\n",
+    "\n",
+    "icl = {\n",
+    "    \"document_outline\": \"The document contains excerpts from FINTRAC regulations designed to combat money laundering and terrorist financing in Canada\",\n",
+    "    \"icl_document\": \"## Overview\\n\\nThis guidance came into effect on June 1, 2021.\\n\\n\\nThis guidance explains the methods that can be used by reporting entities\\n(REs) to verify the identity of a person or an entity.\\n\\n\\n## 1. Meaning of verifying the identity of a person or an entity\\n\\nIt means to use the methods described in this guidance to ensure that the\\ninformation in an identification document or from other informational\\nsources matches the information that the person or entity provided.\\n\\n\\nVerifying identity is a foundational element of Canada's anti-money\\nlaundering and anti-terrorist financing regime and a key component of an\\nRE's relationship with clients. It helps you to know your clients and to\\nunderstand and assess any risk that may be associated to their\\ntransactions or activities.\\n\\n\\n## 2. How to verify the identity of a person\\n\\nYou can use any of the 5 methods described below to identify a person:\\n\\n- 2.1 Government-issued photo identification method\\n\\n- 2.2 Credit file method\\n\\n- 2.3 Dual-process method\\n\\n- 2.4 Affiliate or member method\\n\\n- 2.5 Reliance method\\n\",\n",
+    "    \"icl_query_1\": \"In Canada, what are the methods for verifying someone's identity?\",\n",
+    "    \"icl_query_2\": \"In Canada, why is it important to confirm a client's identity?\",\n",
+    "    \"icl_query_3\": \"In Canada, can I use Reliance method to verify identity of a person?\",\n",
+    "    \"domain\": \"Finance\"\n",
+    "}\n",
+    "\n",
+    "# Map the ICL fields to each document chunk (if you want to use the same ICL for all, as shown here)\n",
+    "seed_data = seed_data.map(lambda x: icl)\n",
+    "\n",
+    "# Save the seed data to a JSONL file for downstream use\n",
+    "seed_data.to_json('seed_data.jsonl', orient='records', lines=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "44f3ff7f",
+   "metadata": {},
+   "source": [
+    "### Next Steps:\n",
+    "- The seed_data.jsonl file is now ready for the knowledge tuning pipeline.\n",
+    "- You can now refer to the [knowledge generation](knowledge_generation.ipynb) notebook"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "sdg_hub",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

{sdg_hub-0.4.1 → sdg_hub-0.4.2}/examples/knowledge_tuning/enhanced_summary_knowledge_tuning/knowledge_generation.ipynb

@@ -4,14 +4,44 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Install SDG\n",
+    "# Enhanced Summary Knowledge Tuning - Data Generation\n",
+    "\n",
+    "## Overview\n",
+    "\n",
+    "This notebook demonstrates how to generate high-quality knowledge tuning datasets using the SDG Hub framework. It creates multiple types of document augmentations and corresponding question-answer pairs that can be used to train or fine-tune language models for enhanced summarization and knowledge extraction capabilities.\n",
+    "\n",
+    "## What This Notebook Does\n",
+    "\n",
+    "This notebook will:\n",
+    "\n",
+    "2. **Generate Four Types of Knowledge Tuning Datasets**:\n",
+    "   - **Extractive Summaries**: Concise summaries that extract key information directly from source documents\n",
+    "   - **Detailed Summaries**: Comprehensive summaries that provide thorough coverage of document content\n",
+    "   - **Key Facts**: Structured fact extraction with corresponding Q&A pairs\n",
+    "   - **Document-Based Q&A**: Question-answer pairs generated directly from document content\n",
+    "\n",
+    "\n",
+    "4. **Output Structured Training Data**:\n",
+    "   - For each augmentation we save JSONL dataset.\n",
+    "   - You can follow [knowledge_mixing](knowledge_mixing.ipynb) to convert it into training dataset\n",
+    "\n",
+    "## Prerequisites\n",
+    "\n",
+    "- SDG Hub installed and configured\n",
+    "- Environment variables set up (see [.env.example](.env.example)). Specifically set the model provider, seed data and output path.\n",
+    "- Document pre-processing completed (run [document_pre_processing.ipynb](document_pre_processing.ipynb) first)\n",
+    "\n",
     "```bash \n",
     "git clone https://github.com/Red-Hat-AI-Innovation-Team/sdg_hub.git\n",
     "cd sdg_hub\n",
     "pip install .[examples]\n",
     "copy the .env.example to .env and set the model endpoint and generation/mixing parameters\n",
     "```\n",
-    "**⚠️ If you haven't already, run the document pre-processing notebook to create the seed data.**"
+    "**⚠️ If you haven't already, run the document pre-processing notebook to create the seed data.**\n",
+    "\n",
+    "## Next Steps\n",
+    "\n",
+    "After running this notebook, use [knowledge_mixing](knowledge_mixing.ipynb) to combine and curate the generated datasets for final model training.\n"
    ]
   },
   {
@@ -50,7 +80,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def create_seed_data(run_on_validation=None, seed_data_path=None):\n",
+    "def create_seed_data_from_quality_benchmark(run_on_validation=None, seed_data_path=None):\n",
     "    \"\"\"\n",
     "    Create seed data from QuALITY Benchmark dataset.\n",
     "    \n",
@@ -135,8 +165,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Create seed data using the function\n",
-    "quality_corpus = create_seed_data().select(range(1))"
+    "# Load seed data. If one is not provided, create it from the quality benchmark dataset.\n",
+    "seed_data_path = os.getenv('SEED_DATA_PATH', 'seed_data.jsonl')\n",
+    "\n",
+    "if not os.path.exists(seed_data_path):\n",
+    "    print(f\"{seed_data_path} not found. Creating seed data...\")\n",
+    "    quality_corpus = create_seed_data_from_quality_benchmark(seed_data_path=seed_data_path)\n",
+    "else:\n",
+    "    print(f\"Loading existing seed data from {seed_data_path}\")\n",
+    "    quality_corpus = load_dataset('json', data_files=seed_data_path, split='train')\n",
+    "\n",
+    "# Subsample the seed data. Useful for debugging.\n",
+    "subsample = int(os.getenv('SEED_DATA_SUBSAMPLE', '0'))\n",
+    "if subsample > 0:\n",
+    "    quality_corpus = quality_corpus.select(range(subsample))"
    ]
   },
   {
@@ -229,23 +271,12 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "# We will use below mapping of flow names to their respective summarization flows\n",
-    "flow_name_map = {\n",
-    "        'Detailed Summary Knowledge Tuning Dataset Generation Flow': 'gen_detailed_summary',\n",
-    "        'Extractive Summary Knowledge Tuning Dataset Generation Flow': 'gen_extractive_summary',\n",
-    "    }"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [],
    "source": [
     "# Get runtime parameters\n",
     "enable_reasoning = os.getenv('ENABLE_REASONING', 'false').lower() in ('1', 'true', 'yes')\n",
-    "number_of_summaries = int(os.getenv('NUMBER_OF_SUMMARIES', '50'))"
+    "number_of_summaries = int(os.getenv('NUMBER_OF_SUMMARIES', '50'))\n",
+    "max_concurrency = int(os.getenv('MAX_CONCURRENCY', '50'))\n",
+    "save_data_path = os.getenv('OUTPUT_DATA_FOLDER', '')"
    ]
   },
   {
@@ -263,21 +294,21 @@
     "flow = set_model_config(flow)\n",
     "number_of_summaries = int(os.getenv('NUMBER_OF_SUMMARIES', '50'))\n",
     "# Generate data for extractive summary\n",
-    "runtime_params = {\n",
-    "    flow_name_map[flow_name]: {\n",
-    "        'n': number_of_summaries\n",
-    "    },\n",
-    "}\n",
     "if enable_reasoning:\n",
     "    # Increase max tokens to accommodate reasoning content\n",
     "    runtime_params = {\n",
     "        'question_generation': {'max_tokens': 1024}, \n",
-    "        flow_name_map[flow_name]: {'n': number_of_summaries, 'max_tokens': 6000}\n",
+    "        'gen_extractive_summary': {'n': number_of_summaries, 'max_tokens': 6000}\n",
     "        }\n",
-    "    \n",
+    "else:\n",
+    "    runtime_params = {\n",
+    "    'gen_extractive_summary': {\n",
+    "        'n': number_of_summaries\n",
+    "    }\n",
+    "}\n",
+    "\n",
+    "extractive_summary_generated_data = flow.generate(quality_corpus, runtime_params=runtime_params, max_concurrency=max_concurrency)\n",
     "\n",
-    "extractive_summary_generated_data = flow.generate(quality_corpus, runtime_params=runtime_params, max_concurrency=2)\n",
-    "save_data_path = os.getenv('OUTPUT_DATA_FOLDER', '')\n",
     "extractive_summary_generated_data.to_json(os.path.join(save_data_path, 'extractive_summary', 'gen.jsonl'), orient='records', lines=True)\n",
     "\n",
     "print(f\"✓ Extractive summary: {len(extractive_summary_generated_data)} records\")\n",
@@ -299,19 +330,19 @@
     "# Set model configuration\n",
     "flow = set_model_config(flow)\n",
     "\n",
-    "runtime_params = ({flow_name_map[flow_name]: {\n",
-    "        'n': number_of_summaries\n",
-    "    }})\n",
-    "\n",
     "if enable_reasoning:\n",
     "    # Increase max tokens to accommodate reasoning content\n",
     "    runtime_params = {\n",
     "        'question_generation': {'max_tokens': 1024}, \n",
-    "        flow_name_map[flow_name]: {'n': number_of_summaries, 'max_tokens': 6000}\n",
+    "        'gen_detailed_summary': {'n': number_of_summaries, 'max_tokens': 6000}\n",
     "        }\n",
+    "else:\n",
+    "    runtime_params = ({'gen_detailed_summary': {\n",
+    "        'n': number_of_summaries\n",
+    "    }})\n",
     "# Generate data for detailed summary\n",
     "detailed_summary_generated_data = flow.generate(quality_corpus, runtime_params=runtime_params, max_concurrency=50)\n",
-    "save_data_path = os.getenv('OUTPUT_DATA_FOLDER', '')\n",
+    "\n",
     "detailed_summary_generated_data.to_json(os.path.join(save_data_path, 'detailed_summary', 'gen.jsonl'), orient='records', lines=True)\n",
     "\n",
     "print(f\"✓ Detailed summary: {len(detailed_summary_generated_data)} records\")\n",
@@ -332,7 +363,7 @@
     "\n",
     "# Set model configuration\n",
     "flow = set_model_config(flow)\n",
-    "\n",
+    "runtime_params = {}\n",
     "if enable_reasoning:\n",
     "    # Increase max tokens for Question Generation to accommodate reasoning content\n",
     "    runtime_params = {\n",
@@ -340,9 +371,8 @@
     "        }\n",
     "\n",
     "# Generate data for key facts summary\n",
-    "key_facts_generated_data = flow.generate(quality_corpus, runtime_params=runtime_params, max_concurrency=50)\n",
+    "key_facts_generated_data = flow.generate(quality_corpus, runtime_params=runtime_params, max_concurrency=max_concurrency)\n",
     "\n",
-    "save_data_path = os.getenv('OUTPUT_DATA_FOLDER', '')\n",
     "key_facts_generated_data.to_json(os.path.join(save_data_path, 'key_facts_to_qa', 'gen.jsonl'), orient='records', lines=True)\n",
     "\n",
     "print(f\"✓ Key facts: {len(key_facts_generated_data)} records\")\n",
@@ -362,16 +392,15 @@
     "\n",
     "# Set model configuration\n",
     "flow = set_model_config(flow)\n",
-    "\n",
+    "runtime_params = {}\n",
     "if enable_reasoning:\n",
     "    # Increase max tokens to accommodate reasoning content\n",
     "    runtime_params = {\n",
-    "        'question_generation': {'max_tokens': 1024}, \n",
+    "        'question_generation': {'max_tokens': 2048}, \n",
     "        }\n",
     "\n",
-    "document_based_generated_data = flow.generate(quality_corpus, runtime_params=runtime_params, max_concurrency=200)\n",
+    "document_based_generated_data = flow.generate(quality_corpus, runtime_params=runtime_params, max_concurrency=max_concurrency)\n",
     "    \n",
-    "save_data_path = os.getenv('OUTPUT_DATA_FOLDER', '')\n",
     "document_based_generated_data.to_json(os.path.join(save_data_path, 'document_based_qa', 'gen.jsonl'), orient='records', lines=True)\n",
     "\n",
     "print(f\"✓ Document based: {len(document_based_generated_data)} records\")\n",
@@ -393,7 +422,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "test_nb",
+   "display_name": "sdg_hub",
    "language": "python",
    "name": "python3"
   },
@@ -407,7 +436,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.8"
+   "version": "3.11.12"
   }
  },
  "nbformat": 4,