Your Personal AI Cardiologist? Fine-Tuning Phi-3 for Heart Disease Prediction

Large Language Models (LLMs) like ChatGPT are incredible at writing poems and summarizing articles. But can they do something more... critical? Like, say, help predict the risk of heart disease? Let's find out!

Cardiovascular diseases are the leading cause of death worldwide. Early detection is a game-changer, and this is where machine learning can lend a hand. We're going to take a powerful, general-purpose LLM and teach it a very specific skill: how to analyze patient data and flag potential heart disease risks.

In this post, we'll roll up our sleeves and walk through fine-tuning the lightweight yet powerful Phi-3-mini model on a real-world medical dataset.

Fine-tuned model is here.

Here's what you'll learn how to do:

• Load a model efficiently using 4-bit quantization (it's easier than it sounds!).
• Set up Low-Rank Adapters (LoRA) to fine-tune the model without needing a supercomputer.
• Load and format the UCI Heart Disease dataset.
• Fine-tune the model using the slick SFTTrainer from the Hugging Face trl library.
• Test our newly specialized model with a sample patient profile.

Ready to play doctor with an AI? Let's get started!

The Setup

First things first, we need to install a few libraries. If you're running this in a Google Colab notebook, just pop these commands into a cell. We're using specific versions to make sure everything plays nicely together.

# We need a few key libraries for this adventure
!pip install transformers==4.46.2 peft==0.13.2 accelerate==1.1.1 trl==0.12.1 bitsandbytes==0.45.2
!pip install datasets==3.1.0 pandas==2.2.2 ucimlrepo==0.0.6

All Aboard! Importing Our Tools

To keep our code clean, let's get all our imports out of the way at the start.

import os
import torch
import pandas as pd
from datasets import Dataset, DatasetDict
from ucimlrepo import fetch_ucirepo
from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl import SFTConfig, SFTTrainer

Loading the Brains: Our Quantized Base Model

We'll be using Microsoft's Phi-3-mini-4k-instruct, a fantastic open model that packs a punch without demanding a ton of resources. But even a "mini" model can be hefty. To make it fit comfortably on a standard GPU, we'll load it in a "quantized" 4-bit format.

Think of quantization as compressing the model's brain. We're taking its knowledge (the model weights) and representing it with less precision. This dramatically reduces memory usage without losing too much of the model's smarts. The BitsAndBytes library makes this a piece of cake.

# Configure the bits and bytes for 4-bit quantization
bnb_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_quant_type="nf4",
   bnb_4bit_use_double_quant=True,
   bnb_4bit_compute_dtype=torch.float32
)

# The model we're going to teach
repo_id = 'microsoft/Phi-3-mini-4k-instruct'

# Load the model with our quantization config
model = AutoModelForCausalLM.from_pretrained(
   repo_id,
   device_map="auto", # Automatically selects the GPU
   quantization_config=bnb_config
)

