CPT (Continual Pre-Training)#
Few months ago I was struck with a client discussion who was interchangeably using CPT and MPT (model post training). I was imagining Low Rank Adaptation (LoRA) and PEFT (Parameter Efficient Fine-Tuning) during the discussion to finetune the modle with additional data, but the clients data corpus was more than 1T tokens - and LoRA and PEFT might not be the best approach…
This led me to explore CPT… Just organizing my notes here for future me…
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import AdamW
from transformers import GPT2Config, GPT2LMHeadModel, GPT2Tokenizer, get_linear_schedule_with_warmup
import random
import numpy as np
from torch.utils.data import Dataset, DataLoader
import re
/Users/n0man/Code/n03an.me/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm
WARNING:torchao.kernel.intmm:Warning: Detected no triton, on systems without Triton certain kernels will not work
Medical Dataset data model for CPT.#
class MedicalDataset(Dataset):
def __init__(self, corpus, tokenizer, max_length=128):
self.corpus = corpus
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.corpus)
def __getitem__(self, idx):
text = self.corpus[idx]
encoding = self.tokenizer(
text,
truncation=True,
max_length=self.max_length,
padding="max_length",
return_tensors="pt"
)
return {
'input_ids': encoding['input_ids'].squeeze(),
'attention_mask': encoding['attention_mask'].squeeze()
}
Synthetic Data generation#
We’ll hand build medical data corpus mimicking EMR data (using domain specific models like LLaMA-Med could be another option - but Hey, we are doing this from scratch :)
import random
from faker import Faker
fake = Faker()
medical_contexts = [
{
"symptoms": ["fever", "cough", "shortness of breath", "chest pain"],
"diagnosis": "pneumonia",
"findings": ["consolidation in the right lower lobe", "pleural effusion", "lung infiltrates"],
"treatments": ["antibiotics", "oxygen therapy", "chest physiotherapy"]
},
{
"symptoms": ["acute chest pain", "shortness of breath", "tachycardia"],
"diagnosis": "pulmonary embolism",
"findings": ["filling defect in the pulmonary artery", "right heart strain", "ventilation-perfusion mismatch"],
"treatments": ["anticoagulant therapy", "thrombolysis", "inferior vena cava filter"]
},
{
"symptoms": ["right lower quadrant pain", "nausea", "vomiting", "fever"],
"diagnosis": "acute appendicitis",
"findings": ["appendix dilation >6mm", "periappendiceal fat stranding", "appendicolith"],
"treatments": ["laparoscopic appendectomy", "IV antibiotics", "NPO with IV fluids"]
},
{
"symptoms": ["unilateral headache", "photophobia", "nausea", "visual aura"],
"diagnosis": "migraine",
"findings": ["normal brain imaging", "no intracranial abnormalities"],
"treatments": ["triptans", "NSAIDs", "anti-emetics", "preventive therapy"]
},
{
"symptoms": ["polyuria", "polydipsia", "fatigue", "blurred vision"],
"diagnosis": "diabetes mellitus",
"findings": ["elevated fasting glucose >126 mg/dL", "HbA1c > 6.5%", "glucosuria"],
"treatments": ["metformin", "insulin therapy", "dietary modification", "exercise program"]
},
{
"symptoms": ["substernal chest pain", "dyspnea on exertion", "diaphoresis"],
"diagnosis": "acute coronary syndrome",
"findings": ["ST-elevation on ECG", "elevated troponin", "regional wall motion abnormality"],
"treatments": ["aspirin", "nitroglycerin", "PCI", "beta-blockers"]
},
{
"symptoms": ["hemoptysis", "weight loss", "night sweats", "persistent cough"],
"diagnosis": "lung cancer",
"findings": ["lung mass on CT", "hilar lymphadenopathy", "positive biopsy for adenocarcinoma"],
"treatments": ["chemotherapy", "radiation therapy", "surgical resection", "immunotherapy"]
},
{
"symptoms": ["jaundice", "right upper quadrant pain", "clay-colored stools"],
"diagnosis": "cholecystitis",
"findings": ["gallbladder wall thickening", "pericholecystic fluid", "positive Murphy's sign"],
"treatments": ["laparoscopic cholecystectomy", "IV antibiotics", "NPO status"]
},
{
"symptoms": ["hematemesis", "melena", "orthostatic hypotension"],
"diagnosis": "peptic ulcer disease",
"findings": ["ulcer crater on endoscopy", "Helicobacter pylori positive", "anemia on CBC"],
"treatments": ["PPI therapy", "H. pylori eradication", "blood transfusion", "endoscopic hemostasis"]
},
{
"symptoms": ["lower extremity edema", "orthopnea", "paroxysmal nocturnal dyspnea"],
"diagnosis": "congestive heart failure",
"findings": ["cardiomegaly on CXR", "elevated BNP", "reduced ejection fraction"],
"treatments": ["diuretics", "ACE inhibitors", "beta-blockers", "salt restriction"]
},
{
"symptoms": ["arthralgia", "morning stiffness", "fatigue"],
"diagnosis": "rheumatoid arthritis",
"findings": ["positive rheumatoid factor", "elevated CRP", "joint erosions on X-ray"],
"treatments": ["methotrexate", "biologics", "NSAIDs", "corticosteroids"]
},
{
"symptoms": ["cognitive decline", "memory loss", "disorientation"],
"diagnosis": "Alzheimer's dementia",
"findings": ["cortical atrophy on MRI", "hippocampal volume loss", "positive amyloid PET"],
"treatments": ["cholinesterase inhibitors", "memantine", "cognitive therapy", "caregiver support"]
},
{
"symptoms": ["hematuria", "flank pain", "fever"],
"diagnosis": "pyelonephritis",
"findings": ["positive urine culture", "renal cortical defects on scan", "leukocytosis"],
"treatments": ["IV antibiotics", "fluid resuscitation", "analgesics", "urine culture follow-up"]
},
{
"symptoms": ["skin rash", "joint pain", "facial butterfly rash"],
"diagnosis": "systemic lupus erythematosus",
"findings": ["positive ANA", "anti-dsDNA antibodies", "low complement levels"],
"treatments": ["hydroxychloroquine", "corticosteroids", "immunosuppressants", "sun protection"]
},
{
"symptoms": ["tremor", "bradykinesia", "rigidity"],
"diagnosis": "Parkinson's disease",
"findings": ["dopamine transporter deficit on DaTscan", "response to levodopa", "resting tremor"],
"treatments": ["levodopa/carbidopa", "dopamine agonists", "physical therapy", "deep brain stimulation"]
}
]
presentation_templates = [
"Patient presented with {symptom1} and {symptom2}.",
"Patient shows {symptom1} along with {symptom2}.",
"Chief complaint: {symptom1} and {symptom2}.",
"Initial symptoms include {symptom1} and {symptom2}.",
"Clinical presentation: {symptom1} and {symptom2}.",
"Symptom onset with {symptom1} and subsequent development of {symptom2}.",
"Patient reports {symptom1} accompanied by {symptom2}.",
"Presenting features: {symptom1} and {symptom2}."
]
finding_templates = [
"CT scan revealed {finding}.",
"Radiology report shows {finding}.",
"Imaging findings: {finding}.",
"Diagnostic imaging demonstrated {finding}.",
"Radiological assessment identified {finding}.",
"CT findings include {finding}.",
"MRI revealed {finding}.",
"Ultrasound showed {finding}.",
"X-ray demonstrated {finding}."
]
diagnosis_templates = [
"Diagnosis: {diagnosis}.",
"Final diagnosis: {diagnosis}.",
"Diagnostic conclusion: {diagnosis}.",
"Clinical diagnosis: {diagnosis}.",
"Impression: {diagnosis}.",
"Working diagnosis: {diagnosis}.",
"Diagnostic assessment confirms {diagnosis}.",
"Diagnostic impression: {diagnosis}."
]
treatment_templates = [
"Treatment plan: {treatment}.",
"Recommended treatment: {treatment}.",
"Management strategy: {treatment}.",
"Therapeutic approach: {treatment}.",
"Prescribed therapy: {treatment}.",
"Treatment initiated: {treatment}.",
"Clinical management includes {treatment}.",
"Therapeutic regimen: {treatment}."
]
medical_corpus = []
for context in medical_contexts:
for _ in range(200):
# Select random symptoms (2-3 symptoms)
num_symptoms = random.randint(2, 3)
symptoms = random.sample(context["symptoms"], num_symptoms)
# Create symptom presentation
if num_symptoms == 2:
symptom_text = random.choice(presentation_templates).format(
symptom1=symptoms[0], symptom2=symptoms[1])
else:
symptom_text = f"Patient presented with {symptoms[0]}, {symptoms[1]}, and {symptoms[2]}."
# Add demographic details for realism
# demographic = f"{fake.random_element(['M', 'F'])}{random.randint(20, 80)}"
# name = fake.name()
# record = f"Case #{fake.random_int(1000,9999)} | {demographic} | Name: {name} | "
record = ""
# Build coherent record
record += (
symptom_text + " " +
random.choice(finding_templates).format(finding=random.choice(context["findings"])) + " " +
random.choice(diagnosis_templates).format(diagnosis=context["diagnosis"]) + " " +
random.choice(treatment_templates).format(treatment=random.choice(context["treatments"]))
)
medical_corpus.append(record)
admission_templates = [
"Patient admitted for management of {condition}.",
"Hospital admission required for {treatment} of {condition}.",
"Admitted to {unit} for {condition} management.",
"Inpatient admission initiated for {condition}.",
"Patient hospitalized for {condition} treatment."
]
discharge_templates = [
"Discharge summary: {condition} resolved with {treatment}. Follow-up in {time}.",
"Discharge plan: Continue {treatment} with PCP follow-up in {time}.",
"Discharge instructions: {treatment} regimen, follow-up in {time}.",
"Discharged in stable condition after {treatment} for {condition}.",
"Hospital course completed for {condition}. Discharge with {treatment}."
]
units = ["medical unit", "ICU", "surgical ward", "cardiology service", "oncology floor"]
timeframes = ["1 week", "2 weeks", "1 month", "3 days", "48 hours"]
for context in medical_contexts:
for _ in range(20):
condition = context["diagnosis"]
treatment = random.choice(context["treatments"])
medical_corpus.append(
random.choice(admission_templates).format(
condition=condition,
treatment=treatment,
unit=random.choice(units))
)
medical_corpus.append(
random.choice(discharge_templates).format(
condition=condition,
treatment=treatment,
time=random.choice(timeframes))
)
print(f"Total medical records generated: {len(medical_corpus)}")
print(f"Sample record:\n{random.choice(medical_corpus)}")
Total medical records generated: 3600
Sample record:
Patient reports weight loss accompanied by hemoptysis. Radiological assessment identified lung mass on CT. Clinical diagnosis: lung cancer. Therapeutic approach: chemotherapy.
Model initialization with pre-trained weights#
Trying distilgpt2 with 82M parameters (pretrained weights) i.e. the model weights already trained on decent internet corpus. BTW, just for fun we could also start gtp-2 from scratch where the weights are randomly initialized and the model knows nothing yet - and train the model on common crawl data ☺️
config = GPT2Config.from_pretrained("distilgpt2") # Gives the model architecture
model = GPT2LMHeadModel(config).to(DEVICE)
but perhaps some other day when I could get my own NVIDIA DGX Spark… 🙂 lets get back to CPT i.e. Continual Pre-Training or domain-adaptive pretraining or further pretraining on new domain
# Configuration
if torch.backends.mps.is_available() and torch.backends.mps.is_built():
DEVICE = "mps" # Apple Metal backend
else:
DEVICE = "cpu"
# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BATCH_SIZE = 4
LEARNING_RATE = 1e-5 # 0.00001
MAX_LENGTH = 256
model = GPT2LMHeadModel.from_pretrained("distilgpt2").to(DEVICE)
tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
tokenizer.pad_token = tokenizer.eos_token
print("Initializing Distil GPT-2-small model with pretrained weights...")
model = GPT2LMHeadModel.from_pretrained("distilgpt2").to(DEVICE)
optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)
# 82M parameters in distilgpt2 - smaller than GPT-2 small 124M
print(f"model parameters in distilgpt2: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")
print(f"Using {DEVICE} for training.")
Initializing Distil GPT-2-small model with pretrained weights...
model parameters in distilgpt2: 81912576
Using mps for training.
Inference function#
def generate_response(model, prompt, max_new_tokens=40):
inputs = tokenizer(
prompt,
return_tensors="pt",
truncation=True,
max_length=MAX_LENGTH,
padding="max_length"
).to(DEVICE)
outputs = model.generate(
inputs.input_ids,
attention_mask=inputs.attention_mask,
max_new_tokens=max_new_tokens,
do_sample=True,
top_k=30, # Reduced top_k for less randomness
temperature=0.5, # Lower temperature for more focused output
top_p=0.95,
no_repeat_ngram_size=3, # Prevent repeating 3-grams
pad_token_id=tokenizer.eos_token_id,
eos_token_id=tokenizer.eos_token_id,
repetition_penalty=1.2, # Increased repetition penalty
num_return_sequences=1
)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Should output very generic as the model has not seen specific medical records yet...
print("Generating a sample response...")
prompt = "Patient presented with"
print(generate_response(model, f"Prompt: {prompt} \nResponse:"))
print(generate_response(model, f"Prompt: Who is spiderman? \nResponse:"))
Generating a sample response...
Prompt: Patient presented with
Response:The patient was admitted to the hospital after being transferred from a local area for treatment. The following day, another nurse who received medical care in Parma had been discharged and is now undergoing an emergency department
Prompt: Who is spiderman?
Response:A. The name of the Spiderman has been changed from "Spiderman" to "Doctor Strange". It was originally based on a comic by Peter Parker, but now it's being used in films
🤦🏻♂️ you see how the responses are very generic indicating the model have decent ability to generate nonsense but appropriate english sentences (quite common for a foundational pre-trained model)
Building Medical dataset#
dataset = MedicalDataset(medical_corpus, tokenizer, MAX_LENGTH)
print("daataset corpus size: ", len(dataset))
print("First sample: ",tokenizer.decode(dataset[0]['input_ids'], skip_special_tokens=True))
print(f"First 10 tokens: {dataset[0]['input_ids'][:10]}")
print(f"First 10 tokens decoded: {tokenizer.decode(dataset[0]['input_ids'][:10], skip_special_tokens=True)}")
daataset corpus size: 3600
First sample: Patient shows chest pain along with fever. Radiology report shows pleural effusion. Impression: pneumonia. Clinical management includes chest physiotherapy.
First 10 tokens: tensor([12130, 1153, 2523, 7721, 2356, 1863, 351, 17372, 13, 5325])
First 10 tokens decoded: Patient shows chest pain along with fever. Rad
Creating Dataloader#
Single data point per iteration causes high variance in gradients resulting loss to fluctuate wildly (toggle high and low). Dataloader helps to a) feeds data in mini-batches, improving training stability and speed b) shuffles data in each epoch, improving generalization
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
print(f"DataLoader created with {len(dataset)} samples")
print(f"Dataloader have {len(dataloader)} mini batches, where each batch has {BATCH_SIZE} samples, , Learning rate: {LEARNING_RATE}.")
DataLoader created with 3600 samples
Dataloader have 900 mini batches, where each batch has 4 samples, , Learning rate: 1e-05.
Model Training#
# Set to training mode
model.train()
print("Model:", "training mode" if model.training else "eval mode")
Model: training mode
We’ll airdrop the dataset into small batches of 4 resulting into 900 batches. The model will go through the entire dataset 4 times (4 epochs)…
Take a note of labels=input_ids in the training loop. Continual pretraining is unsupervised — it just helps the model better understand the new domain (medical records) without any explicit labels. The model learns to predict the next token similar to the input sequence… i.e. it improves the preplexity of the model on the new domain (better fluency) - BUT It doesn’t tell if the model is responding correctly - for that we would need supervised fine-tuning (SFT) on a labeled dataset… Will cover that in separate post… ;)
import time
epoch = 4 # Number of times the model will see the entire dataset
global_step = 0
start_time = time.time()
for i in range(epoch):
epoch_loss = 0
# 900 mini-batches per epoch with each batch containing 4 samples
for batch in dataloader:
input_ids = batch['input_ids'].to(DEVICE)
attention_mask = batch['attention_mask'].to(DEVICE)
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=input_ids
)
loss = outputs.loss
perplexity = torch.exp(loss)
epoch_loss += loss.item()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
optimizer.zero_grad()
global_step += 1
end_time = time.time()
if global_step % 200 == 0:
print(f"Epoch: {i+1}, Step {global_step}/{len(dataloader)*epoch} | Loss: {loss.item():.4f} | Perplexity: {perplexity.item():.2f} | Time taken: {end_time - start_time:.2f}s")
start_time = end_time
`loss_type=None` was set in the config but it is unrecognised.Using the default loss: `ForCausalLMLoss`.
Epoch: 1, Step 200/3600 | Loss: 0.3569 | Perplexity: 1.43 | Time taken: 32.53s
Epoch: 1, Step 400/3600 | Loss: 0.1578 | Perplexity: 1.17 | Time taken: 28.47s
Epoch: 1, Step 600/3600 | Loss: 0.1494 | Perplexity: 1.16 | Time taken: 30.06s
Epoch: 1, Step 800/3600 | Loss: 0.1287 | Perplexity: 1.14 | Time taken: 29.58s
Epoch: 2, Step 1000/3600 | Loss: 0.1194 | Perplexity: 1.13 | Time taken: 31.07s
Epoch: 2, Step 1200/3600 | Loss: 0.1232 | Perplexity: 1.13 | Time taken: 30.13s
Epoch: 2, Step 1400/3600 | Loss: 0.1014 | Perplexity: 1.11 | Time taken: 31.92s
Epoch: 2, Step 1600/3600 | Loss: 0.1037 | Perplexity: 1.11 | Time taken: 41.85s
Epoch: 2, Step 1800/3600 | Loss: 0.1018 | Perplexity: 1.11 | Time taken: 79.67s
Epoch: 3, Step 2000/3600 | Loss: 0.0850 | Perplexity: 1.09 | Time taken: 63.74s
Epoch: 3, Step 2200/3600 | Loss: 0.0706 | Perplexity: 1.07 | Time taken: 42.49s
Epoch: 3, Step 2400/3600 | Loss: 0.0669 | Perplexity: 1.07 | Time taken: 38.55s
Epoch: 3, Step 2600/3600 | Loss: 0.0844 | Perplexity: 1.09 | Time taken: 36.72s
Epoch: 4, Step 2800/3600 | Loss: 0.0908 | Perplexity: 1.10 | Time taken: 35.38s
Epoch: 4, Step 3000/3600 | Loss: 0.0854 | Perplexity: 1.09 | Time taken: 34.67s
Epoch: 4, Step 3200/3600 | Loss: 0.0759 | Perplexity: 1.08 | Time taken: 37.39s
Epoch: 4, Step 3400/3600 | Loss: 0.0712 | Perplexity: 1.07 | Time taken: 36.42s
Epoch: 4, Step 3600/3600 | Loss: 0.0659 | Perplexity: 1.07 | Time taken: 35.50s
Testing the Model post training with medical reasoning prompts#
Given we are still on pre-training, we dont expect the model to be any sensible beyond being just a “medical” token tumbler…
model.eval() # Set model to evaluation mode after training
print("Model:", "training mode" if model.training else "eval mode")
with torch.no_grad():
prompt = "Patient presented with"
print("⭕️ Prompt: ", prompt)
print(f"💬 Response: {generate_response(model, prompt)}")
print(generate_response(model, f"⭕️ Prompt: Who is spiderman? \n💬 Response:"))
Model: eval mode
⭕️ Prompt: Patient presented with
💬 Response: Patient presented with
: shortness of breath, and fever. CT scan revealed pleural effusion on X-ray for pneumonia. Clinical diagnosis: lung cancer. Therapeutic approach includes radiation therapy.
⭕️ Prompt: Who is spiderman?
💬 Response:Asteroid regimen initiated for cholinesterase inhibitors. Recommended treatment: metformin/carbidopa, triptans or immunosuppressants.
So the model seems to be doing better on the trained corpus but notice how model reacted to “Who is Spiderman?” prompt… The model have lost its broken general knowledge it had before CPT - known as catastrophic forgetting (potentially due to overfittting on medical records and our dataset being too small and narrow). The responses are still very robotic i.e. it is able to generate medical tokens but would need further supervised fine-tuning (SFT) on a labeled dataset to improve the reasoning and fluency of the responses… perhaps in a different post… 😉
Lets try a few more reasoning prompts to see how the model responds…
clinical_prompts = [
"Patient presents with fever and cough. Diagnostic considerations:",
"CT scan shows pulmonary embolism. Recommended treatment:",
"Differential diagnosis for chest pain and shortness of breath:",
"Management strategy for acute appendicitis:",
"First-line antibiotic choice for community-acquired pneumonia:",
"Prognostic factors for pneumonia:",
"Potential complications of pulmonary embolism:",
"When to refer a pneumonia patient to ICU:"
]
for prompt in clinical_prompts:
print(f"⭕️ Prompt: '{prompt}'\n💬 Response: {generate_response(model, prompt)}\n")
⭕️ Prompt: 'Patient presents with fever and cough. Diagnostic considerations:'
💬 Response: Patient presents with fever and cough. Diagnostic considerations:
erythematosus . Radiological assessment identified lung infiltrates, pneumonia on arrival in the right lower lobe of blood culture management regimen. Clinical diagnosis: pyelonephritis eradication
⭕️ Prompt: 'CT scan shows pulmonary embolism. Recommended treatment:'
💬 Response: CT scan shows pulmonary embolism. Recommended treatment:
1 mg/dL of PCP follow-up in 1 week, with IV fluids per night for lung cancer management. Clinical regimen includes chemotherapy therapy. Therapeutic approach include radiation restriction and surgical
⭕️ Prompt: 'Differential diagnosis for chest pain and shortness of breath:'
💬 Response: Differential diagnosis for chest pain and shortness of breath:
. Diagnostic impression: pneumonia, right heart strain management . Therapeutic regimen includes oxygen therapy.
⭕️ Prompt: 'Management strategy for acute appendicitis:'
💬 Response: Management strategy for acute appendicitis:
. Follow-up in 2 weeks with NPO status, follow-ups on 1 month after IV antibiotics regimen. Clinical management includes laparoscopic cholecystectomy. Management plan is cort
⭕️ Prompt: 'First-line antibiotic choice for community-acquired pneumonia:'
💬 Response: First-line antibiotic choice for community-acquired pneumonia:
. Following instructions from the right doctor's office, followup in 1 month after antibiotics are initiated. Clinical management includes laparoscopic cholecystectomy. Treatment plan include IV fluids regimen
⭕️ Prompt: 'Prognostic factors for pneumonia:'
💬 Response: Prognostic factors for pneumonia:
. Radiological assessment identified lung infiltrates in the right lower lobe of PCP management, along with oxygen therapy initiated to ward off pulmonary embolism. Diagnosis: acute appendicitis syndrome
⭕️ Prompt: 'Potential complications of pulmonary embolism:'
💬 Response: Potential complications of pulmonary embolism:
. X-ray demonstrated antihistamines, anticoagulant therapy for lung cancer treatment. Impression : systemic lupus erythematosus. Therapeutic approach includes
⭕️ Prompt: 'When to refer a pneumonia patient to ICU:'
💬 Response: When to refer a pneumonia patient to ICU:
. X-ray demonstrated lung infiltrates on arrival in the management of pulmonary embolism. Diagnostic impression: congestive heart failure. Treatment plan includes thrombolysis, beta blocker