Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import xgboost as xgb
 import pandas as pd
 from datasets import load_dataset
+from sklearn.model_selection import train_test_split
 
 # Load the dataset
 dataset = load_dataset("Ammok/hair_health")
@@ -9,13 +10,36 @@ dataset = load_dataset("Ammok/hair_health")
 # Convert to Pandas DataFrame for exploration
 df = pd.DataFrame(dataset['train'])
 
-# Example: Train a simple XGBoost model
-X = df.drop(columns=["target_column"])  # Replace with your feature columns
-y = df["target_column"]  # Replace with your target column
+### PREPROCESSING
 
-# Train a basic XGBoost model (replace with custom model training code)
+# Replace "No Data" entries with NaN for missing values handling
+df.replace("No Data", pd.NA, inplace=True)
+
+# Handle missing numerical values with mean
+df.fillna(df.mean(), inplace=True)
+
+# Handle missing categorical values with mode
+for col in df.select_dtypes(include=['object']).columns:
+    df[col].fillna(df[col].mode()[0], inplace=True)
+
+# One-hot encoding for categorical variables
+categorical_cols = [
+    'Genetics', 'Hormonal Changes', 'Medical Conditions', 
+    'Medications & Treatments', 'Nutritional Deficiencies ', 'Stress',
+    'Poor Hair Care Habits ', 'Environmental Factors', 'Smoking', 'Weight Loss '
+]
+df = pd.get_dummies(df, columns=categorical_cols, drop_first=True)
+
+# Extract features and target
+X = df.drop(columns=["Hair Loss"])
+y = df["Hair Loss"]
+
+# Split the dataset into train and test sets
+X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+
+# Train a basic XGBoost model
 model = xgb.XGBClassifier()
-model.fit(X, y)
+model.fit(X_train, y_train)
 
 # Function for making predictions
 def predict(input_data):