update model to increase accuracy

main.py

@@ -1,10 +1,18 @@
 import gradio as gr
 import numpy as np
 import xgboost as xgb
-from ucimlrepo import fetch_ucirepo
 from sklearn.impute import SimpleImputer
 from sklearn.preprocessing import StandardScaler
 from imblearn.over_sampling import SMOTE
+from sklearn.model_selection import train_test_split, GridSearchCV
+from sklearn.metrics import (
+    accuracy_score,
+    precision_score,
+    recall_score,
+    f1_score,
+    confusion_matrix,
+)
+from ucimlrepo import fetch_ucirepo
 import os
 
 # Paths for saving/loading the model
@@ -15,6 +23,7 @@ heart_disease = fetch_ucirepo(id=45)
 X = heart_disease.data.features
 y = np.ravel(heart_disease.data.targets)
 
+# Preprocessing pipeline
 imputer = SimpleImputer(strategy="mean")
 X = imputer.fit_transform(X)
 scaler = StandardScaler()
@@ -22,26 +31,58 @@ X = scaler.fit_transform(X)
 smote = SMOTE(random_state=42)
 X_resampled, y_resampled = smote.fit_resample(X, y)
 
+# Train-test split
+X_train, X_test, y_train, y_test = train_test_split(
+    X_resampled, y_resampled, test_size=0.2, random_state=42, stratify=y_resampled
+)
+
 # Train or load the model
 if os.path.exists(MODEL_PATH):
     # Load pre-trained model
     model = xgb.Booster()
     model.load_model(MODEL_PATH)
 else:
-    # Train the model
-    dtrain = xgb.DMatrix(X_resampled, label=y_resampled)
+    # Hyperparameter tuning
     param_grid = {
-        "objective": "multi:softmax",
-        "num_class": len(np.unique(y_resampled)),
-        "eval_metric": "mlogloss",
-        "learning_rate": 0.1,
-        "max_depth": 5,
-        "subsample": 0.8,
-        "colsample_bytree": 0.8,
+        "max_depth": [4, 5, 6],
+        "learning_rate": [0.01, 0.05, 0.1],
+        "n_estimators": [100, 200, 300],
+        "subsample": [0.8, 1.0],
+        "colsample_bytree": [0.8, 1.0],
+        "gamma": [0, 1, 5],
+        "lambda": [1, 2, 3],
+        "alpha": [0, 1],
     }
-    model = xgb.train(params=param_grid, dtrain=dtrain, num_boost_round=100)
-    # Save the model
-    model.save_model(MODEL_PATH)
+
+    model = xgb.XGBClassifier(use_label_encoder=False, eval_metric="mlogloss")
+    grid_search = GridSearchCV(
+        estimator=model, param_grid=param_grid, scoring="accuracy", cv=5, verbose=1
+    )
+    grid_search.fit(X_train, y_train)
+
+    # Best model
+    best_model = grid_search.best_estimator_
+    best_model.save_model(MODEL_PATH)
+
+# Load the best model
+model = xgb.Booster()
+model.load_model(MODEL_PATH)
+
+# Evaluate model
+X_test_dmatrix = xgb.DMatrix(X_test)
+y_pred = model.predict(X_test_dmatrix)
+accuracy = accuracy_score(y_test, y_pred)
+precision = precision_score(y_test, y_pred, average="weighted")
+recall = recall_score(y_test, y_pred, average="weighted")
+f1 = f1_score(y_test, y_pred, average="weighted")
+conf_matrix = confusion_matrix(y_test, y_pred)
+
+print(f"Accuracy: {accuracy * 100:.2f}%")
+print(f"Precision: {precision:.2f}")
+print(f"Recall: {recall:.2f}")
+print(f"F1 Score: {f1:.2f}")
+print("Confusion Matrix:")
+print(conf_matrix)
 
 
 # Define prediction function
@@ -49,7 +90,7 @@ def predict(
     age, sex, cp, trestbps, chol, fbs, restecg, thalach, exang, oldpeak, slope, ca, thal
 ):
     # Convert string values to numeric where needed
-    sex = int(sex)  # Convert '1: Male' -> 1
+    sex = int(sex)
     cp = int(cp)
     fbs = int(fbs)
     restecg = int(restecg)
@@ -85,7 +126,7 @@ def predict(
     return int(prediction[0])
 
 
-# Update Gradio interface to return numeric values
+# Gradio interface
 feature_inputs = [
     gr.Number(label="Age (years)"),
     gr.Radio(label="Sex", choices=["0", "1"], type="value"),  # Male: 1, Female: 0
@@ -110,8 +151,6 @@ feature_inputs = [
     gr.Radio(label="Thalassemia (thal)", choices=["0", "1", "2", "3"], type="value"),
 ]
 
-
-# Define the Gradio interface
 interface = gr.Interface(
     fn=predict,
     inputs=feature_inputs,