import json import logging import time 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") 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): config_text = f""" # 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=YOUR_GRETEL_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 start_button: with st.expander("Augmentation Results", expanded=True): st.subheader("Augmentation Results") 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() logs = [] max_log_lines = 50 def custom_log_handler(msg): nonlocal logs logs.append(msg) if len(logs) > max_log_lines: logs = logs[-max_log_lines:] log_text = "\n".join(logs) log_container.text(log_text) handler = StreamlitLogHandler(custom_log_handler) logger = logging.getLogger("navigator_helpers") 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, ) augmented_data = [] 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() augmented_data.append(new_df) augmented_data_placeholder.subheader("Augmented Data") augmented_data_placeholder.dataframe( pd.concat(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!") if stop_button: st.warning("Augmentation stopped by the user.") st.stop() else: st.info( "Please upload a file or select a dataset from Hugging Face to proceed." ) if __name__ == "__main__": main()