Even after this compression, the model still takes up over 2GB of memory. But here's the catch: a quantized model is like a read-only book. We can get information from it, but we can't write new information in (i.e., we can't train it).

To solve this, we need to add a sprinkle of adapters.

Setting Up LoRA: The Smart Way to Fine-Tune

This is where the magic of LoRA (Low-Rank Adaptation) comes in. Instead of trying to update the entire frozen model, we attach tiny, trainable "adapter" layers to it. We only train these small adapters, which represent a tiny fraction of the total model size. It's an incredibly efficient way to teach an old model new tricks.

Setting up LoRA is a simple three-step process:

1. Prep the model for training.
2. Define a LoRA configuration.
3. Apply the config to our model.

# 1. Prepare the model for k-bit training
model = prepare_model_for_kbit_training(model)

# 2. Configure our LoRA adapters
config = LoraConfig(
    r=8,                   # The rank of the adapter. Lower is smaller.
    lora_alpha=16,         # A scaling factor, usually 2*r
    bias="none",
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
    # We need to tell it which layers to attach the adapters to
    target_modules=['o_proj', 'qkv_proj', 'gate_up_proj', 'down_proj'],
)

# 3. Wrap the model with our LoRA config
model = get_peft_model(model, config)

Now, let's see how many parameters we're actually going to train.

train_p, tot_p = model.get_nb_trainable_parameters()
print(f'Trainable parameters:      {train_p/1e6:.2f}M')
print(f'Total parameters:          {tot_p/1e6:.2f}M')
print(f'% of trainable parameters: {100*train_p/tot_p:.2f}%')

Output:

Trainable parameters:      12.58M
Total parameters:          3833.66M
% of trainable parameters: 0.33%

Look at that! We only have to train 0.33% of the total parameters. Our model is now ready for its medical school training. All we need is the textbook—our dataset.

The Curriculum: Formatting Our Heart Disease Dataset

Instead of teaching our model to speak like Yoda, we're giving it a much more serious task. We'll use the UCI Heart Disease dataset, which is a collection of data from 918 patients. It includes 11 clinical features like age, cholesterol, and chest pain type, plus a final diagnosis.

Our goal is to turn each row of this structured data into a natural language question-and-answer pair that the LLM can understand.

First, let's fetch the data using the ucimlrepo library and load it into a pandas DataFrame.

# Fetch dataset from UCI
heart_disease = fetch_ucirepo(id=45)

# Get the features (X) and target (y)
X = heart_disease.data.features
y = heart_disease.data.targets

# Combine them into a single DataFrame
df = pd.concat([X, y], axis=1)

# The target column is named 'num', let's rename it to something clearer
df.rename(columns={'num': 'HeartDisease'}, inplace=True)
# The original dataset uses values 0-4 for the target. We only need a binary classification: 0 (no disease) vs >0 (disease).
df['HeartDisease'] = df['HeartDisease'].apply(lambda x: 1 if x > 0 else 0)

# Let's check our data
print(df.head())

Now for the fun part. We'll create a function that turns each patient's data into a human-readable prompt and a simple completion.

def create_prompt_completion(row):
    # Map coded values to human-readable strings for clarity in the prompt
    sex_map = {'M': 'male', 'F': 'female'}
    chest_pain_map = {'TA': 'Typical Angina', 'ATA': 'Atypical Angina', 'NAP': 'Non-Anginal Pain', 'ASY': 'Asymptomatic'}
    resting_ecg_map = {'Normal': 'Normal', 'ST': 'ST-T wave abnormality', 'LVH': 'probable left ventricular hypertrophy'}
    st_slope_map = {'Up': 'upsloping', 'Flat': 'flat', 'Down': 'downsloping'}
    exercise_angina_map = {'Y': 'Yes', 'N': 'No'}

    # Determine the human-readable string for FastingBS
    fasting_bs_status = "high (greater than 120 mg/dl)" if row['FastingBS'] == 1 else "normal (120 mg/dl or less)"

    prompt = (
        f"A {row['Age']} year old {sex_map[row['Sex']]} has chest pain of type '{chest_pain_map[row['ChestPainType']]}', "
        f"a resting blood pressure of {row['RestingBP']} mmHg, and cholesterol levels of {row['Cholesterol']} mg/dl. "
        f"Their fasting blood sugar is {fasting_bs_status}. "
        f"Resting ECG results are '{resting_ecg_map[row['RestingECG']]}'. "
        f"The maximum heart rate achieved is {row['MaxHR']}. "
        f"Exercise induced angina is '{exercise_angina_map[row['ExerciseAngina']]}'. "
        f"The oldpeak value is {row['Oldpeak']}. "
        f"The slope of the peak exercise ST segment is '{st_slope_map[row['ST_Slope']]}'. "
        f"Does this individual likely have heart disease?"
    )
    completion = 'Yes, could be Heart disease' if row['HeartDisease'] == 1 else 'No, likely not Heart disease'
    return {'prompt': prompt, 'completion': completion}

# Apply the function to create our prompt/completion pairs
formatted_df = df.apply(create_prompt_completion, axis=1, result_type='expand')

# Convert the pandas DataFrame to a Hugging Face Dataset
dataset = Dataset.from_pandas(formatted_df)

# Let's see an example!
print(dataset[0])

This gives us a prompt like this:

'prompt': "A 40 year old male has chest pain of type 'Atypical Angina', a resting blood pressure of 140 mmHg, and cholesterol levels of 289 mg/dl. Their fasting blood sugar is normal (120 mg/dl or less). Resting ECG results are 'Normal'. The maximum heart rate achieved is 172. Exercise induced angina is 'No'. The oldpeak value is 0.0. The slope of the peak exercise ST segment is 'upsloping'. Does this individual likely have heart disease?"
'completion': 'No, likely not Heart disease'

Perfect! The SFTTrainer works best with a "conversational" format, so we'll quickly map our prompt and completion columns into the right structure.

def format_dataset(examples):
    converted_sample = [
        {"role": "user", "content": examples["prompt"]},
        {"role": "assistant", "content": examples["completion"]},
    ]
    return {'messages': converted_sample}

dataset = dataset.map(format_dataset, remove_columns=['prompt', 'completion'])
print(dataset[0]['messages'])

Loading the Tokenizer

Before we train, we need the model's tokenizer. The tokenizer converts our text prompts into a sequence of numbers (tokens) that the model can process. For instruction-tuned models like Phi-3, it also contains a chat_template that wraps our conversation with special tokens so the model knows who is talking (`