Update app.py

app.py

@@ -190,6 +190,51 @@ def train_model(_):
     wandb_run = wandb.init(project="huggingface-data-analysis", name=f"Optuna_Run_{run_counter}", reinit=True)
     run_counter += 1
 
+      
+        X_train, X_test, y_train, y_test = prepare_data()
+        
+        def objective(trial):
+            params = {
+                "n_estimators": trial.suggest_int("n_estimators", 50, 200),
+                "max_depth": trial.suggest_int("max_depth", 3, 10),
+            }
+            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(objective, n_trials=15)
+
+        best_params = study.best_params
+        model = RandomForestClassifier(**best_params)
+        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(_):