fix logging bug, download logs and data

app.py

@@ -1,6 +1,9 @@
 import json
 import logging
+import os
+import tempfile
 import time
+import zipfile
 from io import StringIO
 
 import pandas as pd
@@ -8,8 +11,8 @@ import requests
 import streamlit as st
 from datasets import load_dataset
 from gretel_client import Gretel
-
-from navigator_helpers import DataAugmentationConfig, DataAugmenter, StreamlitLogHandler
+from navigator_helpers import (DataAugmentationConfig, DataAugmenter,
+                               StreamlitLogHandler)
 
 # Create a StringIO buffer to capture the logging output
 log_buffer = StringIO()
@@ -421,9 +424,19 @@ def main():
         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:
-            with st.expander("Augmentation Results", expanded=True):
-                st.subheader("Augmentation Results")
+            # 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:
@@ -433,20 +446,23 @@ def main():
                     )
                 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)
+                    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)
 
-                handler = StreamlitLogHandler(custom_log_handler)
+                # 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,
@@ -463,7 +479,6 @@ def main():
                     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):
@@ -476,10 +491,12 @@ def main():
                             verbose=True,
                         )
                         new_df = augmenter.augment()
-                        augmented_data.append(new_df)
-                        augmented_data_placeholder.subheader("Augmented Data")
+                        st.session_state.augmented_data.append(new_df)
+                        augmented_data_placeholder.subheader("Synthetic Data")
                         augmented_data_placeholder.dataframe(
-                            pd.concat(augmented_data, ignore_index=True)
+                            pd.concat(
+                                st.session_state.augmented_data, ignore_index=True
+                            )
                         )
                         progress = (index + 1) / num_records
                         progress_bar.progress(progress)
@@ -500,13 +517,71 @@ def main():
                         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, "augmented_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="augmentation_results.zip",
+                        mime="application/zip",
+                    )
+
             st.stop()
-        else:
-            st.info(
-                "Please upload a file or select a dataset from Hugging Face to proceed."
-            )
+
+    else:
+        st.info(
+            "Please upload a file or select a dataset from Hugging Face to proceed."
+        )
 
 
 if __name__ == "__main__":