Update app.py

app.py

@@ -198,59 +198,75 @@ def prepare_data(df, target_column=None):
     return train_test_split(X, y, test_size=0.3, random_state=42)
 
 
-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
-
-        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
+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),
+    }
 
-    except Exception as e:
-        print(f"Training Error: {e}")
-        return {}, pd.DataFrame()
+    # 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
 
 
 def explainability(_):
@@ -351,9 +367,13 @@ with gr.Blocks() as demo:
         explain_btn = gr.Button("SHAP + LIME Explainability")
         shap_img = gr.Image(label="SHAP Summary Plot")
         lime_img = gr.Image(label="LIME Explanation")
+        metrics_output = gr.JSON(label="Evaluation Metrics")
+        trials_output = gr.Dataframe(label="Classification Report")
+        conf_matrix_img = gr.Image(label="Confusion Matrix")
+        metric_plot_img = gr.Image(label="Metric Bar Plot")
 
     agent_btn.click(fn=analyze_data, inputs=[file_input], outputs=[insights_output, visual_output])
-    train_btn.click(fn=train_model, inputs=[file_input], outputs=[metrics_output, trials_output])
+    train_btn.click(fn=train_model, inputs=[file_input, ab_dropdown], outputs=[metrics_output, trials_output, conf_matrix_img, metric_plot_img])
     explain_btn.click(fn=explainability, inputs=[], outputs=[shap_img, lime_img])
 
 demo.launch(debug=True)