Update app.py

app.py

@@ -13,7 +13,11 @@ description = "This model predicts the Premium Amount. Drag and drop any slice f
 # Load configuration
 with open("./config.json") as f:
     config_dict = json.load(f)
-headers = config_dict["sklearn"]["columns"]
+all_headers = config_dict["sklearn"]["columns"]
+
+# Filter headers to only include those present in the dataset
+headers = [col for col in all_headers if col in df.columns]
+
 
 # Load and prepare dataset
 df = datasets.load_dataset("silvaKenpachi/mental_health")["train"].to_pandas()