Update app.py

app.py

@@ -184,12 +184,19 @@ def compare_models():
     return results_df
     
 # 1. prepare_data should come first
-def prepare_data(df, target_column="target"):
+def prepare_data(df, target_column=None):
     from sklearn.model_selection import train_test_split
+
+    # If no target column is specified, select the first object column or the last column
+    if target_column is None:
+        target_column = df.select_dtypes(include=['object']).columns[0] if len(df.select_dtypes(include=['object']).columns) > 0 else df.columns[-1]
+    
     X = df.drop(columns=[target_column])
     y = df[target_column]
+    
     return train_test_split(X, y, test_size=0.2, random_state=42)
 
+
 def train_model(_):
     try:
         wandb.login(key=os.environ.get("WANDB_API_KEY"))