Update app.py

app.py

@@ -1,13 +1,15 @@
-import gradio as gr
-import pandas as pd
-import numpy as np
+import os
 import io
 import base64
+import warnings
 from typing import Optional, Tuple
-import plotly.express as px
+
+import gradio as gr
+import pandas as pd
+import numpy as np
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
-import warnings
+
 warnings.filterwarnings("ignore")
 
 # Import Mostly AI SDK
@@ -18,96 +20,142 @@ except ImportError:
     MOSTLY_AI_AVAILABLE = False
     print("Warning: Mostly AI SDK not available. Please install with: pip install mostlyai[local]")
 
+
 class SyntheticDataGenerator:
     def __init__(self):
         self.mostly = None
         self.generator = None
         self.original_data = None
-        
-    def initialize_mostly_ai(self):
+
+    def initialize_mostly_ai(self) -> Tuple[bool, str]:
         """Initialize Mostly AI SDK"""
         if not MOSTLY_AI_AVAILABLE:
             return False, "Mostly AI SDK not installed. Please install with: pip install mostlyai[local]"
-        
         try:
             self.mostly = MostlyAI(local=True, local_port=8080)
             return True, "Mostly AI SDK initialized successfully."
         except Exception as e:
             return False, f"Failed to initialize Mostly AI SDK: {str(e)}"
-    
-    
-    def train_generator(self, data: pd.DataFrame, name: str, epochs: int = 10, max_training_time: int = 60, batch_size: int = 32, value_protection: bool = True) -> Tuple[bool, str]:
+
+    def train_generator(
+        self,
+        data: pd.DataFrame,
+        name: str,
+        epochs: int = 10,
+        max_training_time: int = 60,
+        batch_size: int = 32,
+        value_protection: bool = True,
+    ) -> Tuple[bool, str]:
         """Train the synthetic data generator"""
         if not self.mostly:
             return False, "Mostly AI SDK not initialized. Please initialize the SDK first."
-        
         try:
             self.original_data = data
-            train_config = {'tables':
-                            [
-                                {
-                                    'name': name,
-                                    'data': data,
-                                    'tabular_model_configuration':
-                                    {
-                                        'max_epochs': epochs,
-                                        'max_training_time': max_training_time,
-                                        'value_protection': value_protection,
-                                        'batch_size': batch_size
-                                    }
-                                }
-                            ]
-                           }
-                            
-            self.generator = self.mostly.train(
-                config = train_config
-            )
+            train_config = {
+                "tables": [
+                    {
+                        "name": name,
+                        "data": data,
+                        "tabular_model_configuration": {
+                            "max_epochs": epochs,
+                            "max_training_time": max_training_time,
+                            "value_protection": value_protection,
+                            "batch_size": batch_size,
+                        },
+                    }
+                ]
+            }
+
+            self.generator = self.mostly.train(config=train_config)
             return True, f"Training completed successfully. Model name: {name}"
         except Exception as e:
             return False, f"Training failed with error: {str(e)}"
-    
-    def generate_synthetic_data(self, size: int) -> Tuple[pd.DataFrame, str]:
+
+    def generate_synthetic_data(self, size: int) -> Tuple[Optional[pd.DataFrame], str]:
         """Generate synthetic data"""
         if not self.generator:
             return None, "No trained generator available. Please train a model first."
-        
         try:
             synthetic_data = self.mostly.generate(self.generator, size=size)
             df = synthetic_data.data()
             return df, f"Synthetic data generated successfully. {len(df)} records created."
         except Exception as e:
             return None, f"Synthetic data generation failed with error: {str(e)}"
-    
-    def get_quality_report(self) -> str:
-        """Get quality assurance report"""
+
+    # ---- Report helpers (new) ----
+    def get_quality_report_text(self) -> str:
+        """Return a concise status about the report."""
         if not self.generator:
             return "No trained generator available. Please train a model first."
-        
         try:
-            report = self.generator.reports(display=False)
-            return str(report)
+            _ = self.generator.reports(display=False)  # builds report internally
+            return "Quality report generated. Use the button to download."
         except Exception as e:
             return f"Failed to generate quality report: {str(e)}"
