Medical Edge Model - Advanced Reasoning Training Pipeline

Complete training pipeline for first-aid reasoning model using SFT, reward modeling, and Socratic tree reasoning.

Quick Setup

1. Install Dependencies

pip install -r requirements-stable.txt

2. Configure Environment

Copy the example environment file and add your OpenAI API key:

cp .env.example .env
# Edit .env and add your OPENAI_API_KEY

3. Run Pipeline Sequentially

Execute scripts in order from the pipeline/ directory:

cd pipeline

# Step 1: Load raw dataset
python 1_Dataset_Loader.py

# Step 2: Clean and split dataset
python 2_Dataset_Cleaning.py

# Step 3: Generate teacher responses (requires OpenAI API key)
python 3_Teacher_Distillation.py --dataset all --limit 10  # Test with 10 samples first

# Step 4: Train SFT model
python 4_SFT_Training.py

# Step 5: Train reward model
python 5_Reward_Model.py --mode full

# Step 6: Run Socratic tree reasoning
python 6_Socratic_Tree_Reasoning.py --num_questions 3

# Step 7: Evaluate performance
python 7_Test_Performance.py

# Step 8: (Optional Visualizations)
python 8_Optional_Visuals.py

Pipeline Overview

1. Dataset Loader - Downloads and extracts raw first-aid Q&A data
2. Dataset Cleaning - Filters and splits data (train/test)
3. Teacher Distillation - Generates high-quality reasoning responses via GPT-4o-mini
4. SFT Training - Fine-tunes base model on distilled data
5. Reward Model - Trains preference model with Bradley-Terry loss
6. Socratic Reasoning - Implements tree-based reasoning with reward scoring
7. Performance Testing - Evaluates final model quality
8. Optional Visualizing - Generate images to see performance metric comparisons

Memory Efficient Design

• Uses LoRA adapters instead of full model fine-tuning
• Shared base model between SFT and reward scoring
• ~3.14GB memory usage vs ~6GB for merged models (48% savings)

Notes

• Models saved to models/ directory
• Processed data saved to data/ directory
• Requires GPU for training (CPU supported but slow)
• Teacher distillation costs vary for different-sized datasets and teacher models

Repository Structure

10:27/
├── pipeline/          # Training scripts (run in order)
├── data/              # Processed datasets
├── models/            # Trained models (qlora adapters only)
├── visuals/           # Visual artifacts
├── requirements-stable.txt
└── .env              # Your API keys (git ignored)