SmolLM-135M_Med-CPT

Github: https://github.com/Aadit032/SmolLM-135M_Med

Medical-domain continued pre-training (CPT) pipeline for HuggingFaceTB's SmolLM-135M model. Trains on PubMed, PMC, Medline, and FineWeb datasets using LoRA adapters with the Unsloth framework.

---

Pipeline

┌──────────────────┐     ┌──────────────────┐
│  1. LOAD MODEL   │ ──► │  2. BASELINE     │
│  (4-bit NF4)     │     │  EVAL (untrained)│
└──────────────────┘     └──────────────────┘
                                 │
                                 ▼
                          ┌──────────────────┐
                          │  Perplexity      │
                          │  + Benchmarks    │
                          └──────────────────┘
                                 │
                                 ▼
┌──────────────────┐     ┌──────────────────┐
│  3. TRAIN        │ ◄── │  4. DATA PHASE   │
│  (LoRA + CPT)    │     │  (inside train)  │
└──────────────────┘     └──────────────────┘
        │
        ▼
┌──────────────────┐     ┌──────────────────┐
│  5. SAVE         │ ──► │  6. POST-TRAIN   │
│  Merged 16-bit   │     │  EVAL (trained)  │
└──────────────────┘     └──────────────────┘
                                 │
                                 ▼
                          ┌──────────────────┐
                          │  Perplexity       │
                          │  + Benchmarks     │
                          └──────────────────┘

Step-by-Step

1. Load Model (model_utils.py) — Loads SmolLM-135M in 4-bit (NF4) via Unsloth with bfloat16 compute dtype.

2. Baseline Eval — Evaluates the untrained base model on perplexity (PubMed Abstracts, Medline) and benchmarks (PubMedQA, MedMCQA).

3. Data Phase (data.py) — Downloads and tokenizes biomedical datasets (PubMed, PMC, Medline, FineWeb; 200K samples), splits 90/10 train/val, writes to text files, loads into HuggingFace Datasets. Called inside train.py.

4. Training (train.py) — Loads base model, attaches LoRA adapters (rank 32), tokenizes datasets into packed sequences, runs 1 epoch of CPT with UnslothTrainer.

5. Export — Saves a merged 16-bit model to SmolLM-135M_Med_Merged/.

6. Post-Training Eval — Re-runs perplexity and benchmarks on the trained model.

---

Configuration

All paths and hyperparameters are set in config.yaml:

---

Training Details

Model

• Base: HuggingFaceTB/SmolLM-135M (135M parameters)
• Precision: 4-bit loaded (NF4), merged to 16-bit on save
• Sequence length: 512 tokens

LoRA Configuration

Training Hyperparameters

Datasets (200,000 total samples)

• Split: 90% train / 10% validation (per-source random split)
• Pre-tokenization: Tokenized with SmolLM-135M tokenizer to count tokens

Evaluation

• Perplexity: Sliding-window (window=512, stride=256) on PubMed Abstracts and Medline (1,000 samples each)
• Benchmarks: PubMedQA (yes/no/maybe) and MedMCQA (4-option MCQ) — 200 samples each
• Results: Saved to ./results/ as JSON with _untrained and _trained suffixes

---

Usage

# Run the full pipeline
uv run main.py

Dependencies

• Python >= 3.13
• unsloth
• datasets
• omegaconf
• evaluate

Install with uv:

uv sync

---

Project Structure

├── main.py              # Pipeline entry point
├── train.py             # LoRA training with UnslothTrainer
├── data.py              # Dataset download & preprocessing
├── model_utils.py       # Shared model loading & config
├── config.yaml          # Configuration (paths & hyperparams)
├── pyproject.toml       # Project metadata & dependencies
├── evals/
│   ├── __init__.py
│   ├── benchmarks.py    # PubMedQA & MedMCQA evaluation
│   └── perplexity.py    # Sliding-window perplexity
└── results/             # Evaluation outputs (JSON: _untrained / _trained)