Update app.py

app.py

@@ -41,12 +41,22 @@ X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_
 model = xgb.XGBClassifier()
 model.fit(X_train, y_train)
 
-# Function for making predictions
 def predict(input_data):
-    data = pd.DataFrame([input_data], columns=X.columns)
+    # Remove features the user wants to omit (e.g., 'Gender')
+    for col in X.columns:
+        if input_data.get(col) is None:
+            input_data.pop(col, None)
+
+    # Keep only the features the user provided
+    data = pd.DataFrame([input_data], columns=input_data.keys())
+
+    # Align with model columns and fill missing required columns with defaults
+    data = data.reindex(columns=X.columns, fill_value=X.mean())
+
     prediction = model.predict(data)
     return prediction[0]
 
+
 # Set up Gradio interface for data exploration
 def explore_data(row_number):
     return df.iloc[row_number].to_dict()