Update app.py

app.py

@@ -198,75 +198,59 @@ def prepare_data(df, target_column=None):
     return train_test_split(X, y, test_size=0.3, random_state=42)
 
 
-def train_model(file, ab_choice="A"):
-    df = pd.read_csv(file.name)
-    df = preprocess_data(df)
-    X = df.drop("target", axis=1)
-    y = df["target"]
-
-    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
-
-    # SmolAgent
-    tuner = SmolAgent(model="gpt-4")
-    model = tuner.fit(X_train, y_train)
-
-    # Evaluate
-    y_pred = model.predict(X_test)
-    y_proba = model.predict_proba(X_test) if hasattr(model, "predict_proba") else None
-
-    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),
-    }
+def train_model(_):
+    try:
+        wandb.login(key=os.environ.get("WANDB_API_KEY"))
+        wandb_run = wandb.init(
+            project="huggingface-data-analysis",
+            name=f"Optuna_Run_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
+            reinit=True
+        )
+
+        X_train, X_test, y_train, y_test = prepare_data(df_global)
+
+        def objective(trial):
+            params = {
+                "n_estimators": trial.suggest_int("n_estimators", 50, 200),
+                "max_depth": trial.suggest_int("max_depth", 3, 10),
+            }
+            model = RandomForestClassifier()
+            score = cross_val_score(model, X_train, y_train, cv=3).mean()
+            wandb.log({**params, "cv_score": score})
+            return score
 
-    # Optional AUC
-    if y_proba is not None:
-        try:
-            if y_proba.shape[1] == 2:
-                metrics["roc_auc"] = roc_auc_score(y_test, y_proba[:, 1])
-            else:
-                metrics["roc_auc_ovr"] = roc_auc_score(y_test, y_proba, multi_class="ovr")
-        except Exception:
-            pass
-
-    # Classification report
-    metrics["classification_report"] = classification_report(y_test, y_pred, output_dict=True)
-
-    # Confusion Matrix Plot
-    fig, ax = plt.subplots(figsize=(6, 4))
-    ConfusionMatrixDisplay.from_predictions(y_test, y_pred, ax=ax)
-    plt.title("Confusion Matrix")
-    conf_matrix_path = "./conf_matrix.png"
-    plt.savefig(conf_matrix_path)
-    plt.close()
-
-    # Bar Plot of Metrics
-    def plot_metrics(metrics):
-        plt.figure(figsize=(6, 3))
-        keys = [k for k in metrics if isinstance(metrics[k], (int, float))]
-        values = [metrics[k] for k in keys]
-        plt.barh(keys, values, color="skyblue")
-        plt.xlabel("Score")
-        plt.title("Model Performance Metrics")
-        path = "./metrics_plot.png"
-        plt.tight_layout()
-        plt.savefig(path)
-        plt.close()
-        return path
-
-    metrics_plot_path = plot_metrics(metrics)
-
-    # Log to WandB
-    wandb.init(project="ab-test", name=f"variant_{ab_choice}", reinit=True)
-    wandb.log({**metrics, "confusion_matrix": wandb.Image(conf_matrix_path),
-               "metrics_plot": wandb.Image(metrics_plot_path),
-               "ab_variant": ab_choice})
-    wandb.finish()
-
-    # Return everything to Gradio
-    return metrics, pd.DataFrame.from_dict(metrics["classification_report"]).T, conf_matrix_path, metrics_plot_path
+        study = optuna.create_study(direction="maximize")
+        study.optimize(objective, n_trials=15)
+
+        best_params = study.best_params
+        model = RandomForestClassifier()
+        model.fit(X_train, y_train)
+        y_pred = model.predict(X_test)
+
+
+        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()
+
+        # Top 7 trials
+        top_trials = sorted(study.trials, key=lambda x: x.value, reverse=True)[:7]
+        trial_rows = []
+        for t in top_trials:
+            row = t.params.copy()
+            row["score"] = t.value
+            trial_rows.append(row)
+        trials_df = pd.DataFrame(trial_rows)
+
+        return metrics, trials_df
+
+    except Exception as e:
+        print(f"Training Error: {e}")
+        return {}, pd.DataFrame()
 
 
 def explainability(_):
@@ -359,25 +343,17 @@ with gr.Blocks() as demo:
             agent_btn = gr.Button("Run AI Agent (5 Insights + 5 Visualizations)")
 
     with gr.Row():
-        ab_dropdown = gr.Dropdown(choices=["A", "B"], label="Choose Model Variant", value="A")
         train_btn = gr.Button("Train Model with Optuna + WandB")
-
-    with gr.Row():
         metrics_output = gr.JSON(label="Performance Metrics")
         trials_output = gr.DataFrame(label="Top 7 Hyperparameter Trials")
 
-    with gr.Row():
-        conf_matrix_img = gr.Image(label="Confusion Matrix")
-        metric_plot_img = gr.Image(label="Metric Bar Plot")
-
     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")
 
-    # Button logic
     agent_btn.click(fn=analyze_data, inputs=[file_input], outputs=[insights_output, visual_output])
-    train_btn.click(fn=train_model, inputs=[file_input, ab_dropdown], outputs=[metrics_output, trials_output, conf_matrix_img, metric_plot_img])
+    train_btn.click(fn=train_model, inputs=[file_input], outputs=[metrics_output, trials_output])
     explain_btn.click(fn=explainability, inputs=[], outputs=[shap_img, lime_img])
 
-demo.launch(debug=True)
+demo.launch(debug=True)