-    
+
+    def get_quality_report_file(self) -> Optional[str]:
+        """
+        Generate/export the report and return a file path for download.
+        Tries to find an existing ZIP; otherwise saves a TXT fallback.
+        """
+        if not self.generator:
+            return None
+        try:
+            rep = self.generator.reports(display=False)
+
+            # 1) If a string path to a .zip is returned
+            if isinstance(rep, str) and rep.endswith(".zip") and os.path.exists(rep):
+                return rep
+
+            # 2) If the object exposes a path-like attribute
+            for attr in ("archive_path", "zip_path", "path", "file_path"):
+                if hasattr(rep, attr):
+                    p = getattr(rep, attr)
+                    if isinstance(p, str) and os.path.exists(p):
+                        return p
+
+            # 3) If the object can save/export itself
+            target_zip = "/mnt/data/quality_report.zip"
+            if hasattr(rep, "save"):
+                try:
+                    rep.save(target_zip)
+                    if os.path.exists(target_zip):
+                        return target_zip
+                except Exception:
+                    pass
+            if hasattr(rep, "export"):
+                try:
+                    rep.export(target_zip)
+                    if os.path.exists(target_zip):
+                        return target_zip
+                except Exception:
+                    pass
+
+            # 4) Fallback: write string representation
+            target_txt = "/mnt/data/quality_report.txt"
+            with open(target_txt, "w", encoding="utf-8") as f:
+                f.write(str(rep))
+            return target_txt
+
+        except Exception:
+            return None
+
     def estimate_memory_usage(self, df: pd.DataFrame) -> str:
         """Estimate memory usage for the dataset"""
         if df is None or df.empty:
             return "No data available to analyze."
-        
-        # Calculate approximate memory usage
+
         memory_mb = df.memory_usage(deep=True).sum() / (1024 * 1024)
         rows, cols = len(df), len(df.columns)
-        
-        # Estimate training memory (roughly 3-5x the data size)
         estimated_training_mb = memory_mb * 4
-        
+
         if memory_mb < 100:
             status = "Good"
         elif memory_mb < 500:
             status = "Large"
         else:
             status = "Very Large"
-        
+
         return f"""
 Memory Usage Estimate:
 - Data size: {memory_mb:.1f} MB
@@ -116,196 +164,180 @@ Memory Usage Estimate:
 - Rows: {rows:,} | Columns: {cols}
         """.strip()
 
+
 # Initialize the generator
 generator = SyntheticDataGenerator()
 
+# ---- Wrapper functions for Gradio ----
+def initialize_sdk() -> str:
+    ok, msg = generator.initialize_mostly_ai()
+    return ("Success: " if ok else "Error: ") + msg
 
-def initialize_sdk() -> Tuple[str, str]:
-    """Initialize the Mostly AI SDK"""
-    success, message = generator.initialize_mostly_ai()
-    status = "Success" if success else "Error"
-    return status, message
 
-def train_model(data: pd.DataFrame, model_name: str, epochs: int, max_training_time: int, batch_size: int, value_protection: bool) -> Tuple[str, str]:
-    """Train the synthetic data generator"""
+def train_model(
+    data: pd.DataFrame,
+    model_name: str,
+    epochs: int,
+    max_training_time: int,
+    batch_size: int,
+    value_protection: bool,
+) -> str:
     if data is None or data.empty:
-        return "Error", "No data provided. Please upload or create sample data first."
-    
-    success, message = generator.train_generator(data, model_name, epochs, max_training_time, batch_size, value_protection)
-    status = "Success" if success else "Error"
-    return status, message
-
-def generate_data(size: int) -> Tuple[pd.DataFrame, str]:
-    """Generate synthetic data"""
-    if generator.generator is None:
-        return None, "Error: No trained model available. Please train a model first."
-    
+        return "Error: No data provided. Please upload or create sample data first."
+    ok, msg = generator.train_generator(
+        data, model_name, epochs, max_training_time, batch_size, value_protection
+    )
+    return ("Success: " if ok else "Error: ") + msg
+
+
+def generate_data(size: int) -> Tuple[Optional[pd.DataFrame], str]:
     synthetic_df, message = generator.generate_synthetic_data(size)
