app.py

import json
import logging
import os
import tempfile
import time
import zipfile
from io import StringIO

import pandas as pd
import requests
import streamlit as st
from datasets import load_dataset
from gretel_client import Gretel
from navigator_helpers import DataAugmentationConfig, DataAugmenter, StreamlitLogHandler

# Create a StringIO buffer to capture the logging output
log_buffer = StringIO()

# Create a handler to redirect logging output to the buffer
handler = logging.StreamHandler(log_buffer)
handler.setLevel(logging.INFO)

# Set up the logger
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.addHandler(handler)


SAMPLE_DATASET_URL = "https://gretel-public-website.s3.us-west-2.amazonaws.com/datasets/llm-training-data/dolly-examples-qa-with-context.csv"
WELCOME_MARKDOWN = """
Gretel Navigator is an interface designed to help you create high-quality, diverse training data examples through synthetic data generation techniques. It aims to assist in scenarios where you have limited training data or want to enhance the quality and diversity of your existing dataset.

## 🎯 Key Use Cases

1. **Augment Existing Training Data**: Expand your existing training data with additional synthetic examples generated by Gretel Navigator. This can help improve the robustness and generalization of your AI models.

2. **Create Diverse Training or Evaluation Data**: Generate diverse training or evaluation data from plain text or seed examples. This ensures your AI models are exposed to a wide range of scenarios and edge cases during training.

3. **Address Data Limitations**: Generate additional examples to fill gaps in your dataset, particularly for underrepresented classes, rare events, or challenging scenarios. This helps improve your model's ability to handle diverse real-world situations.

4. **Mitigate Bias and Toxicity**: Generate training examples that are unbiased and non-toxic by incorporating diverse perspectives and adhering to ethical guidelines. This promotes fairness and responsible AI development.

5. **Enhance Model Performance**: Improve the performance of your AI models across various tasks by training them on diverse synthetic data generated by Gretel Navigator.

## 🔧 Getting Started

To start using Gretel Navigator, you'll need:

1. A Gretel account (free accounts are available).
2. Seed text or input/output pairs to create or augment AI training data.

## 📂 Input Data Formats

Gretel Navigator supports the following formats for input data:

- Existing AI training or evaluation data formats:
  - Input/Output pair format (or instruction/response) with any number of ground truth or "context fields".
  - Plain text data.
- File formats:
  - Hugging Face dataset
  - CSV
  - JSON
  - JSONL

## 📤 Output

Gretel Navigator generates one additional training example per row in the input/output pair format. You can specify requirements for the input and output pairs in the configuration. Run the process multiple times to scale your data to any desired level.

## 🌟 AI Alignment Techniques

Gretel Navigator incorporates AI alignment techniques to generate high-quality synthetic data:

- Diverse Instruction and Response Generation
- AI-Aligning-AI Methodology (AAA) for iterative data quality enhancement
- Quality Evaluation
- Bias and Toxicity Detection

By leveraging these techniques, Gretel Navigator helps you create training data that leads to more robust, unbiased, and high-performing AI models.

---

Ready to enhance your AI training data and unlock the full potential of your models? Let's get started with Gretel Navigator! 🚀
"""


