Build a Heart Disease Prediction Model
Train a classifier on real clinical data to predict heart disease.
Introduction
30 Second Summary
Heart disease is the leading cause of death worldwide. Doctors already use clinical measurements like blood pressure, cholesterol, and heart rate to assess risk, but what if a computer could learn from hundreds of real patient records to predict heart disease automatically?
In this project, you will build a machine learning model that predicts whether a patient has heart disease based on real clinical data from 303 patients.
What You'll Build
A complete ML pipeline that loads real patient data, visualizes it, preprocesses it, trains two classifiers, evaluates them with cross-validation, and exports the best model as a file ready for API integration.
By the end of this project, you'll have:
- A Google Colab notebook containing your full ML pipeline, from raw data to saved model.
- Exploratory data analysis visualizations including distribution plots and a correlation heatmap.
- Two trained classification models (Logistic Regression and Random Forest) compared side-by-side with confusion matrices.
- A saved .joblib model file ready for API integration in a follow-on project.
- Secret Mission: Tune your Random Forest model with GridSearchCV to squeeze out better performance.
Are there any prerequisites?
You should be comfortable with basic Python (variables, functions, and loops). No prior machine learning experience is needed.
Set Up Your Colab Notebook
Every machine learning project starts with a workspace where you can write code, run experiments, and see results instantly.
In this step, you will open Google Colab, create a fresh notebook, and confirm that all the libraries you need are already available.
In this step, get ready to:
- Open Google Colab and create a new notebook.
- Verify that all required ML libraries are pre-installed.
- Understand the dataset you will be working with.
Open Google Colab and create a new notebook
Google Colab is a free, browser-based Python environment hosted by Google. It comes with popular data science libraries pre-installed, so you can start coding immediately without managing any local setup.
- Open your browser and navigate to colab.research.google.com.
- Sign in with your Google account if prompted.
- Click New Notebook on the Colab landing page.
- A blank notebook opens with a default name like Untitled0.ipynb. Click the notebook name at the top-left of the page.
- Replace the default name with Heart Disease Prediction and press Enter.
What is Google Colab?
Google Colab is a free cloud-hosted Jupyter notebook environment. It runs Python in your browser with no local installation required.
Colab autosaves your work to Google Drive every few seconds, so you never need to manually save.
Verify your pre-installed libraries
Colab comes with all the data science libraries this project needs already installed. You will import them now and print their versions to confirm everything is ready.
- Click into the empty code cell in your notebook and paste the following code:
# --- Cell 1: Import Libraries ---
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report, ConfusionMatrixDisplay
from sklearn.pipeline import Pipeline
import joblib
print("All libraries imported successfully!")
print(f"scikit-learn version: {__import__('sklearn').__version__}")
print(f"pandas version: {pd.__version__}")- Run the cell by pressing Shift+Enter or clicking the play button on the left side of the cell.
What does this code do?
This cell imports all the libraries you will use throughout the project. Each one has a specific role:
- pandas loads and manipulates tabular data (your patient dataset).
- numpy handles numerical operations under the hood.
- matplotlib and seaborn create visualizations like distribution plots and heatmaps.
- sklearn (scikit-learn) provides the machine learning tools: splitting data, scaling features, training models, and evaluating results.
- joblib saves your trained model to a file so it can be loaded later.
You should see output that reads All libraries imported successfully! followed by version numbers for scikit-learn and pandas.
Seeing an ImportError?
If any import fails, your Colab runtime may need a restart. Click Runtime in the top menu, then Restart session. Run the cell again after the restart completes.
If the error persists, Colab may have updated its default packages. Add a cell above with !pip install scikit-learn pandas matplotlib seaborn and run it before re-running your imports.
Help me debug my import error in Google Colab.
✔️ Awesome, I've got everything!
Great. Your notebook is ready and all libraries are confirmed.
ⓧ I'd like to double check the full code
Here is the complete code for Cell 1. Make sure your cell matches this exactly:
# --- Cell 1: Import Libraries ---
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report, ConfusionMatrixDisplay
from sklearn.pipeline import Pipeline
import joblib
print("All libraries imported successfully!")
print(f"scikit-learn version: {__import__('sklearn').__version__}")
print(f"pandas version: {pd.__version__}")Understand the dataset you will use
Before you load any data, it helps to know what you are working with. This project uses a real clinical dataset from the UCI Machine Learning Repository.
What is the UCI Heart Disease dataset?
The Cleveland Heart Disease dataset contains records from 303 real patients. Each record has 13 clinical features (like age, blood pressure, cholesterol, and maximum heart rate) plus a target column indicating whether heart disease is present.
The target value ranges from 0 (no disease) to 4 (severe). You will simplify this to a binary classification problem: 0 means no disease, and 1 means disease is present.
The dataset is free and publicly available for research. You will load it directly from the UCI repository URL in the next step.
Your workspace is set up and every tool is confirmed working. Next up, you will load the heart disease dataset and explore it visually to understand the patterns hiding in the data.
Load and Explore the Dataset
Your Colab notebook is ready with all libraries imported. Now it's time to bring in the real clinical data you'll be training your model on.
In this step, you'll load the UCI Cleveland Heart Disease dataset, check it for missing values, and create visualizations that reveal which features relate most to heart disease.
In this step, get ready to:
- Load the UCI Heart Disease dataset into a pandas DataFrame.
- Explore missing values and binarize the target column.
- Create distribution plots and a correlation heatmap.
Load the dataset from the UCI repository
The UCI Cleveland Heart Disease dataset contains records from 303 real patients. Each patient has 13 clinical features (age, blood pressure, cholesterol, and more) plus a target column indicating disease presence.
The raw data file has no column headers and uses ? characters to represent missing values. You'll handle both when loading the file.
- Load the dataset into a pandas DataFrame by pasting this code into a new cell:
# --- Cell 2: Load Dataset ---
url = "https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/processed.cleveland.data"
column_names = [
'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs',
'restecg', 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'
]
df = pd.read_csv(url, names=column_names, na_values='?')
print(f"Dataset shape: {df.shape}")
print(f"\nFirst 5 rows:")
print(df.head())
print(f"\nData types:")
print(df.dtypes)What does this code do?
- url points to the raw dataset hosted on the UCI Machine Learning Repository.
- column_names provides the 14 column headers manually since the raw file has none.
- na_values='?' tells pandas to treat any ? character as a missing value (NaN) instead of a string.
- The print statements display the dataset's dimensions, first 5 rows, and column data types so you can verify the load worked.
- Run the cell by pressing Shift+Enter.
- You should see output showing a shape of (303, 14), a table of the first 5 rows with columns like age and chol, and a list of data types.
Seeing a connection error?
The dataset loads from a URL, so you need an active internet connection. If Colab shows a connection error, click Runtime in the top menu, then Disconnect and delete runtime. Then click Connect to start a fresh session and re-run both cells from the top.
help me debug a connection error when loading data from a URL in Google Colab
Explore missing values and binarize the target
Before building any model, you need to understand what's in your data. Two critical checks: are there missing values that could crash your training, and what does the target column actually represent?
- Check for missing values and examine the target distribution by pasting this into a new cell:
# --- Cell 3: Initial Exploration ---
print(f"Missing values per column:")
print(df.isnull().sum())
print(f"\nTarget distribution (original):")
print(df['target'].value_counts())What does this code do?
- df.isnull().sum() counts how many missing values exist in each column. You'll see that ca and thal have a few missing entries.
- value_counts() shows the distribution of the target column. The original values range from 0 (no disease) to 4 (severe disease).
- Run the cell. You should see that ca has 4 missing values and thal has 2. The target column shows values 0 through 4.
The target has five severity levels, but for this project you only need to predict a yes/no answer: does the patient have heart disease or not? This simplification is called binarization.
- Binarize the target column by pasting this into a new cell:
# --- Cell 4: Binarize Target ---
df['target'] = df['target'].apply(lambda x: 1 if x > 0 else 0)
print(f"Binary target distribution:")
print(df['target'].value_counts())
print(f"\nDisease present: {df['target'].sum()}")
print(f"No disease: {(df['target'] == 0).sum()}")What does this code do?
The lambda function converts any target value greater than 0 to 1 (disease present), and keeps 0 as 0 (no disease). This turns a multi-class problem into a binary classification problem.
This approach matches how all published experiments on this dataset work. The UCI documentation defines 0 as less than 50% diameter narrowing and 1+ as greater than 50% narrowing.
- Run the cell. You should see roughly 138 patients with no disease and 165 with disease present.
Seeing a KeyError on 'target'?
This means Cell 2 didn't run successfully. Go back to Cell 2 and re-run it first. The DataFrame needs to exist before you can transform its columns.
help me fix a KeyError when accessing a column in my pandas DataFrame
Create exploratory visualizations
Numbers alone don't reveal patterns easily. Visualizations let you see at a glance which clinical features differ between patients with and without heart disease.
- Create distribution plots for six key features by pasting this into a new cell:
# --- Cell 5: EDA - Distribution Plots ---
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
sns.histplot(data=df, x='age', hue='target', kde=True, ax=axes[0, 0])
axes[0, 0].set_title('Age Distribution by Target')
sns.histplot(data=df, x='trestbps', hue='target', kde=True, ax=axes[0, 1])
axes[0, 1].set_title('Resting Blood Pressure by Target')
sns.histplot(data=df, x='chol', hue='target', kde=True, ax=axes[0, 2])
axes[0, 2].set_title('Cholesterol by Target')
sns.histplot(data=df, x='thalach', hue='target', kde=True, ax=axes[1, 0])
axes[1, 0].set_title('Max Heart Rate by Target')
sns.histplot(data=df, x='oldpeak', hue='target', kde=True, ax=axes[1, 1])
axes[1, 1].set_title('ST Depression by Target')
sns.countplot(data=df, x='cp', hue='target', ax=axes[1, 2])
axes[1, 2].set_title('Chest Pain Type by Target')
plt.tight_layout()
plt.show()What do these plots show?
- Each plot splits patients by target (0 = no disease, 1 = disease) so you can spot visual differences.
- hue='target' colors the bars differently for each group.
- kde=True adds a smooth curve showing the distribution shape, making it easier to compare the two groups.
- Look for features where the two colored distributions are clearly separated. Those features will be most useful to your model.
- Run the cell. You should see a 2x3 grid of histograms with overlapping orange and blue distributions for each feature.
Now you'll generate a correlation heatmap to see how all features relate to each other and to the target.
- Generate the correlation heatmap by pasting this into a new cell:
# --- Cell 6: EDA - Correlation Heatmap ---
plt.figure(figsize=(12, 8))
sns.heatmap(df.corr(numeric_only=True), annot=True, cmap='coolwarm', center=0, fmt='.2f')
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.show()How do I read a correlation heatmap?
Each cell shows the correlation between two features on a scale from -1 to +1. Values close to +1 mean the features increase together. Values close to -1 mean one increases as the other decreases. Values near 0 mean no linear relationship.
Look at the target row at the bottom. Features with high absolute correlation to the target (positive or negative) are the most predictive for your model.
- Run the cell. You should see a color-coded grid with numbers in each square. Red means positive correlation, blue means negative.
Seeing a blank or overlapping plot?
Make sure plt.tight_layout() is included before plt.show(). Without it, subplot labels can overlap. If the plot is blank, check that you ran Cell 2 and Cell 4 first so df exists with data.
help me fix blank or overlapping matplotlib plots in Colab
✔️ Awesome, I've got everything!
Great. Make sure all cells (1 through 6) have run successfully before moving on.
ⓧ I'd like to double check the full code
Here's the complete notebook code up to this point. Compare your cells against this:
# --- Cell 1: Import Libraries ---
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report, ConfusionMatrixDisplay
from sklearn.pipeline import Pipeline
import joblib
print("All libraries imported successfully!")
print(f"scikit-learn version: {__import__('sklearn').__version__}")
print(f"pandas version: {pd.__version__}")
# --- Cell 2: Load Dataset ---
url = "https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/processed.cleveland.data"
column_names = [
'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs',
'restecg', 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'
]
df = pd.read_csv(url, names=column_names, na_values='?')
print(f"Dataset shape: {df.shape}")
print(f"\nFirst 5 rows:")
print(df.head())
print(f"\nData types:")
print(df.dtypes)
# --- Cell 3: Initial Exploration ---
print(f"Missing values per column:")
print(df.isnull().sum())
print(f"\nTarget distribution (original):")
print(df['target'].value_counts())
# --- Cell 4: Binarize Target ---
df['target'] = df['target'].apply(lambda x: 1 if x > 0 else 0)
print(f"Binary target distribution:")
print(df['target'].value_counts())
print(f"\nDisease present: {df['target'].sum()}")
print(f"No disease: {(df['target'] == 0).sum()}")
# --- Cell 5: EDA - Distribution Plots ---
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
sns.histplot(data=df, x='age', hue='target', kde=True, ax=axes[0, 0])
axes[0, 0].set_title('Age Distribution by Target')
sns.histplot(data=df, x='trestbps', hue='target', kde=True, ax=axes[0, 1])
axes[0, 1].set_title('Resting Blood Pressure by Target')
sns.histplot(data=df, x='chol', hue='target', kde=True, ax=axes[0, 2])
axes[0, 2].set_title('Cholesterol by Target')
sns.histplot(data=df, x='thalach', hue='target', kde=True, ax=axes[1, 0])
axes[1, 0].set_title('Max Heart Rate by Target')
sns.histplot(data=df, x='oldpeak', hue='target', kde=True, ax=axes[1, 1])
axes[1, 1].set_title('ST Depression by Target')
sns.countplot(data=df, x='cp', hue='target', ax=axes[1, 2])
axes[1, 2].set_title('Chest Pain Type by Target')
plt.tight_layout()
plt.show()
# --- Cell 6: EDA - Correlation Heatmap ---
plt.figure(figsize=(12, 8))
sns.heatmap(df.corr(numeric_only=True), annot=True, cmap='coolwarm', center=0, fmt='.2f')
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.show()You now have a clear picture of your data. You know which columns have missing values and which features are most correlated with heart disease. Next up, you'll clean the data and split it for model training.
Preprocess the Data
Your visualizations revealed which features correlate most with heart disease. Now it's time to clean the data and prepare it for model training.
Machine learning models cannot handle missing values or unscaled features. In this step, you'll fill in the gaps, separate features from the target, and split the data into training and test sets the way a professional data scientist would.
In this step, get ready to:
- Fill missing values in the dataset with median imputation.
- Split the data into training and test sets with stratification.
- Understand why scikit-learn Pipelines prevent data leakage.
Handle missing values
Back in Step 2, you saw that the ca and thal columns contain missing values (originally represented as "?" characters in the raw file). Models cannot train on rows with gaps, so you need to fill them in.
- Create a new code cell in your notebook and fill the missing values by running this code:
# --- Cell 7: Handle Missing Values ---
print(f"Missing values before cleaning:")
print(df.isnull().sum()[df.isnull().sum() > 0])
df['ca'] = df['ca'].fillna(df['ca'].median())
df['thal'] = df['thal'].fillna(df['thal'].median())
print(f"\nTotal missing values after cleaning: {df.isnull().sum().sum()}")What does this code do?
First, it prints only the columns that have missing values so you can confirm which ones need fixing.
Then .fillna(df['ca'].median()) replaces each missing value with the median of that column. The median is a safe choice here because the dataset only has a handful of missing values (4 in ca and 2 in thal), and the median is not skewed by outliers.
- Run the cell by pressing **Shift+Enter**.
- You should see the missing value counts for ca and thal before cleaning, followed by Total missing values after cleaning: 0.
Still seeing missing values after cleaning?
Make sure you ran all previous cells in order. If you restarted the runtime, Colab loses all variables. Go to **Runtime**, then **Run all** to re-execute every cell from the top.
Check that your column names match exactly. The column is ca (lowercase), not Ca or CA.
Help me debug my missing value imputation code
Split into features and target, then train/test split
Your model needs to learn from some data and be tested on data it has never seen. This is how you measure whether it actually learned patterns or just memorized the training examples.
- Create a new code cell and split your data into training and test sets by running this code:
# --- Cell 8: Train/Test Split ---
X = df.drop('target', axis=1)
y = df['target']
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"Training set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")
print(f"\nTraining target distribution:")
print(y_train.value_counts())What does this code do?
- X = df.drop('target', axis=1) creates the features matrix. This is everything the model uses to make predictions (all 13 clinical features).
- y = df['target'] extracts the target vector. This is the answer the model is trying to predict (0 = no disease, 1 = disease present).
- test_size=0.2 holds back 20% of the data for testing. The model never sees this portion during training.
- stratify=y preserves the same ratio of disease/no-disease patients in both the training and test sets. Without this, you could accidentally get a test set with mostly one class.
- random_state=42 makes the split reproducible. Anyone running this code gets the same split.
- Run the cell by pressing **Shift+Enter**.
- You should see Training set size: 242 and Test set size: 61, followed by the target distribution showing both classes represented in the training data.
Seeing a NameError for train_test_split?
This means Cell 1 (imports) was not run in this session. Go back to your first cell and run it again to load all the library imports.
Help me fix my train_test_split error in Google Colab
Understand why we use a Pipeline for preprocessing
Before you train models in the next step, there is one more preprocessing concept to understand. Features like age (range 29-77) and cholesterol (range 126-564) are on very different scales. Many ML algorithms perform better when features are scaled to a similar range.
What is data leakage?
If you scale your data before splitting, the scaler "sees" statistics (mean, standard deviation) from the test set. This gives the model a subtle unfair advantage because it learned information about data it is supposed to have never seen.
The scikit-learn documentation states: "Pipelines help avoid leaking statistics from your test data into the trained model in cross-validation, by ensuring that the same samples are used to train the transformers and predictors."
Why use a Pipeline instead of scaling manually?
A Pipeline chains a scaler and a classifier into one object. When you call .fit(), the Pipeline fits the scaler on training data only, transforms it, and then trains the classifier. This eliminates any chance of data leakage.
This is the professional approach that interviewers expect to see. You will build Pipelines in the next step when you train your models.
✔️ Awesome, I've got everything!
Great. Double check that both Cell 7 and Cell 8 ran without errors before moving on.
ⓧ I'd like to double check the full code
Here is all the code in your notebook up to this point. Each section corresponds to a separate code cell.
# --- Cell 1: Import Libraries ---
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report, ConfusionMatrixDisplay
from sklearn.pipeline import Pipeline
import joblib
print("All libraries imported successfully!")
print(f"scikit-learn version: {__import__('sklearn').__version__}")
print(f"pandas version: {pd.__version__}")
# --- Cell 2: Load Dataset ---
url = "https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/processed.cleveland.data"
column_names = [
'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs',
'restecg', 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'
]
df = pd.read_csv(url, names=column_names, na_values='?')
print(f"Dataset shape: {df.shape}")
print(f"\nFirst 5 rows:")
print(df.head())
print(f"\nData types:")
print(df.dtypes)
# --- Cell 3: Initial Exploration ---
print(f"Missing values per column:")
print(df.isnull().sum())
print(f"\nTarget distribution (original):")
print(df['target'].value_counts())
# --- Cell 4: Binarize Target ---
df['target'] = df['target'].apply(lambda x: 1 if x > 0 else 0)
print(f"Binary target distribution:")
print(df['target'].value_counts())
print(f"\nDisease present: {df['target'].sum()}")
print(f"No disease: {(df['target'] == 0).sum()}")
# --- Cell 5: EDA - Distribution Plots ---
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
sns.histplot(data=df, x='age', hue='target', kde=True, ax=axes[0, 0])
axes[0, 0].set_title('Age Distribution by Target')
sns.histplot(data=df, x='trestbps', hue='target', kde=True, ax=axes[0, 1])
axes[0, 1].set_title('Resting Blood Pressure by Target')
sns.histplot(data=df, x='chol', hue='target', kde=True, ax=axes[0, 2])
axes[0, 2].set_title('Cholesterol by Target')
sns.histplot(data=df, x='thalach', hue='target', kde=True, ax=axes[1, 0])
axes[1, 0].set_title('Max Heart Rate by Target')
sns.histplot(data=df, x='oldpeak', hue='target', kde=True, ax=axes[1, 1])
axes[1, 1].set_title('ST Depression by Target')
sns.countplot(data=df, x='cp', hue='target', ax=axes[1, 2])
axes[1, 2].set_title('Chest Pain Type by Target')
plt.tight_layout()
plt.show()
# --- Cell 6: EDA - Correlation Heatmap ---
plt.figure(figsize=(12, 8))
sns.heatmap(df.corr(numeric_only=True), annot=True, cmap='coolwarm', center=0, fmt='.2f')
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.show()
# --- Cell 7: Handle Missing Values ---
print(f"Missing values before cleaning:")
print(df.isnull().sum()[df.isnull().sum() > 0])
df['ca'] = df['ca'].fillna(df['ca'].median())
df['thal'] = df['thal'].fillna(df['thal'].median())
print(f"\nTotal missing values after cleaning: {df.isnull().sum().sum()}")
# --- Cell 8: Train/Test Split ---
X = df.drop('target', axis=1)
y = df['target']
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"Training set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")
print(f"\nTraining target distribution:")
print(y_train.value_counts())Your data is clean, split, and ready for action. Next up, you'll train two different classification models and see which one predicts heart disease more accurately.
Train and Compare Models
Your data is clean, split, and ready to learn from. Now it's time for the exciting part: training machine learning models that can actually predict heart disease from patient measurements.
You will train two different classifiers and compare them side-by-side. This mirrors what data scientists do in real projects. A single model gives you one answer, but comparing models tells you which approach generalizes best to patients the model has never seen.
In this step, get ready to:
- Train a Logistic Regression pipeline on the training data.
- Train a Random Forest pipeline on the training data.
- Compare both models with 5-fold cross-validation.
Train a Logistic Regression pipeline
Logistic Regression is the classic baseline for binary classification problems. It models the probability that a patient has heart disease as a function of their clinical features.
- Create a new code cell in your notebook and add this code to train the Logistic Regression model:
# --- Cell 9: Train Logistic Regression ---
lr_pipeline = Pipeline([
('scaler', StandardScaler()),
('classifier', LogisticRegression(random_state=42, max_iter=1000))
])
lr_pipeline.fit(X_train, y_train)
lr_predictions = lr_pipeline.predict(X_test)
print("Logistic Regression Results:")
print(f"Accuracy: {accuracy_score(y_test, lr_predictions):.4f}")
print(f"\nClassification Report:")
print(classification_report(y_test, lr_predictions, target_names=['No Disease', 'Disease']))What does this code do?
The Pipeline bundles two steps into one object: first StandardScaler normalizes the features (so large values like cholesterol don't dominate), then LogisticRegression learns the decision boundary.
The classification_report shows precision (how many predicted-disease patients actually had disease), recall (how many real disease patients were caught), and F1-score (the balance between both).
- Run the cell.
- You should see an accuracy score (a decimal like 0.8361) followed by a table showing precision, recall, and F1-score for both "No Disease" and "Disease" classes.
Seeing a NameError?
If you see NameError: name 'X_train' is not defined, make sure you ran the Cell 8 (train/test split) cell first. Colab cells must be run in order.
If you see NameError: name 'Pipeline' is not defined, go back and run Cell 1 (imports) again.
Help me debug my Logistic Regression pipeline error.
Train a Random Forest pipeline
Random Forest is an ensemble method that builds many decision trees and averages their predictions. It often outperforms Logistic Regression on complex patterns because it can capture non-linear relationships between features.
- Create a new code cell and add this code to train the Random Forest model:
# --- Cell 10: Train Random Forest ---
rf_pipeline = Pipeline([
('scaler', StandardScaler()),
('classifier', RandomForestClassifier(n_estimators=100, random_state=42))
])
rf_pipeline.fit(X_train, y_train)
rf_predictions = rf_pipeline.predict(X_test)
print("Random Forest Results:")
print(f"Accuracy: {accuracy_score(y_test, rf_predictions):.4f}")
print(f"\nClassification Report:")
print(classification_report(y_test, rf_predictions, target_names=['No Disease', 'Disease']))What does this code do?
n_estimators=100 tells the Random Forest to build 100 individual decision trees. Each tree sees a random subset of the data and features, then the forest averages their votes for the final prediction.
The structure is identical to the Logistic Regression pipeline. Only the classifier step changes. This makes comparing models clean and fair because both use the same scaling.
- Run the cell.
- You should see an accuracy score and classification report table. Compare the numbers to Logistic Regression above.
Accuracy looks identical to Logistic Regression?
Double check that you used RandomForestClassifier and not LogisticRegression inside this pipeline. It is also normal for accuracies to be close on small datasets.
Help me compare my two model results.
Compare models with 5-fold cross-validation
A single train/test split can be misleading. If you got lucky (or unlucky) with which patients landed in the test set, your accuracy could be artificially high or low.
Cross-validation solves this by splitting the data 5 different ways and averaging the results. This gives you a much more reliable estimate of how well your model generalizes.
- Create a new code cell and add this code to run cross-validation on both models:
# --- Cell 11: Cross-Validation Comparison ---
lr_cv_scores = cross_val_score(lr_pipeline, X, y, cv=5, scoring='accuracy')
rf_cv_scores = cross_val_score(rf_pipeline, X, y, cv=5, scoring='accuracy')
print("5-Fold Cross-Validation Results:")
print(f"\nLogistic Regression: {lr_cv_scores.mean():.4f} (+/- {lr_cv_scores.std():.4f})")
print(f"Random Forest: {rf_cv_scores.mean():.4f} (+/- {rf_cv_scores.std():.4f})")What does this code do?
cross_val_score splits the full dataset (X and y) into 5 folds. It trains on 4 folds and tests on the remaining 1, rotating through all 5 combinations.
The mean tells you the average accuracy across all 5 runs. The standard deviation (the +/- number) tells you how stable the model is. A lower standard deviation means the model performs consistently regardless of which patients are in the test set.
- Run the cell.
- You should see two lines showing the mean accuracy and standard deviation for each model. The model with the higher mean and lower standard deviation generalizes better.
Cell taking a long time?
Cross-validation trains each model 5 times, so it takes a few seconds longer than a single fit. With 303 rows and 13 features, it should still complete within 5 seconds on Colab.
If you see a timeout, check that your Colab runtime is still connected by looking at the green checkmark in the top right.
Help me figure out why my cross-validation is slow or erroring.
Why cross-validate on the full dataset?
Notice that cross_val_score receives X and y (the full 303 patients), not just the training split. Cross-validation creates its own internal train/test splits within each fold, so it needs all the data to work with.
The Pipeline handles scaling correctly within each fold. The scaler is fit only on the fold's training portion, preventing data leakage.
✔️ Awesome, I've got everything!
Great. Double check you have all three cells (9, 10, and 11) run in order with output visible.
ⓧ I'd like to double check the full code
Here is the full notebook code through Cell 11. Compare your cells against this:
# --- Cell 1: Import Libraries ---
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report, ConfusionMatrixDisplay
from sklearn.pipeline import Pipeline
import joblib
print("All libraries imported successfully!")
print(f"scikit-learn version: {__import__('sklearn').__version__}")
print(f"pandas version: {pd.__version__}")
# --- Cell 2: Load Dataset ---
url = "https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/processed.cleveland.data"
column_names = [
'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs',
'restecg', 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'
]
df = pd.read_csv(url, names=column_names, na_values='?')
print(f"Dataset shape: {df.shape}")
print(f"\nFirst 5 rows:")
print(df.head())
print(f"\nData types:")
print(df.dtypes)
# --- Cell 3: Initial Exploration ---
print(f"Missing values per column:")
print(df.isnull().sum())
print(f"\nTarget distribution (original):")
print(df['target'].value_counts())
# --- Cell 4: Binarize Target ---
df['target'] = df['target'].apply(lambda x: 1 if x > 0 else 0)
print(f"Binary target distribution:")
print(df['target'].value_counts())
print(f"\nDisease present: {df['target'].sum()}")
print(f"No disease: {(df['target'] == 0).sum()}")
# --- Cell 5: EDA - Distribution Plots ---
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
sns.histplot(data=df, x='age', hue='target', kde=True, ax=axes[0, 0])
axes[0, 0].set_title('Age Distribution by Target')
sns.histplot(data=df, x='trestbps', hue='target', kde=True, ax=axes[0, 1])
axes[0, 1].set_title('Resting Blood Pressure by Target')
sns.histplot(data=df, x='chol', hue='target', kde=True, ax=axes[0, 2])
axes[0, 2].set_title('Cholesterol by Target')
sns.histplot(data=df, x='thalach', hue='target', kde=True, ax=axes[1, 0])
axes[1, 0].set_title('Max Heart Rate by Target')
sns.histplot(data=df, x='oldpeak', hue='target', kde=True, ax=axes[1, 1])
axes[1, 1].set_title('ST Depression by Target')
sns.countplot(data=df, x='cp', hue='target', ax=axes[1, 2])
axes[1, 2].set_title('Chest Pain Type by Target')
plt.tight_layout()
plt.show()
# --- Cell 6: EDA - Correlation Heatmap ---
plt.figure(figsize=(12, 8))
sns.heatmap(df.corr(numeric_only=True), annot=True, cmap='coolwarm', center=0, fmt='.2f')
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.show()
# --- Cell 7: Handle Missing Values ---
print(f"Missing values before cleaning:")
print(df.isnull().sum()[df.isnull().sum() > 0])
df['ca'] = df['ca'].fillna(df['ca'].median())
df['thal'] = df['thal'].fillna(df['thal'].median())
print(f"\nTotal missing values after cleaning: {df.isnull().sum().sum()}")
# --- Cell 8: Train/Test Split ---
X = df.drop('target', axis=1)
y = df['target']
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"Training set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")
print(f"\nTraining target distribution:")
print(y_train.value_counts())
# --- Cell 9: Train Logistic Regression ---
lr_pipeline = Pipeline([
('scaler', StandardScaler()),
('classifier', LogisticRegression(random_state=42, max_iter=1000))
])
lr_pipeline.fit(X_train, y_train)
lr_predictions = lr_pipeline.predict(X_test)
print("Logistic Regression Results:")
print(f"Accuracy: {accuracy_score(y_test, lr_predictions):.4f}")
print(f"\nClassification Report:")
print(classification_report(y_test, lr_predictions, target_names=['No Disease', 'Disease']))
# --- Cell 10: Train Random Forest ---
rf_pipeline = Pipeline([
('scaler', StandardScaler()),
('classifier', RandomForestClassifier(n_estimators=100, random_state=42))
])
rf_pipeline.fit(X_train, y_train)
rf_predictions = rf_pipeline.predict(X_test)
print("Random Forest Results:")
print(f"Accuracy: {accuracy_score(y_test, rf_predictions):.4f}")
print(f"\nClassification Report:")
print(classification_report(y_test, rf_predictions, target_names=['No Disease', 'Disease']))
# --- Cell 11: Cross-Validation Comparison ---
lr_cv_scores = cross_val_score(lr_pipeline, X, y, cv=5, scoring='accuracy')
rf_cv_scores = cross_val_score(rf_pipeline, X, y, cv=5, scoring='accuracy')
print("5-Fold Cross-Validation Results:")
print(f"\nLogistic Regression: {lr_cv_scores.mean():.4f} (+/- {lr_cv_scores.std():.4f})")
print(f"Random Forest: {rf_cv_scores.mean():.4f} (+/- {rf_cv_scores.std():.4f})")You now have two trained models with solid cross-validation scores. Next up, you will visualize their performance with confusion matrices and export the best model as a file ready for API integration.
Evaluate and Export Your Model
You've trained two models and compared their cross-validation scores. Now it's time to see exactly where each model succeeds and fails.
In this step, you'll visualize both models with confusion matrices, pick the best performer, and export it as a file you can load into any Python application.
In this step, get ready to:
- Visualize both models with confusion matrices.
- Save the best model as a .joblib file.
- Load and test the saved model on a sample patient.
Visualize model performance with confusion matrices
Accuracy gives you one number, but it hides the details. A confusion matrix breaks down predictions into four categories: true positives, false positives, true negatives, and false negatives.
- Create a new code cell and paste the following code to plot confusion matrices for both models side-by-side:
# --- Cell 12: Confusion Matrix Visualization ---
fig, axes = plt.subplots(1, 2, figsize=(12, 5))
ConfusionMatrixDisplay.from_predictions(
y_test, lr_predictions,
display_labels=['No Disease', 'Disease'],
ax=axes[0],
cmap='Blues'
)
axes[0].set_title('Logistic Regression')
ConfusionMatrixDisplay.from_predictions(
y_test, rf_predictions,
display_labels=['No Disease', 'Disease'],
ax=axes[1],
cmap='Blues'
)
axes[1].set_title('Random Forest')
plt.tight_layout()
plt.show()What does this code do?
ConfusionMatrixDisplay.from_predictions takes the true labels and predictions, then renders a grid showing how many patients were correctly or incorrectly classified.
The four quadrants tell you: how many healthy patients were correctly identified (true negative, top-left), how many sick patients were correctly identified (true positive, bottom-right), how many healthy patients were wrongly flagged (false positive, top-right), and how many sick patients were missed (false negative, bottom-left).
- Run the cell.
- You should see two blue heatmap grids side-by-side, one labeled Logistic Regression and one labeled Random Forest.
Why are false negatives more dangerous here?
In medical diagnosis, a false negative means a patient with heart disease is told they are healthy. That missed diagnosis could delay life-saving treatment.
A false positive (healthy patient told they might be sick) leads to extra tests, which is inconvenient but not dangerous. When evaluating your matrices, pay close attention to the bottom-left cell.
Don't see the confusion matrices?
Make sure you ran all previous cells first. The variables lr_predictions and rf_predictions must exist from Cells 9 and 10.
If you see a NameError, click Runtime then Run all to re-run every cell from the top.
Help me debug my confusion matrix code
Save the best model with joblib
Your trained model only lives in memory right now. If you close the notebook, it disappears. Saving the model to a file means you can load it later in a completely different application, like an API server.
- Create a new code cell and paste the following code to save the best-performing pipeline:
# --- Cell 13: Save Best Model ---
best_pipeline = lr_pipeline
joblib.dump(best_pipeline, 'heart_disease_model.joblib')
print("Model saved as 'heart_disease_model.joblib'")What does this code do?
best_pipeline = lr_pipeline assigns the Logistic Regression pipeline as the model to save. Based on your cross-validation results, Logistic Regression tends to generalize more consistently on this small dataset.
joblib.dump() serializes the entire pipeline (the StandardScaler and the trained LogisticRegression) into a single .joblib file. This file can be loaded in any Python environment to make predictions without retraining.
Why joblib instead of pickle?
joblib is optimized for objects containing large numpy arrays, which is what scikit-learn models store internally. It is the recommended serialization approach in the scikit-learn documentation.
- Run the cell.
- You should see the message Model saved as 'heart_disease_model.joblib' printed in the output.
- Click the folder icon in the left sidebar of Colab to open the file browser. You should see heart_disease_model.joblib listed there.
Don't see the file in the sidebar?
Click the refresh button (circular arrow icon) at the top of the file browser panel. Colab sometimes needs a manual refresh to show newly created files.
Help me find my saved model file in Colab
Test the saved model by loading and predicting
A saved model is only useful if you can load it back and get predictions. This final test simulates what a production API would do: load the model file, receive patient data, and return a prediction.
- Create a new code cell and paste the following code to load the model and predict on a sample patient:
# --- Cell 14: Test Saved Model ---
loaded_model = joblib.load('heart_disease_model.joblib')
sample_patient = X_test.iloc[0:1]
prediction = loaded_model.predict(sample_patient)
probability = loaded_model.predict_proba(sample_patient)
print(f"Sample patient features:")
print(sample_patient.to_dict('records')[0])
print(f"\nPrediction: {'Disease Present' if prediction[0] == 1 else 'No Disease'}")
print(f"Probability of disease: {probability[0][1]:.4f}")What does this code do?
joblib.load() reads the saved file back into memory as a fully functional pipeline. No retraining needed.
X_test.iloc[0:1] grabs the first patient from the test set as a DataFrame row. predict() returns a binary label (0 or 1). predict_proba() returns a confidence score between 0 and 1, showing how certain the model is about its prediction.
- Run the cell.
- You should see the patient's feature values printed, followed by a prediction (either Disease Present or No Disease) and a probability score.
Seeing a FileNotFoundError?
This means Cell 13 did not run successfully. Go back and run Cell 13 first to create the heart_disease_model.joblib file, then re-run Cell 14.
Help me fix my model loading error
✔️ Awesome, I've got everything!
Great. Double check you've run all cells in order from Cell 1 through Cell 14.
ⓧ I'd like to double check the full code
Here is the complete notebook code from Cell 1 through Cell 14. Compare your cells against this:
# --- Cell 1: Import Libraries ---
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report, ConfusionMatrixDisplay
from sklearn.pipeline import Pipeline
import joblib
print("All libraries imported successfully!")
print(f"scikit-learn version: {__import__('sklearn').__version__}")
print(f"pandas version: {pd.__version__}")
# --- Cell 2: Load Dataset ---
url = "https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/processed.cleveland.data"
column_names = [
'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs',
'restecg', 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'
]
df = pd.read_csv(url, names=column_names, na_values='?')
print(f"Dataset shape: {df.shape}")
print(f"\nFirst 5 rows:")
print(df.head())
print(f"\nData types:")
print(df.dtypes)
# --- Cell 3: Initial Exploration ---
print(f"Missing values per column:")
print(df.isnull().sum())
print(f"\nTarget distribution (original):")
print(df['target'].value_counts())
# --- Cell 4: Binarize Target ---
df['target'] = df['target'].apply(lambda x: 1 if x > 0 else 0)
print(f"Binary target distribution:")
print(df['target'].value_counts())
print(f"\nDisease present: {df['target'].sum()}")
print(f"No disease: {(df['target'] == 0).sum()}")
# --- Cell 5: EDA - Distribution Plots ---
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
sns.histplot(data=df, x='age', hue='target', kde=True, ax=axes[0, 0])
axes[0, 0].set_title('Age Distribution by Target')
sns.histplot(data=df, x='trestbps', hue='target', kde=True, ax=axes[0, 1])
axes[0, 1].set_title('Resting Blood Pressure by Target')
sns.histplot(data=df, x='chol', hue='target', kde=True, ax=axes[0, 2])
axes[0, 2].set_title('Cholesterol by Target')
sns.histplot(data=df, x='thalach', hue='target', kde=True, ax=axes[1, 0])
axes[1, 0].set_title('Max Heart Rate by Target')
sns.histplot(data=df, x='oldpeak', hue='target', kde=True, ax=axes[1, 1])
axes[1, 1].set_title('ST Depression by Target')
sns.countplot(data=df, x='cp', hue='target', ax=axes[1, 2])
axes[1, 2].set_title('Chest Pain Type by Target')
plt.tight_layout()
plt.show()
# --- Cell 6: EDA - Correlation Heatmap ---
plt.figure(figsize=(12, 8))
sns.heatmap(df.corr(numeric_only=True), annot=True, cmap='coolwarm', center=0, fmt='.2f')
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.show()
# --- Cell 7: Handle Missing Values ---
print(f"Missing values before cleaning:")
print(df.isnull().sum()[df.isnull().sum() > 0])
df['ca'] = df['ca'].fillna(df['ca'].median())
df['thal'] = df['thal'].fillna(df['thal'].median())
print(f"\nTotal missing values after cleaning: {df.isnull().sum().sum()}")
# --- Cell 8: Train/Test Split ---
X = df.drop('target', axis=1)
y = df['target']
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"Training set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")
print(f"\nTraining target distribution:")
print(y_train.value_counts())
# --- Cell 9: Train Logistic Regression ---
lr_pipeline = Pipeline([
('scaler', StandardScaler()),
('classifier', LogisticRegression(random_state=42, max_iter=1000))
])
lr_pipeline.fit(X_train, y_train)
lr_predictions = lr_pipeline.predict(X_test)
print("Logistic Regression Results:")
print(f"Accuracy: {accuracy_score(y_test, lr_predictions):.4f}")
print(f"\nClassification Report:")
print(classification_report(y_test, lr_predictions, target_names=['No Disease', 'Disease']))
# --- Cell 10: Train Random Forest ---
rf_pipeline = Pipeline([
('scaler', StandardScaler()),
('classifier', RandomForestClassifier(n_estimators=100, random_state=42))
])
rf_pipeline.fit(X_train, y_train)
rf_predictions = rf_pipeline.predict(X_test)
print("Random Forest Results:")
print(f"Accuracy: {accuracy_score(y_test, rf_predictions):.4f}")
print(f"\nClassification Report:")
print(classification_report(y_test, rf_predictions, target_names=['No Disease', 'Disease']))
# --- Cell 11: Cross-Validation Comparison ---
lr_cv_scores = cross_val_score(lr_pipeline, X, y, cv=5, scoring='accuracy')
rf_cv_scores = cross_val_score(rf_pipeline, X, y, cv=5, scoring='accuracy')
print("5-Fold Cross-Validation Results:")
print(f"\nLogistic Regression: {lr_cv_scores.mean():.4f} (+/- {lr_cv_scores.std():.4f})")
print(f"Random Forest: {rf_cv_scores.mean():.4f} (+/- {rf_cv_scores.std():.4f})")
# --- Cell 12: Confusion Matrix Visualization ---
fig, axes = plt.subplots(1, 2, figsize=(12, 5))
ConfusionMatrixDisplay.from_predictions(
y_test, lr_predictions,
display_labels=['No Disease', 'Disease'],
ax=axes[0],
cmap='Blues'
)
axes[0].set_title('Logistic Regression')
ConfusionMatrixDisplay.from_predictions(
y_test, rf_predictions,
display_labels=['No Disease', 'Disease'],
ax=axes[1],
cmap='Blues'
)
axes[1].set_title('Random Forest')
plt.tight_layout()
plt.show()
# --- Cell 13: Save Best Model ---
best_pipeline = lr_pipeline
joblib.dump(best_pipeline, 'heart_disease_model.joblib')
print("Model saved as 'heart_disease_model.joblib'")
# --- Cell 14: Test Saved Model ---
loaded_model = joblib.load('heart_disease_model.joblib')
sample_patient = X_test.iloc[0:1]
prediction = loaded_model.predict(sample_patient)
probability = loaded_model.predict_proba(sample_patient)
print(f"Sample patient features:")
print(sample_patient.to_dict('records')[0])
print(f"\nPrediction: {'Disease Present' if prediction[0] == 1 else 'No Disease'}")
print(f"Probability of disease: {probability[0][1]:.4f}")Secret mission
Hyperparameter Tuning with GridSearchCV
Your Random Forest used default hyperparameters. Can you systematically search for the best combination of settings to squeeze out better accuracy? This is the technique that separates junior from intermediate data scientists.
Clean Up Your Resources
Clean Up Your Resources
This project runs entirely in Google Colab, so there are no ongoing costs. Decide whether to keep your notebook and model files, pause the runtime to free up compute resources, or delete everything entirely.
Resources you used:
- Google Colab notebook ("Heart Disease Prediction") stored in your Google Drive.
- Saved model files (heart_disease_model.joblib and heart_disease_model_tuned.joblib) in the Colab runtime.
Keep everything running
No action needed. Choose this if you plan to continue building on this notebook or want to reference it later.
- Your notebook stays in Google Drive under the Colab Notebooks folder.
- The .joblib model files persist as long as your Colab runtime stays connected. If you want to keep them permanently, download them to your computer before disconnecting.
Pause - I'll come back to this later
Disconnect the Colab runtime to free up compute resources. Your notebook code stays saved in Google Drive, but the runtime's temporary files (including your model files) will be deleted.
- Download your model files first if you want to keep them. Click the folder icon in the left sidebar, right-click each .joblib file, and select Download.
- Click Runtime in the top menu bar.
- Select Disconnect and delete runtime.
- Your notebook code is still saved in Google Drive. Next time you open it, click Runtime then Run all to re-execute all cells.
Delete - I don't want to use this again
Remove the notebook and all associated runtime files completely.
- In Google Colab, click File in the top menu bar.
- Select Move to trash. This moves the notebook to your Google Drive trash.
- To permanently delete, open Google Drive, navigate to Trash in the left sidebar, right-click the notebook, and select Delete forever.
- The .joblib model files are automatically deleted when the runtime disconnects. No extra cleanup needed for those.
Nice Work!
Nice Work!
You just built a complete machine learning pipeline that predicts heart disease from real clinical data. From raw patient records to a saved, production-ready model, you now have the end-to-end workflow that data science interviewers look for.
You've learned how to:
- Build a complete ML pipeline from loading and exploring the UCI Heart Disease dataset through preprocessing, training, evaluation, and model export.
- Train and compare Logistic Regression and Random Forest classifiers using scikit-learn Pipelines and 5-fold cross-validation to select the best model.
- Evaluate model performance with confusion matrices and classification reports, then export the trained pipeline as a .joblib file ready for API integration.
- Secret Mission: Tuned hyperparameters with GridSearchCV to systematically search for the best Random Forest configuration and squeeze out better performance.
Ready to quiz yourself?