Update app.py

app.py

@@ -1,204 +1,164 @@
-import gradio as gr
-from smolagents import HfApiModel, CodeAgent
-from huggingface_hub import login
 import os
-import shutil
-import wandb
-import time
-import psutil
+import gradio as gr
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+import shap
+import lime.lime_tabular
 import optuna
+import wandb
 import ast
-import pandas as pd
-from sklearn.model_selection import train_test_split
+from smolagents import HfApiModel, CodeAgent
+from huggingface_hub import login
 from sklearn.ensemble import RandomForestClassifier
+from sklearn.model_selection import train_test_split, cross_val_score
 from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
-import shap
-import lime
-import lime.lime_tabular
-import matplotlib.pyplot as plt
 
-# Authenticate Hugging Face
+# Authenticate Hugging Face Hub
 hf_token = os.getenv("HF_TOKEN")
 login(token=hf_token)
 
-# Initialize Model
+# Setup SmolAgent with LLM
 model = HfApiModel("mistralai/Mixtral-8x7B-Instruct-v0.1", token=hf_token)
-
-def format_analysis_report(raw_output, visuals):
+agent = CodeAgent(
+    tools=[],
+    model=model,
+    additional_authorized_imports=["numpy", "pandas", "matplotlib.pyplot", "seaborn", "sklearn"],
+    max_iterations=10,
+)
+
+# Data cleaning function
+def clean_data(df):
+    df = df.dropna(how='all', axis=1).dropna(how='all', axis=0)
+    df = df.fillna(df.mean(numeric_only=True))
+    df = df.select_dtypes(include=[np.number])
+    return df
+
+# Global dataframe
+df_global = None
+
+# Upload and clean
+def upload_file(file):
+    global df_global
+    ext = os.path.splitext(file.name)[-1]
+    df = pd.read_csv(file.name) if ext == ".csv" else pd.read_excel(file.name)
+    df = clean_data(df)
+    df_global = df
+    return df.head()
+
+# Run SmolAgent for analysis
+def run_agent(_):
     try:
-        if isinstance(raw_output, dict):
-            analysis_dict = raw_output
-        else:
-            try:
-                analysis_dict = ast.literal_eval(str(raw_output))
-            except (SyntaxError, ValueError) as e:
-                print(f"Error parsing CodeAgent output: {e}")
-                return str(raw_output), visuals
-                
-        report = f"""
-        <div style="font-family: Arial, sans-serif; padding: 20px; color: #333;">
-            <h1 style="color: #2B547E; border-bottom: 2px solid #2B547E; padding-bottom: 10px;">📊 Data Analysis Report</h1>
-            <div style="margin-top: 25px; background: #f8f9fa; padding: 20px; border-radius: 8px;">
-                <h2 style="color: #2B547E;">🔍 Key Observations</h2>
-                {format_observations(analysis_dict.get('observations', {}))}
-            </div>
-            <div style="margin-top: 30px;">
-                <h2 style="color: #2B547E;">💡 Insights & Visualizations</h2>
-                {format_insights(analysis_dict.get('insights', {}), visuals)}
-            </div>
-        </div>
-        """
-        return report, visuals
+        output = agent.run(
+            df_global,
+            instructions="Generate 5 data insights and 5 data visualizations. Visualizations should be saved in current working directory."
+        )
+        return str(output)
     except Exception as e:
-        print(f"Error in format_analysis_report: {e}")
-        return str(raw_output), visuals
-
-def format_observations(observations):
-    return '\n'.join([
-        f"""
-        <div style="margin: 15px 0; padding: 15px; background: white; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.05);">
-            <h3 style="margin: 0 0 10px 0; color: #4A708B;">{key.replace('_', ' ').title()}</h3>
-            <pre style="margin: 0; padding: 10px; background: #f8f9fa; border-radius: 4px;">{value}</pre>
-        </div>
-        """ for key, value in observations.items() if 'proportions' in key
-    ])
-
-def format_insights(insights, visuals):
-    return '\n'.join([
-        f"""
-        <div style="margin: 20px 0; padding: 20px; background: white; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.05);">
-            <div style="display: flex; align-items: center; gap: 10px;">
-                <div style="background: #2B547E; color: white; width: 30px; height: 30px; border-radius: 50%; display: flex; align-items: center; justify-content: center;">{idx+1}</div>
-                <p style="margin: 0; font-size: 16px;">{insight}</p>
-            </div>
-            {f'<img src="/file={visuals[idx]}" style="max-width: 100%; height: auto; margin-top: 10px; border-radius: 6px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">' if idx < len(visuals) else ''}
-        </div>
-        """ for idx, (key, insight) in enumerate(insights.items())
-    ])
-
-def analyze_data(csv_file, additional_notes=""):
-    start_time = time.time()
-    process = psutil.Process(os.getpid())
-    initial_memory = process.memory_info().rss / 1024 ** 2
-
-    if os.path.exists('./figures'):
-        shutil.rmtree('./figures')
-    os.makedirs('./figures', exist_ok=True)
-
-    wandb.login(key=os.environ.get('WANDB_API_KEY'))
-    run = wandb.init(project="huggingface-data-analysis", config={
-        "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
-        "additional_notes": additional_notes,
-        "source_file": csv_file.name if csv_file else None
-    })
-
-    agent = CodeAgent(tools=[], model=model, additional_authorized_imports=["numpy", "pandas", "matplotlib.pyplot", "seaborn", "sklearn"])
-    analysis_result = agent.run("""
-        You are an expert data analyst. Perform comprehensive analysis including:
-        1. Basic statistics and data quality checks
-        2. 3 insightful analytical questions about relationships in the data
-        3. Visualization of key patterns and correlations
-        4. Actionable real-world insights derived from findings.
-        Generate publication-quality visualizations and save to './figures/'.
-        Return the analysis results as a python dictionary that can be parsed by ast.literal_eval().
-        The dictionary should have the following structure:
-        {
-            'observations': {
-                'observation_1_key': 'observation_1_value',
-                ...
-            },
-            'insights': {
-                'insight_1_key': 'insight_1_value',
-                ...
-            }
+        return f"SmolAgent Error: {str(e)}"
+
+# Train model + Optuna + WandB
+def train_model(_):
+    wandb.login(key=os.environ.get("WANDB_API_KEY"))
+    wandb_run = wandb.init(project="huggingface-data-analysis", name="Optuna_Run", reinit=True)
+
+    target = df_global.columns[-1]
+    X = df_global.drop(target, axis=1)
+    y = df_global[target]
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
+
+    def objective(trial):
+        params = {
+            "n_estimators": trial.suggest_int("n_estimators", 50, 200),
+            "max_depth": trial.suggest_int("max_depth", 3, 10),
         }
-    """, additional_args={"additional_notes": additional_notes, "source_file": csv_file})
-
-    execution_time = time.time() - start_time
-    final_memory = process.memory_info().rss / 1024 ** 2
-    memory_usage = final_memory - initial_memory
-    wandb.log({"execution_time_sec": execution_time, "memory_usage_mb": memory_usage})
-
-    visuals = [os.path.join('./figures', f) for f in os.listdir('./figures') if f.endswith(('.png', '.jpg', '.jpeg'))]
-    for viz in visuals:
-        wandb.log({os.path.basename(viz): wandb.Image(viz)})
-
-    run.finish()
-    return format_analysis_report(analysis_result, visuals)
-
-def objective(trial, X_train, y_train, X_test, y_test):
-    n_estimators = trial.suggest_int("n_estimators", 50, 200)
-    max_depth = trial.suggest_int("max_depth", 3, 10)
-    
-    model = RandomForestClassifier(n_estimators=n_estimators, max_depth=max_depth, random_state=42)
-    model.fit(X_train, y_train)
-    predictions = model.predict(X_test)
-
-    return accuracy_score(y_test, predictions)
-
-def tune_hyperparameters(csv_file, n_trials: int):
-    df = pd.read_csv(csv_file)
-    y = df.iloc[:, -1]
-    X = df.iloc[:, :-1]
-    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+        model = RandomForestClassifier(**params)
+        score = cross_val_score(model, X_train, y_train, cv=3).mean()
+        wandb.log(params | {"cv_score": score})
+        return score
 
     study = optuna.create_study(direction="maximize")
-    study.optimize(lambda trial: objective(trial, X_train, y_train, X_test, y_test), n_trials=n_trials)
+    study.optimize(objective, n_trials=15)
 
     best_params = study.best_params
-    model = RandomForestClassifier(**best_params, random_state=42)
+    model = RandomForestClassifier(**best_params)
     model.fit(X_train, y_train)
-    predictions = model.predict(X_test)
+    y_pred = model.predict(X_test)
 
-    accuracy = accuracy_score(y_test, predictions)
-    precision = precision_score(y_test, predictions, average='weighted', zero_division=0)
-    recall = recall_score(y_test, predictions, average='weighted', zero_division=0)
-    f1 = f1_score(y_test, predictions, average='weighted', zero_division=0)
+    metrics = {
+        "accuracy": accuracy_score(y_test, y_pred),
+        "precision": precision_score(y_test, y_pred, average="weighted", zero_division=0),
+        "recall": recall_score(y_test, y_pred, average="weighted", zero_division=0),
+        "f1_score": f1_score(y_test, y_pred, average="weighted", zero_division=0),
+    }
+    wandb.log(metrics)
+    wandb_run.finish()
 
