Update cancer.py

cancer.py

@@ -36,13 +36,20 @@ def preprocess_data(df):
     return train_test_split(X, y, test_size=0.2, random_state=23), preprocess
 
 # Train Models
-def train_model(model, X_train, y_train, preprocess):
+# Train Model
+def train_model(x_train, y_train, preprocess, model_name):
+    models = {
+        'Decision Tree': DecisionTreeClassifier(),
+        'Logistic Regression': LogisticRegression(),
+        'KNN': KNeighborsClassifier(),
+        'Random Forest': RandomForestClassifier(),
+        'XGBoost': XGBClassifier()
+    }
     pipeline = Pipeline([
         ('preprocessor', preprocess),
-        ('classifier', model)
+        ('classifier', models[model_name])
     ])
-    pipeline.fit(X_train, y_train)
-    return pipeline
+    pipeline.fit(x_train, y_train)
 
 # Streamlit UI
 st.set_page_config(page_title='Cancer Prediction App', layout='wide')