Spaces:
Sleeping
Sleeping
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) | |
# Create a ZIP file containing the logs and augmented data | |
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") | |
# Download the ZIP file | |
with open(zip_file_path, "rb") as zip_file: | |
st.download_button( | |
label="πΎ Download Synthetic Data and Logs", | |
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() | |