Update app.py

app.py

@@ -10,6 +10,10 @@ pipe = joblib.load("./model.pkl")
 title = "Premium Amount Prediction"
 description = "This model predicts the Premium Amount. Drag and drop any slice from the dataset or edit values as you wish in the dataframe component below."
 
+# Load and prepare dataset
+df = datasets.load_dataset("silvaKenpachi/mental_health")["train"].to_pandas()
+df.dropna(axis=0, inplace=True)
+
 # Load configuration
 with open("./config.json") as f:
     config_dict = json.load(f)
@@ -18,11 +22,6 @@ 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()
-df.dropna(axis=0, inplace=True)
-
 # Define input and output interfaces
 inputs = [gr.Dataframe(headers=headers, row_count=(2, "dynamic"), col_count=(len(headers), "fixed"), label="Input Data", interactive=True)]
 outputs = [gr.Dataframe(row_count=(2, "dynamic"), col_count=(1, "fixed"), label="Predictions", headers=["Depression"])]