-    if synthetic_df is not None:
-        status = "Success"
-    else:
-        status = "Error"
-    
+    status = "Success" if synthetic_df is not None else "Error"
     return synthetic_df, f"{status}: {message}"
 
-def get_quality_report() -> str:
-    """Get quality report"""
-    return generator.get_quality_report()
 
-def create_comparison_plot(original_df: pd.DataFrame, synthetic_df: pd.DataFrame) -> go.Figure:
-    """Create comparison plots between original and synthetic data"""
+def get_quality_report_and_file():
+    """
+    Return (status_text, download_component_update)
+    The second value updates the DownloadButton with the file path and visibility.
+    """
+    status = generator.get_quality_report_text()
+    path = generator.get_quality_report_file()
+    if path:
+        return status, gr.update(value=path, visible=True)
+    else:
+        # keep it hidden if we don't have a file
+        return status, gr.update(visible=False)
+
+
+def create_comparison_plot(original_df: pd.DataFrame, synthetic_df: pd.DataFrame) -> Optional[go.Figure]:
     if original_df is None or synthetic_df is None:
         return None
-    
-    # Select numeric columns for comparison
+
     numeric_cols = original_df.select_dtypes(include=[np.number]).columns.tolist()
-    
     if not numeric_cols:
         return None
-    
-    # Create subplots
+
     n_cols = min(3, len(numeric_cols))
     n_rows = (len(numeric_cols) + n_cols - 1) // n_cols
-    
-    fig = make_subplots(
-        rows=n_rows, 
-        cols=n_cols,
-        subplot_titles=numeric_cols[:n_rows*n_cols]
-    )
-    
-    for i, col in enumerate(numeric_cols[:n_rows*n_cols]):
+
+    fig = make_subplots(rows=n_rows, cols=n_cols, subplot_titles=numeric_cols[: n_rows * n_cols])
+
+    for i, col in enumerate(numeric_cols[: n_rows * n_cols]):
         row = i // n_cols + 1
         col_idx = i % n_cols + 1
-        
-        # Add original data histogram
+
         fig.add_trace(
-            go.Histogram(
-                x=original_df[col],
-                name=f'Original {col}',
-                opacity=0.7,
-                nbinsx=20
-            ),
-            row=row, col=col_idx
+            go.Histogram(x=original_df[col], name=f"Original {col}", opacity=0.7, nbinsx=20),
+            row=row,
+            col=col_idx,
         )
-        
-        # Add synthetic data histogram
         fig.add_trace(
-            go.Histogram(
-                x=synthetic_df[col],
-                name=f'Synthetic {col}',
-                opacity=0.7,
-                nbinsx=20
-            ),
-            row=row, col=col_idx
+            go.Histogram(x=synthetic_df[col], name=f"Synthetic {col}", opacity=0.7, nbinsx=20),
+            row=row,
+            col=col_idx,
         )
-    
-    fig.update_layout(
-        title="Original vs Synthetic Data Comparison",
-        height=300 * n_rows,
-        showlegend=True
-    )
-    
+
+    fig.update_layout(title="Original vs Synthetic Data Comparison", height=300 * n_rows, showlegend=True)
     return fig
 
-def download_csv(df: pd.DataFrame) -> str:
-    """Convert DataFrame to CSV for download"""
+
+def download_csv(df: pd.DataFrame) -> Optional[str]:
     if df is None or df.empty:
         return None
-    
-    csv = df.to_csv(index=False)
-    return csv
+    # Write CSV to a stable path so DownloadButton can fetch it
+    path = "/mnt/data/synthetic_data.csv"
+    df.to_csv(path, index=False)
+    return path
+
 
-# Create the Gradio interface
+# ---- UI ----
 def create_interface():
     with gr.Blocks(title="MOSTLY AI Synthetic Data Generator", theme=gr.themes.Soft()) as demo:
-
-        # display image above tabs
-        gr.Image(value="https://img.mailinblue.com/8225865/images/content_library/original/6880d164e4e4ea1a183ad4c0.png", show_label=False, elem_id="header-image")
+        # Header image
+        gr.Image(
+            value="https://img.mailinblue.com/8225865/images/content_library/original/6880d164e4e4ea1a183ad4c0.png",
+            show_label=False,
+            elem_id="header-image",
+        )
 
         # README
-        gr.Markdown("""
+        gr.Markdown(
+            """
         # Synthetic Data SDK by MOSTLY AI Demo Space
 
         [Documentation](https://mostly-ai.github.io/mostlyai/) | [Technical White Paper](https://arxiv.org/abs/2508.00718) | [Usage Examples](https://mostly-ai.github.io/mostlyai/usage/) | [Free Cloud Service](https://app.mostly.ai/)
         
         A Python toolkit for generating high-fidelity, privacy-safe synthetic data. 
-        """)
-        
+        """
+        )
+
         with gr.Tab("Quick Start"):
             gr.Markdown("### Initialize the SDK and upload your data")
-            
             with gr.Row():
                 with gr.Column():
                     init_btn = gr.Button("Initialize Mostly AI SDK", variant="primary")
                     init_status = gr.Textbox(label="Initialization Status", interactive=False)
-                
                 with gr.Column():
-                    gr.Markdown("""
+                    gr.Markdown(
+                        """
                     **Next Steps:**
                     1. Initialize the SDK (click button above)
                     2. Go to "Upload Data and Train Model" tab to upload your CSV file
                     3. Train a model on your data
                     4. Generate synthetic data
-                    """)
-        
+                    """
+                    )
+
         with gr.Tab("Upload Data and Train Model"):
             gr.Markdown("### Upload your CSV file to generate synthetic data")
-            
-            gr.Markdown("""
+            gr.Markdown(
+                """
             **File Requirements:**
             - Format: CSV with header row
             - Size: Optimized for Hugging Face Spaces (2 vCPU, 16GB RAM)
-            """)
-            
-            file_upload = gr.File(
-                label="Upload CSV File",
-                file_types=[".csv"],
-                file_count="single"
+            """
             )
-            
+
+            file_upload = gr.File(label="Upload CSV File", file_types=[".csv"], file_count="single")
             uploaded_data = gr.Dataframe(label="Uploaded Data", interactive=False)
-            
             memory_info = gr.Markdown(label="Memory Usage Info", visible=False)
-            
+
             with gr.Row():
                 with gr.Column():
                     model_name = gr.Textbox(
-                        value="My Synthetic Model",
-                        label="Model Name",
-                        placeholder="Enter a name for your model"
+                        value="My Synthetic Model", label="Model Name", placeholder="Enter a name for your model"
                     )
                     epochs = gr.Slider(1, 200, value=100, step=1, label="Training Epochs")
                     max_training_time = gr.Slider(1, 1000, value=60, step=1, label="Maximum Training Time")
                     batch_size = gr.Slider(8, 1024, value=32, step=8, label="Training Batch Size")
                     value_protection = gr.Checkbox(label="Value Protection", info="Enable Value Protection")
                     train_btn = gr.Button("Train Model", variant="primary")
-                
                 with gr.Column():
                     train_status = gr.Textbox(label="Training Status", interactive=False)
-                    quality_report = gr.Textbox(label="Quality Report", lines=10, interactive=False)
-            
-            get_report_btn = gr.Button("Get Quality Report", variant="secondary")
-        
+                    quality_report = gr.Textbox(label="Quality Report", lines=8, interactive=False)
+
+            with gr.Row():
+                get_report_btn = gr.Button("Get Quality Report", variant="secondary")
+                report_download_btn = gr.DownloadButton("Download Quality Report", visible=False)
+
         with gr.Tab("Generate Data"):
             gr.Markdown("### Generate synthetic data from your trained model")
-            
             with gr.Row():
                 with gr.Column():
                     gen_size = gr.Slider(10, 1000, value=100, step=10, label="Number of Records to Generate")
                     generate_btn = gr.Button("Generate Synthetic Data", variant="primary")
-                
                 with gr.Column():
                     gen_status = gr.Textbox(label="Generation Status", interactive=False)
-            
+
             synthetic_data = gr.Dataframe(label="Synthetic Data", interactive=False)
-            
             with gr.Row():
-                download_btn = gr.DownloadButton("Download CSV", variant="secondary")
+                download_btn = gr.DownloadButton("Download CSV", file_name="synthetic_data.csv", variant="secondary")
                 comparison_plot = gr.Plot(label="Data Comparison")
 
-        # README
-        gr.Markdown("""
-            
+        # README footer
+        gr.Markdown(
+            """
         **Modes of operation:**
         - **LOCAL mode** trains and generates synthetic data on your own compute resources.
         - **CLIENT mode** connects to a remote MOSTLY AI platform for training and generation.
@@ -325,75 +357,52 @@ def create_interface():
         The open source Synthetic Data SDK by MOSTLY AI powers the MOSTLY AI Platform and MOSTLY AI Assistant.
         
         Sign up for free and try the [MOSTLY AI Platform](https://app.mostly.ai/) today!
-        """)
-        
-        # Event handlers
-        init_btn.click(
-            initialize_sdk,
-            outputs=[init_status, init_status]
+        """
         )
-        
+
+        # ---- Event handlers ----
+        init_btn.click(initialize_sdk, outputs=[init_status])
+
         train_btn.click(
             train_model,
             inputs=[uploaded_data, model_name, epochs, max_training_time, batch_size, value_protection],
-            outputs=[train_status, train_status]
+            outputs=[train_status],
         )
-        
+
+        # Build + expose quality report for download
         get_report_btn.click(
-            get_quality_report,
-            outputs=[quality_report]
-        )
-        
-        generate_btn.click(
-            generate_data,
-            inputs=[gen_size],
-            outputs=[synthetic_data, gen_status]
-        )
-        
-        # Update download button when synthetic data changes
-        synthetic_data.change(
-            download_csv,
-            inputs=[synthetic_data],
-            outputs=[download_btn]
+            get_quality_report_and_file,
+            outputs=[quality_report, report_download_btn],
         )
-        
-        # Create comparison plot when both datasets are available
+
+        # Generate data
+        generate_btn.click(generate_data, inputs=[gen_size], outputs=[synthetic_data, gen_status])
+
+        # Update CSV DownloadButton whenever synthetic data changes
+        synthetic_data.change(download_csv, inputs=[synthetic_data], outputs=[download_btn])
+
+        # Build comparison plot when both datasets are available
         synthetic_data.change(
-            create_comparison_plot,
-            inputs=[uploaded_data, synthetic_data],
-            outputs=[comparison_plot]
+            create_comparison_plot, inputs=[uploaded_data, synthetic_data], outputs=[comparison_plot]
         )
-        
+
         # Handle file upload with size and column limits
         def process_uploaded_file(file):
             if file is None:
                 return None, "No file uploaded.", gr.update(visible=False)
-            
             try:
-                # Read the CSV file
                 df = pd.read_csv(file.name)
-                
                 success_msg = f"File uploaded successfully. {len(df)} rows × {len(df.columns)} columns"
-                
-                memory_info = generator.estimate_memory_usage(df)
-                
-                return df, success_msg, gr.update(value=memory_info, visible=True)
-                
+                mem_info = generator.estimate_memory_usage(df)
+                return df, success_msg, gr.update(value=mem_info, visible=True)
             except Exception as e:
                 return None, f"Error reading file: {str(e)}", gr.update(visible=False)
-        
-        file_upload.change(
-            process_uploaded_file,
-            inputs=[file_upload],
-            outputs=[uploaded_data, train_status, memory_info]
-        )
-    
+
+        file_upload.change(process_uploaded_file, inputs=[file_upload], outputs=[uploaded_data, train_status, memory_info])
+
     return demo
 
+
 if __name__ == "__main__":
     demo = create_interface()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=True
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)