def main():
    st.set_page_config(page_title="Gretel", layout="wide")
    st.title("🎨 Gretel Navigator: Enhance Your AI Training Data")
    st.write(
        "Generate diverse synthetic training data from text or existing datasets to improve the performance and robustness of your AI models."
    )

    with st.expander("Introduction", expanded=False):
        st.markdown(WELCOME_MARKDOWN)

    st.subheader("Step 1: API Key Validation")
    with st.expander("API Key Configuration", expanded=True):
        api_key = st.text_input(
            "Enter your Gretel API key (Get a free API key at: https://console.gretel.ai/users/me/key)",
            value="",
            type="password",
            help="Your Gretel API key is required to authenticate and use Gretel Navigator. If you don't have one yet, sign up for a free account at https://console.gretel.ai to get started.",
        )
        if "gretel" not in st.session_state:
            st.session_state.gretel = None
        if st.button("Validate API Key"):
            if api_key:
                try:
                    st.session_state.gretel = Gretel(api_key=api_key, validate=True)
                    st.success("API key validated. Connection successful!")
                except Exception as e:
                    st.error(f"Error connecting to Gretel: {str(e)}")
            else:
                st.warning("Please enter your Gretel API key to proceed.")
        if st.session_state.gretel is None:
            st.stop()

    st.subheader("Step 2: Data Source Selection")
    with st.expander("Data Source", expanded=True):
        data_source = st.radio(
            "Select data source",
            options=[
                "Upload a file",
                "Select a dataset from Hugging Face",
                "Use a sample dataset",
            ],
            help="Choose whether to upload a file, select a dataset from Hugging Face, or use a sample dataset",
        )

        df = None
        if data_source == "Upload a file":
            uploaded_file = st.file_uploader(
                "Upload a CSV, JSON, or JSONL file",
                type=["csv", "json", "jsonl"],
                help="Upload the dataset file in CSV, JSON, or JSONL format",
            )

            if uploaded_file is not None:
                if uploaded_file.name.endswith(".csv"):
                    df = pd.read_csv(uploaded_file)
                elif uploaded_file.name.endswith(".json"):
                    df = pd.read_json(uploaded_file)
                elif uploaded_file.name.endswith(".jsonl"):
                    df = pd.read_json(uploaded_file, lines=True)
                st.success(f"File uploaded successfully: {uploaded_file.name}")

        elif data_source == "Select a dataset from Hugging Face":
            huggingface_dataset = st.text_input(
                "Hugging Face Dataset Repository",
                help="Enter the name of the Hugging Face dataset repository (e.g., 'squad')",
            )

            huggingface_split = st.selectbox(
                "Dataset Split",
                options=["train", "validation", "test"],
                help="Select the dataset split to use",
            )

            if st.button("Load Hugging Face Dataset"):
                if huggingface_dataset:
                    try:
                        with st.spinner("Loading dataset from Hugging Face..."):
                            dataset = load_dataset(
                                huggingface_dataset, split=huggingface_split
                            )
                            df = dataset.to_pandas()
                        st.success(
                            f"Dataset loaded from Hugging Face repository: {huggingface_dataset}"
                        )
                    except Exception as e:
                        st.error(f"Error loading dataset from Hugging Face: {str(e)}")
                else:
                    st.warning("Please provide a Hugging Face dataset repository name.")

        elif data_source == "Use a sample dataset":
            st.write("Try a sample dataset to get started quickly.")
            if st.button("Try Sample Dataset"):
                try:
                    df = pd.read_csv(SAMPLE_DATASET_URL)
                    st.success("Sample dataset loaded successfully.")
                except Exception as e:
                    st.error(f"Error downloading sample dataset: {str(e)}")

        if df is not None:
            st.session_state.df = df
            st.session_state.selected_fields = list(df.columns)
            st.write(
                f"Loaded dataset with {len(df)} rows and {len(df.columns)} columns."
            )
        else:
            df = st.session_state.get("df")

    st.subheader("Step 3: Data Preview and Configuration")
    if df is not None:
        with st.expander("Data Preview", expanded=True):
            st.dataframe(df.head())

        with st.expander("Input Fields Selection", expanded=True):
            st.write(
                "Select the context fields to provide the LLM access to for generating input/output pairs. This can include existing instructions and responses. All selected fields will be treated as ground truth data."
            )

            selected_fields = []
            for column in df.columns:
                if st.checkbox(
                    column,
                    value=column in st.session_state.get("selected_fields", []),
                    key=f"checkbox_{column}",
                ):
                    selected_fields.append(column)

            st.session_state.selected_fields = selected_fields

        with st.expander("Advanced Options", expanded=False):

            output_instruction_field = st.text_input(
                "Synthetic instruction field",
                value=st.session_state.get("output_instruction_field", "instruction"),
                help="Specify the name of the output field for generated instructions",
            )
            st.session_state.output_instruction_field = output_instruction_field

            output_response_field = st.text_input(
                "Synthetic response field",
                value=st.session_state.get("output_response_field", "response"),
                help="Specify the name of the output field for generated responses",
            )
            st.session_state.output_response_field = output_response_field

            num_records = st.number_input(
                "Max number of records from input data to process",
                min_value=1,
                max_value=len(df),
                value=len(df),
                help="Specify the number of records to process",
            )
            st.session_state.num_records = num_records

            num_instructions = st.number_input(
                "Number of diverse candidate instructions",
                min_value=1,
                value=st.session_state.get("num_instructions", 5),
                help="Specify the number of instructions to generate",
            )
            st.session_state.num_instructions = num_instructions

            num_responses = st.number_input(
                "Number of diverse candidateresponses",
                min_value=1,
                value=st.session_state.get("num_responses", 5),
                help="Specify the number of responses to generate",
            )
            st.session_state.num_responses = num_responses

            temperature = st.slider(
                "Temperature",
                min_value=0.0,
                max_value=1.0,
                value=st.session_state.get("temperature", 0.8),
                step=0.1,
                help="Adjust the temperature for response generation",
            )
            st.session_state.temperature = temperature

            max_tokens_instruction = st.slider(
                "Max tokens (instruction)",
                min_value=1,
                max_value=1024,
                value=st.session_state.get("max_tokens_instruction", 100),
                help="Specify the maximum number of tokens for instructions",
            )
            st.session_state.max_tokens_instruction = max_tokens_instruction

            max_tokens_response = st.slider(
                "Max tokens (response)",
                min_value=1,
                max_value=1024,
                value=st.session_state.get("max_tokens_response", 100),
                help="Specify the maximum number of tokens for responses",
            )
            st.session_state.max_tokens_response = max_tokens_response

        with st.expander("Model Configuration", expanded=True):
            st.markdown("### Primary Navigator Models")

            tabular_models = st.session_state.gretel.factories.get_navigator_model_list(
                "tabular"
            )
            navigator_tabular = st.selectbox(
                "Navigator Tabular",
                options=tabular_models,
                index=st.session_state.get("navigator_tabular_index", 0),
                help="Select the primary Navigator tabular model",
            )
            st.session_state.navigator_tabular_index = tabular_models.index(
                navigator_tabular
            )

            nl_models = st.session_state.gretel.factories.get_navigator_model_list(
                "natural_language"
            )
            navigator_llm = st.selectbox(
                "Navigator LLM",
                options=nl_models,
                index=st.session_state.get("navigator_llm_index", 0),
                help="Select the primary Navigator LLM",
            )
            st.session_state.navigator_llm_index = nl_models.index(navigator_llm)

            st.markdown("---")
            st.markdown("### AI Align AI (AAA)")
            st.write(
                "AI Align AI (AAA) is a technique that iteratively improves the quality and coherence of generated outputs by using multiple LLMs for co-teaching and self-teaching. Enabling AAA will enhance the overall quality of the synthetic data, but it may slow down the generation process."
            )

            use_aaa = st.checkbox(
                "Use AI Align AI (AAA)",
                value=st.session_state.get("use_aaa", True),
                help="Enable or disable the use of AI Align AI.",
            )
            st.session_state.use_aaa = use_aaa

            co_teach_llms = []  # Initialize co_teach_llms with an empty list

            if use_aaa:
                st.markdown("#### Navigator Co-teaching LLMs")
                st.write(
                    "Select additional Navigator LLMs for co-teaching in AAA. It is recommended to use different LLMs than the primary Navigator LLM for this step."
                )

                for model in nl_models:
                    if model != navigator_llm:
                        if st.checkbox(model, value=True, key=f"checkbox_{model}"):
                            co_teach_llms.append(model)
                    else:
                        if st.checkbox(model, value=False, key=f"checkbox_{model}"):
                            co_teach_llms.append(model)
                st.session_state.co_teach_llms = co_teach_llms

            st.markdown("---")
            st.markdown("### Format Prompts")

            instruction_format_prompt = st.text_area(
                "Instruction Format Prompt",
                value=st.session_state.get(
                    "instruction_format_prompt",
                    "A well-formulated question or command in everyday English.",
                ),
                help="Specify the format prompt for instructions",
            )
            st.session_state.instruction_format_prompt = instruction_format_prompt

            response_format_prompt = st.text_area(
                "Response Format Prompt",
                value=st.session_state.get(
                    "response_format_prompt",
                    "A well-formulated response to the question in everyday English.",
                ),
                help="Specify the format prompt for responses",
            )
            st.session_state.response_format_prompt = response_format_prompt

        with st.expander("Download SDK Code", expanded=False):
            st.markdown("### Ready to generate data at scale?")
            st.write(
                "Get started with your current configuration using the SDK code below:"
            )

            config_text = f"""
        #!pip install -Uqq git+https://github.com/gretelai/navigator-helpers.git

        import logging
        import sys
        import pandas as pd

        from navigator_helpers import DataAugmentationConfig, DataAugmenter

        # Configure the logger
        logger = logging.getLogger()
        logger.setLevel(logging.INFO)

        DATASET = "YOUR_DATASET"
        API_KEY = "YOUR_API_KEY"

        df = pd.read_csv(DATASET)

        # Create the data augmentation configuration
        config = DataAugmentationConfig(
            input_fields={st.session_state.selected_fields},
            output_instruction_field="{output_instruction_field}",
            output_response_field="{output_response_field}",
            num_instructions={num_instructions},
            num_responses={num_responses},
            temperature={temperature},
            max_tokens_instruction={max_tokens_instruction},
            max_tokens_response={max_tokens_response},
            api_key=API_KEY,
            navigator_tabular="{navigator_tabular}",
            navigator_llm="{navigator_llm}",
            co_teach_llms={co_teach_llms},
            instruction_format_prompt='''{instruction_format_prompt}''',
            response_format_prompt='''{response_format_prompt}'''
        )

        # Create the data augmenter and perform augmentation
        augmenter = DataAugmenter(
            df,
            config,
            use_aaa={use_aaa},
            output_file="results.csv",
            verbose=True,
        )
        new_df = augmenter.augment()
        """
            st.code(config_text, language="python")
            st.download_button(
                label="Download SDK Code",
                data=config_text,
                file_name="data_augmentation_code.py",
                mime="text/plain",
            )

        start_stop_container = st.empty()

        col1, col2 = st.columns(2)
        with col1:
            start_button = st.button("🚀 Start")
        with col2:
            stop_button = st.button("🛑 Stop")

        if "logs" not in st.session_state:
            st.session_state.logs = []

        if "augmented_data" not in st.session_state:
            st.session_state.augmented_data = []

        if start_button:
            # Clear the augmented data and logs before starting a new generation
            st.session_state.augmented_data = []
            st.session_state.logs = []

            with st.expander("Synthetic Data", expanded=True):
                st.subheader("Synthetic Data Generation")
                progress_bar = st.progress(0)
                tab1, tab2 = st.tabs(["Augmented Data", "Logs"])
                with tab1:
                    augmented_data_placeholder = st.empty()
                    st.info(
                        "Click on the 'Logs' tab to see and debug real-time logging for each record as it is generated by the agents."
                    )
                with tab2:
                    log_container = st.empty()
                    max_log_lines = 50

                def custom_log_handler(msg):
                    st.session_state.logs.append(msg)
                    displayed_logs = st.session_state.logs[-max_log_lines:]
                    log_text = "\n".join(displayed_logs)
                    log_container.text(log_text)

                # Remove the previous log handler if it exists
                logger = logging.getLogger("navigator_helpers")
                for handler in logger.handlers:
                    if isinstance(handler, StreamlitLogHandler):
                        logger.removeHandler(handler)

                handler = StreamlitLogHandler(custom_log_handler)
                logger.addHandler(handler)

                config = DataAugmentationConfig(
                    input_fields=selected_fields,
                    output_instruction_field=output_instruction_field,
                    output_response_field=output_response_field,
                    num_instructions=num_instructions,
                    num_responses=num_responses,
                    temperature=temperature,
                    max_tokens_instruction=max_tokens_instruction,
                    max_tokens_response=max_tokens_response,
                    api_key=api_key,
                    navigator_tabular=navigator_tabular,
                    navigator_llm=navigator_llm,
                    co_teach_llms=co_teach_llms,
                    instruction_format_prompt=instruction_format_prompt,
                    response_format_prompt=response_format_prompt,
                )
                start_time = time.time()
                with st.spinner("Generating synthetic data..."):
                    for index in range(num_records):
                        row = df.iloc[index]
                        augmenter = DataAugmenter(
                            pd.DataFrame([row]),
                            config,
                            use_aaa=use_aaa,
                            output_file="results.csv",
                            verbose=True,
                        )
                        new_df = augmenter.augment()
                        st.session_state.augmented_data.append(new_df)
                        augmented_data_placeholder.subheader("Synthetic Data")
                        augmented_data_placeholder.dataframe(
                            pd.concat(
                                st.session_state.augmented_data, ignore_index=True
                            )
                        )
                        progress = (index + 1) / num_records
                        progress_bar.progress(progress)

                        elapsed_time = time.time() - start_time
                        records_processed = index + 1
                        records_remaining = num_records - records_processed
                        est_time_per_record = (
                            elapsed_time / records_processed
                            if records_processed > 0
                            else 0
                        )
                        est_time_remaining = est_time_per_record * records_remaining

                        progress_text = f"Progress: {progress:.2%} | Records Processed: {records_processed} | Records Remaining: {records_remaining} | Est. Time per Record: {est_time_per_record:.2f}s | Est. Time Remaining: {est_time_remaining:.2f}s"
                        progress_bar.text(progress_text)

                        time.sleep(0.1)
                logger.removeHandler(handler)
                st.success("Data augmentation completed!")
            st.stop()

        if stop_button:
            st.warning("Augmentation stopped by the user.")

            # Get the complete logs from the session state
            complete_logs = st.session_state.logs

            # Convert complete logs to JSONL format
            log_jsonl = "\n".join([json.dumps({"log": log}) for log in complete_logs])

            # Convert augmented data to JSONL format if it exists
            if st.session_state.augmented_data:
                augmented_df = pd.concat(
                    st.session_state.augmented_data, ignore_index=True
                )
                if not augmented_df.empty:
                    augmented_data_jsonl = "\n".join(
                        [
                            json.dumps(row.to_dict())
                            for _, row in augmented_df.iterrows()
                        ]
                    )
                else:
                    augmented_data_jsonl = None
            else:
                augmented_data_jsonl = None

            # Create a temporary directory to store the files
            with tempfile.TemporaryDirectory() as temp_dir:
                # Write the complete logs to a file
                log_file_path = os.path.join(temp_dir, "complete_logs.jsonl")
                with open(log_file_path, "w") as log_file:
                    log_file.write(log_jsonl)

                # Write the augmented data to a file if it exists
                if augmented_data_jsonl:
                    augmented_data_file_path = os.path.join(
                        temp_dir, "synthetic_data.jsonl"
                    )
                    with open(augmented_data_file_path, "w") as augmented_data_file:
                        augmented_data_file.write(augmented_data_jsonl)

                # Write the SDK code to a file
                sdk_file_path = os.path.join(temp_dir, "data_augmentation_code.py")
                with open(sdk_file_path, "w") as sdk_file:
                    sdk_file.write(config_text)

                # Create a ZIP file containing the logs, augmented data, and SDK code
                zip_file_path = os.path.join(temp_dir, "augmentation_results.zip")
                with zipfile.ZipFile(zip_file_path, "w") as zip_file:
                    zip_file.write(log_file_path, "complete_logs.jsonl")
                    if augmented_data_jsonl:
                        zip_file.write(augmented_data_file_path, "augmented_data.jsonl")
                    zip_file.write(sdk_file_path, "data_augmentation_code.py")

                # Download the ZIP file
                with open(zip_file_path, "rb") as zip_file:
                    st.download_button(
                        label="💾 Download Synthetic Data, Logs, and SDK Code",
                        data=zip_file.read(),
                        file_name="gretel_synthetic_data.zip",
                        mime="application/zip",
                    )

            st.stop()

    else:
        st.info(
            "Please upload a file or select a dataset from Hugging Face to proceed."
        )


if __name__ == "__main__":
    main()