SmolLM-135M_MedicalQA-SFT

Github: https://github.com/Aadit032

Medical-domain adaptation pipeline for Aadit-032/SmolLM-135M_MedicalQA-SFT. The project runs a two-stage training workflow — continued pre-training (CPT) on biomedical corpora, then supervised fine-tuning (SFT) on medical reasoning Q&A — with comprehensive evaluation at each stage.

---

Overview

Both stages use LoRA adapters via Unsloth with 4-bit (NF4) loading and 16-bit merged export.

---

Pipelines

CPT pipeline (cpt.py)

Load base model → baseline evals → CPT training → merge & save → post-training evals

1. Load SmolLM-135M in 4-bit via Unsloth
2. Run baseline evals (perplexity, medical benchmarks, generation, lm-eval)
3. Download & tokenize 200K biomedical samples (cpt_data.py)
4. Train LoRA adapters for 1 epoch (cpt_train.py)
5. Merge and save to SmolLM-135M_Med_Merged/
6. Re-run all evals and compare before/after

SFT pipeline (sft.py)

Load CPT model → pre-SFT evals → SFT training → merge & save → post-SFT evals

1. Load Aadit-032/SmolLM-135M_Med-CPT (post-CPT checkpoint)
2. Run pre-SFT evals on the CPT model
3. Load SFT dataset from sft_data.py (medical-o1 by default)
4. Train LoRA adapters with response-only loss masking (sft_train.py)
5. Merge and save to SmolLM-135M_Med-SFT-Merged/ (also hosted on Hugging Face)
6. Re-run all evals and compare pre-SFT vs post-SFT

---

Usage

# Install dependencies
uv sync

# --- CPT ---
uv run cpt.py              # full CPT pipeline
uv run cpt_train.py        # CPT training only
uv run cpt_data.py         # download & prepare CPT data only

# --- SFT ---
uv run sft_data.py         # build/load SFT data (defaults to medical_o1)
uv run sft.py              # full SFT pipeline
uv run sft_train.py        # SFT training only

SFT data options

# Default: medical-o1 reasoning dataset (cached after first run)
uv run sft_data.py --loader medical_o1

# Synthetic: CPT split → semantic chunks → OpenRouter Q&A generation
export OPENROUTER_API_KEY="your-key"
uv run sft_data.py --loader synthetic --max-train-chunks 50

# Rebuild from scratch (ignore cached JSONL)
uv run sft_data.py --rebuild

# Choose output format: chat (default), alpaca, or raw
uv run sft_data.py --format alpaca

---

Configuration

All settings live in config.yaml:

---

Training Details

CPT (cpt_train.py)

LoRA

Hyperparameters

CPT data (200,000 samples, 90/10 train/val split)

SFT (sft_train.py)

LoRA

Hyperparameters

SFT data loaders (sft_data.py)

Both loaders append 20 general instruction pairs (math, science, geography, refusal examples) to prevent catastrophic forgetting.

SFT data is exported in three formats: alpaca, raw (Question: / Answer:), and chat (ChatML-style).

---

Evaluation

All results are saved to ./results/ as JSON.

Medical benchmarks use single-pass log-prob scoring — one forward pass per question by batching all answer choices together.

---

Project Structure

├── cpt.py               # CPT pipeline entry point
├── cpt_train.py         # CPT LoRA training
├── cpt_data.py          # CPT dataset download & preprocessing
├── sft.py               # SFT pipeline entry point
├── sft_train.py         # SFT LoRA training
├── sft_data.py          # SFT dataset loaders (medical_o1 / synthetic)
├── model_utils.py       # Model loading (base + SFT checkpoint)
├── config.yaml          # All configuration
├── pyproject.toml       # Dependencies
├── evals/
│   ├── benchmarks.py    # PubMedQA & MedMCQA
│   ├── perplexity.py    # Sliding-window perplexity
│   ├── generation.py    # General + medical generation eval
│   └── lm_eval.py       # General capability benchmarks
├── data/
│   ├── train.txt        # CPT training data
│   ├── val.txt          # CPT validation data
│   └── sft/             # Cached SFT datasets (JSONL)
└── results/             # Evaluation outputs

---

Dependencies

• Python >= 3.13
• unsloth
• datasets
• omegaconf
• evaluate
• lm-eval

Requires a CUDA GPU. For the synthetic SFT loader, set OPENROUTER_API_KEY in your environment.

---

CPT Results

CPT achieved its primary objective — domain adaptation — with ~20% perplexity reduction on biomedical text. Downstream medical QA did not improve proportionally, motivating the SFT stage on medical reasoning data.

---

SFT Results

Evaluated on the CPT checkpoint (pre-SFT) vs the merged SFT model (post-SFT). Full outputs are in ./results/.

Model: Aadit-032/SmolLM-135M_MedicalQA-SFT

SFT recovered PubMedQA accuracy toward the untrained baseline (49.5%) while keeping biomedical perplexity stable. MedMCQA was essentially unchanged — a harder 4-option MCQ benchmark that likely needs more targeted training data or longer fine-tuning.

General capability (lm-eval, post-SFT)

Generation samples (generation_pre_sft.json vs generation_sft.json) show modest gains in instruction-following structure, but outputs remain repetitive at 135M scale — expected for a model this size without RLHF or larger SFT corpora.