Update app.py

app.py

@@ -420,7 +420,6 @@ def prepare_data(df):
     
     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"))
@@ -439,36 +438,30 @@ def train_model(_):
             }
             model = RandomForestClassifier(**params)
             score = cross_val_score(model, X_train, y_train, cv=3).mean()
-        if wandb.run is None:
-            wandb.init(project="model_optimization", name=f"optuna_trial_{trial.number}", reinit=True)
             wandb.log({**params, "cv_score": score})
-            return score
+            return score  # ✅ Must be returned here
 
         study = optuna.create_study(direction="maximize")
         study.optimize(objective, n_trials=15)
 
         best_params = study.best_params
-        model = RandomForestClassifier()
+        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)
+        trial_rows = [dict(**t.params, score=t.value) for t in top_trials]
         trials_df = pd.DataFrame(trial_rows)
 
         return metrics, trials_df
@@ -478,6 +471,7 @@ def train_model(_):
         return {}, pd.DataFrame()
 
 
+
 def explainability(_):
     import warnings
     warnings.filterwarnings("ignore")