-    wandb.log({
-        "best_params": best_params,
-        "accuracy": accuracy,
-        "precision": precision,
-        "recall": recall,
-        "f1": f1,
-    })
+    top_trials = pd.DataFrame(study.trials_dataframe().sort_values(by="value", ascending=False).head(7))
+    return metrics, top_trials
+
+# SHAP & LIME
+def explainability(_):
+    target = df_global.columns[-1]
+    X = df_global.drop(target, axis=1)
+    y = df_global[target]
+
+    model = RandomForestClassifier()
+    model.fit(X, y)
 
     # SHAP
     explainer = shap.TreeExplainer(model)
-    shap_values = explainer.shap_values(X_test)
-    shap.summary_plot(shap_values, X_test, show=False)
-    shap_fig_path = "./figures/shap_summary.png"
+    shap_values = explainer.shap_values(X)
+    shap.summary_plot(shap_values, X, show=False)
+    shap_fig_path = "./shap_plot.png"
     plt.savefig(shap_fig_path)
-    wandb.log({"shap_summary": wandb.Image(shap_fig_path)})
     plt.clf()
 
     # LIME
-    lime_explainer = lime.lime_tabular.LimeTabularExplainer(X_train.values, feature_names=X_train.columns, class_names=['target'], mode='classification')
-    lime_exp = lime_explainer.explain_instance(X_test.iloc[0].values, model.predict_proba)
+    lime_explainer = lime.lime_tabular.LimeTabularExplainer(X.values, feature_names=X.columns, class_names=['target'], mode="classification")
+    lime_exp = lime_explainer.explain_instance(X.iloc[0].values, model.predict_proba)
     lime_fig = lime_exp.as_pyplot_figure()
-    lime_path = "./figures/lime_explanation.png"
-    lime_fig.savefig(lime_path)
-    wandb.log({"lime_explanation": wandb.Image(lime_path)})
+    lime_fig_path = "./lime_plot.png"
+    lime_fig.savefig(lime_fig_path)
     plt.clf()
 
-    return f"Best Hyperparameters: {best_params}<br>Accuracy: {accuracy}<br>Precision: {precision}<br>Recall: {recall}<br>F1-score: {f1}"
+    # Log to wandb
+    wandb.init(project="huggingface-data-analysis", name="Explainability", reinit=True)
+    wandb.log({
+        "shap_summary": wandb.Image(shap_fig_path),
+        "lime_explanation": wandb.Image(lime_fig_path)
+    })
+    wandb.finish()
+
+    return shap_fig_path, lime_fig_path
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## 📊 AI-Powered Data Analysis with Hyperparameter Optimization")
 
-# Gradio Interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## 📊 AI Data Analysis Agent with Hyperparameter Optimization")
     with gr.Row():
         with gr.Column():
-            file_input = gr.File(label="Upload CSV Dataset", type="filepath")
-            notes_input = gr.Textbox(label="Dataset Notes (Optional)", lines=3)
-            analyze_btn = gr.Button("Analyze", variant="primary")
-            optuna_trials = gr.Number(label="Number of Hyperparameter Tuning Trials", value=10)
-            tune_btn = gr.Button("Optimize Hyperparameters", variant="secondary")
+            file_input = gr.File(label="Upload CSV or Excel", type="file")
+            upload_btn = gr.Button("Upload & Clean")
+            df_output = gr.DataFrame(label="Cleaned Data Preview")
+
         with gr.Column():
-            analysis_output = gr.Markdown("### Analysis results will appear here...")
-            optuna_output = gr.HTML(label="Hyperparameter Tuning Results")
-            gallery = gr.Gallery(label="Data Visualizations", columns=2)
+            insights_output = gr.Textbox(label="Insights from SmolAgent", lines=15)
+            agent_btn = gr.Button("Run AI Agent (5 Insights + 5 Visualizations)")
+
+    with gr.Row():
+        train_btn = gr.Button("Train Model with Optuna + WandB")
+        metrics_output = gr.JSON(label="Performance Metrics")
+        trials_output = gr.DataFrame(label="Top 7 Hyperparameter Trials")
+
+    with gr.Row():
+        explain_btn = gr.Button("SHAP + LIME Explainability")
+        shap_img = gr.Image(label="SHAP Summary Plot")
+        lime_img = gr.Image(label="LIME Explanation")
 
-    analyze_btn.click(fn=analyze_data, inputs=[file_input, notes_input], outputs=[analysis_output, gallery])
-    tune_btn.click(fn=tune_hyperparameters, inputs=[file_input, optuna_trials], outputs=[optuna_output])
+    upload_btn.click(fn=upload_file, inputs=file_input, outputs=df_output)
+    agent_btn.click(fn=run_agent, inputs=df_output, outputs=insights_output)
+    train_btn.click(fn=train_model, inputs=df_output, outputs=[metrics_output, trials_output])
+    explain_btn.click(fn=explainability, inputs=df_output, outputs=[shap_img, lime_img])
 
-demo.launch(debug=True)
+demo.launch()