Hierarchical Medical Response Generation

11-667 Large Language Models: Methods and Applications - Mini Project

This project aims to fine-tune small language models to generate medical consultation responses in a structured, hierarchical format following the Pyramid Principle and MECE (Mutually Exclusive, Collectively Exhaustive) framework.

📋 Project Overview

Objectives

1. Structured Output Generation: Train models to organize medical responses into multi-level (≥3 levels) hierarchical lists with ≤5 points per level
2. Medical Domain Application: Apply structured reasoning to doctor-patient conversations in healthcare contexts
3. Model Comparison: Compare small foundation models' ability to learn and generate structured outputs

Models

• Qwen2.5-0.5B (base): Smaller, less instruction-tuned, more sensitive to prompt format
• DeepSeek-R1 1.5B (instruct): Larger, better instruction adherence, more stable structure output

Dataset

Patient-Doctor-Conversation - Contains ~3,300 medical consultation dialogues with patient descriptions, conversations, and severity labels.

🚀 Quick Start

Prerequisites

# Conda environment (recommended)
conda activate 11667

# Or create new environment
conda create -n 11667 python=3.10
conda activate 11667

Installation

# Install dependencies
pip install -r requirements.txt

# For Modal fine-tuning:
pip install modal

Environment Setup

Set your DeepSeek API key in ~/.zshrc (or ~/.bashrc):

export DEEPSEEK_API_KEY="your_api_key_here"

Then reload:

source ~/.zshrc

📁 Project Structure

mini_project/
├── README.md                    # Project documentation
├── requirements.txt             # Python dependencies
├── main.py                      # Original test script (legacy)
├── config/                      # Configuration files
│   ├── __init__.py
│   └── settings.py              # All configuration settings
├── src/                         # Source code
│   ├── __init__.py
│   ├── data_loader.py           # Dataset loading and caching
│   ├── converter.py             # API client for hierarchical conversion
│   ├── processor.py             # Multi-process batch processor
│   └── utils/                   # Utility functions
│       └── __init__.py
├── scripts/                     # Executable scripts
│   ├── download_dataset.py      # Download and cache dataset
│   └── convert_dataset.py       # Convert dataset to hierarchical format
└── data/                        # Data directory
    ├── raw/                     # Downloaded raw dataset (cached)
    └── processed/               # Converted hierarchical datasets
        └── hierarchical_dataset/  # HuggingFace format for fine-tuning

🔄 Dataset Workflow

Step 1: Download Dataset (Optional)

Download and cache the dataset without processing:

python scripts/download_dataset.py

Step 2: Convert to Hierarchical Format

Convert doctor responses to hierarchical structure:

# Test mode: Process only 16 examples
python scripts/convert_dataset.py

# Full dataset: Edit config/settings.py and set TEST_MODE = False
python scripts/convert_dataset.py

Features

• Multi-process Parallel Processing: Concurrent API calls (4 workers by default)
• Automatic Checkpointing: Saves progress every 10 samples (configurable)
• Retry Logic: Handles API failures with exponential backoff
• Progress Tracking: Real-time progress bar with success rate statistics
• Format Preservation: Maintains original HuggingFace dataset structure

Configuration

Adjust settings in config/settings.py:

# Processing Configuration
MAX_WORKERS = 4      # Number of parallel API requests
TEST_MODE = True     # Set to False for full dataset
TEST_SIZE = 16       # Number of examples in test mode
SAVE_INTERVAL = 10   # Save checkpoint every N samples
MAX_RETRIES = 3      # Retry failed API calls
TIMEOUT = 60         # API call timeout in seconds

Output Format

The processed dataset maintains the same structure as the original:

{
  "Description": "what does abutment of the nerve root mean",
  "Patient": "hi doctor I am just wondering...",
  "Doctor": "- Assessment\n  - Query Understanding\n    - ...",  # Hierarchical format
  "Status": "medium severity"
}

Note: Failed conversions fall back to original doctor responses.

📊 Hierarchical Structure Criteria

All converted responses follow:

1. Depth: Minimum 3 levels of hierarchy
2. Breadth: Maximum 5 points per level
3. Pyramid Principle: Top-down reasoning (answer first, then supporting details)
4. MECE: Mutually exclusive, collectively exhaustive categorization

Example:

- Main Assessment
  - Diagnosis Category 1
    - Specific symptom observation
    - Relevant medical indicator
  - Diagnosis Category 2
    - Supporting evidence
- Treatment Recommendations
  - Primary intervention
    - Medication details
    - Dosage information
  - Secondary considerations

🛠️ Development Workflow

1. Data Preparation

# Download dataset
python scripts/download_dataset.py

# Test conversion (16 examples)
python scripts/convert_dataset.py

# Full conversion: Set TEST_MODE = False in config/settings.py
python scripts/convert_dataset.py

2. Model Fine-tuning

(To be implemented)

• Load base models (Qwen2.5-0.5B, DeepSeek-R1 1.5B)
• Fine-tune on hierarchical dataset from data/processed/hierarchical_dataset/
• Evaluate structure quality and medical accuracy

3. Model Fine-tuning on Modal

Important correction based on actual reference data: Memory scales as batch_size × seq_len², not model size. Using your 120MB model reference (batch=50, seq=1024 → 70GB), corrected batch sizes below.

Setup (One-time)

pip install modal
modal token new  # Authenticate

# Create volumes
modal volume create medical-dataset-volume
modal volume create medical-models-volume

# Upload dataset
modal volume put medical-dataset-volume \
  data/processed/hierarchical_dataset_clean /hierarchical_dataset_clean

Recommended Training Plans

Plan A: Safest ⭐ (RECOMMENDED)

Model: Qwen2.5-0.5B (Full Fine-tuning)
batch_size=16, context_length=512, epochs=3
Total iterations: 558 (2,980 samples / 16 batch × 3)
GPU: A100 (80GB)
Time: 15-20 minutes
Cost: $0.71
Peak memory: 4-6 GB
Command: modal run --detach modal_qwen_finetune_full.py::train

Plan B: Maximum Quality

Model: DeepSeek-R1 1.5B (Full + Gradient Checkpointing)
batch_size=4, context_length=512, epochs=3
Total iterations: 2,235
GPU: A100 (80GB)
Time: 45-60 minutes
Cost: $1.50
Peak memory: 8-10 GB
Command: modal run --detach modal_deepseek_finetune_full.py::train

Plan C: Budget

Model: Qwen2.5-0.5B (LoRA)
batch_size=32, context_length=512, epochs=3
Total iterations: 280
GPU: A10G (24GB)
Time: 8-10 minutes
Cost: $0.06
Peak memory: 2-3 GB
Command: modal run --detach modal_qwen_finetune_lora.py::train

Correct Modal Commands

cd scripts

# Qwen full fine-tuning (safest)
modal run --detach modal_qwen_finetune_full.py::train

# OR: DeepSeek full fine-tuning (best quality)
modal run --detach modal_deepseek_finetune_full.py::train

# OR: Qwen LoRA (fastest/cheapest)
modal run --detach modal_qwen_finetune_lora.py::train

# Monitor progress
modal apps list           # Get APP_ID
modal logs APP_ID         # View logs

# Download results
modal volume ls medical-models-volume
modal volume get medical-models-volume qwen_full_TIMESTAMP/final ./outputs/

Memory Reality (Corrected)

Why these batch sizes?

Your 120MB model: batch=50, seq=1024 → 70GB peak memory

This comes from attention matrices (Q @ K^T): batch × heads × seq_len²

• Not model size! Model weights are only ~500MB
• Quadratic scaling: reducing seq_len by 2x saves 75% activation memory
• Using seq=512 instead of 1024 brings memory down 10x

Qwen2.5-0.5B (500M params):

• batch=16, seq=512 → ~4-6GB ✓
• batch=32, seq=512 → ~8-10GB ✓
• batch=8, seq=1024 → ~8-12GB (risky)

DeepSeek-R1 1.5B (1.5B params):

• batch=4, seq=512 → ~6-8GB ✓ (gradient checkpointing helps)
• batch=8, seq=512 → ~8-12GB ✓
• batch=4, seq=1024 + checkpointing → ~25-30GB ⚠️ (tight)

Expected Results

Based on your 120MB model (20,000 iters → 0.02 loss):

Our Qwen setup (558 iterations):

• Should reach 0.08-0.15 loss (hierarchical generation is harder)
• Training should plateau after epoch 2
• Structure quality improves non-linearly with loss

4. Model Inference & Testing

Test fine-tuned models locally with automatic comparison to base models.

Quick Test

# Test Qwen full fine-tuned model
python scripts/test_inference.py qwen full

# Test Qwen LoRA
python scripts/test_inference.py qwen lora

# Test DeepSeek full
python scripts/test_inference.py deepseek full

# Test DeepSeek LoRA
python scripts/test_inference.py deepseek lora

# With custom token limit
python scripts/test_inference.py qwen full --max-tokens 1024

Features

• Automatic Path Inference: Models automatically loaded from outputs/{model}_{type}_final/
• Device Detection: Auto-detects CUDA, MPS (M1), or CPU
• Always Compares: Automatically loads and compares with base model
• LoRA Verification: For LoRA models, verifies weights are actually trained (non-zero)
• Structure Analysis: Checks if fine-tuned model adds hierarchical bullets (evidence of training effectiveness)

Expected Output

For each test, you'll see:

1. Model Loading: Confirms fine-tuned model loaded successfully
2. LoRA Verification (if applicable): Shows LoRA layer statistics
3. Fine-tuned Response: Generation from your trained model
4. Base Model Response: Generation from untuned base model
5. Comparison Analysis:
   • Identifies if outputs differ
   • Checks for hierarchical structure (bullets)
   • Provides evidence of fine-tuning effectiveness

Directory Structure

Ensure your fine-tuned models are saved in:

outputs/
├── qwen_full_final/      # For: python test_inference.py qwen full
├── qwen_lora_final/      # For: python test_inference.py qwen lora
├── deepseek_full_final/  # For: python test_inference.py deepseek full
└── deepseek_lora_final/  # For: python test_inference.py deepseek lora

Each directory should contain:

• config.json / model.safetensors or .bin (model weights)
• tokenizer.model (tokenizer)
• adapter_config.json (for LoRA only)

5. Evaluation

Comprehensive evaluation of all 4 fine-tuned models on a high-quality evaluation dataset.

Evaluation Dataset Generation (Two-Step Process)

Step 1: Generate Plain Responses

python scripts/generate_evaluation_dataset.py --num-samples 40

• Generates 40 new medical Q&A scenarios using few-shot prompting
• Uses real training examples as templates for style consistency
• Features: 20 concurrent workers, rate-limited API calls, full logging
• Output: data/processed/evaluation_dataset_synthetic/

Step 2: Convert to Hierarchical Format

python scripts/convert_evaluation_dataset.py

• Converts plain responses to hierarchical structure
• Uses proven conversion logic (99.73% success rate from training)
• Output: data/processed/evaluation_dataset_hierarchical/

Model Evaluation

Evaluate all 4 fine-tuned models on the evaluation dataset:

# Full evaluation (40 samples per model, ~15-30 minutes)
python scripts/evaluate_models_on_dataset.py

# Quick test (5 samples per model, ~2-3 minutes)
python scripts/evaluate_models_on_dataset.py --max-samples 5

# Specific models only
python scripts/evaluate_models_on_dataset.py --models qwen_full qwen_lora

Metrics Computed:

• Text Similarity: String Similarity, ROUGE-L, BLEU
• Structural Quality: MECE Score, Depth Score, Constraint Score, Grouping Score
• Performance: Success Rate, Generation Time, Response Length
• Breakdown: By severity level (low/medium/high)

Output:

• evaluation_results/model_evaluation_YYYYMMDD_HHMMSS.json
• Contains detailed metrics for all samples and summary statistics per model

Evaluation Documentation

For detailed evaluation methodology, metrics interpretation, and troubleshooting, see: 📖 README_EVALUATION.md

This comprehensive guide includes:

• Dataset generation process with quality assurance
• Detailed explanation of all evaluation metrics
• How to run evaluations and interpret results
• Troubleshooting common issues
• Comparative analysis of model performance

💻 Hardware Requirements & Training Feasibility

AWS g5.2xlarge Compatibility (24GB VRAM, 32GB RAM)

Both target models can be fine-tuned on g5.2xlarge instances:

Qwen2.5-0.5B - Comfortable Fit ✅

Configuration	VRAM Usage	Batch Size	Training Time (3 epochs)
Full Fine-tuning (FP16)	~12GB	8	~20 minutes
LoRA Fine-tuning	~4GB	32+	~15 minutes

Recommended Configuration:

TrainingArguments(
    per_device_train_batch_size=8,
    gradient_accumulation_steps=2,  # Effective batch = 16
    fp16=True,
    max_seq_length=512,
    learning_rate=2e-5,
)

DeepSeek-R1-1.5B - Requires Optimization ✅

Configuration	VRAM Usage	Batch Size	Training Time (3 epochs)
Full Fine-tuning (FP16 + Gradient Checkpointing)	~22GB	2	~40 minutes
LoRA Fine-tuning	~6GB	16+	~25 minutes

Recommended Configuration:

TrainingArguments(
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,  # Effective batch = 16
    gradient_checkpointing=True,    # Required!
    fp16=True,
    max_seq_length=512,
    learning_rate=1e-5,
)

Memory Breakdown

Full Fine-tuning Memory Components:

• Model weights (FP16): 0.5B → 0.93GB, 1.5B → 2.79GB
• Gradients (FP16): Same as weights
• Optimizer states (AdamW): 8 bytes/param
• Activations: 2-4GB (varies with batch size)

Optimization Techniques:

1. Gradient Checkpointing (Required for DeepSeek-R1-1.5B)

   • Saves 30-50% activation memory
   • ~20% slower training

   model.gradient_checkpointing_enable()

2. Flash Attention (Optional, provides 20-40% speedup)

   pip install flash-attn --no-build-isolation

   model = AutoModelForCausalLM.from_pretrained(
       model_name,
       attn_implementation="flash_attention_2"
   )

3. Mixed Precision Training (Required)

   • Use FP16 for all training
   • Reduces memory by 50% compared to FP32

📊 Dataset Statistics

Final Conversion Results

Metric	Count	Percentage
Original Dataset	3,325 samples	100%
Successfully Converted	3,311 samples	99.73%
Failed Conversions	9 samples	0.27%
Final Training Dataset	3,311 samples	-

Failed Samples

9 samples failed conversion because their Doctor field was NULL in the original dataset (data quality issue, not conversion failure). These have been filtered out from the training dataset.

Failed sample indices: [98, 99, 733, 785, 1060, 1069, 1236, 1381, 1565]

Dataset Format

The processed dataset uses Apache Arrow format (HuggingFace standard):

• Format: .arrow binary columnar format
• Organization: Each sample is an independent record (no separators needed)
• Access: dataset[0], dataset[1], etc. or dataset[:10] for batch
• Training Ready: Direct compatibility with HuggingFace Trainer
• Size: ~8.7 MB (3,311 samples)

Directory structure:

data/processed/hierarchical_dataset_clean/
├── data-00000-of-00001.arrow  # Main data file
├── dataset_info.json           # Schema and metadata
└── state.json                  # Dataset state

Fields:

• Description: Medical question/issue (string)
• Patient: Patient's detailed description (string)
• Doctor: Hierarchical medical response (string, markdown format)
• Status: Severity level (string: "low/medium/high severity")
• _original_doctor: Original doctor response for comparison (string)
• _conversion_success: Conversion status flag (boolean)

📈 Expected Outcomes

1. Structured Training Data: 3,311 hierarchically formatted medical responses ✅
2. Fine-tuned Models: Two models capable of generating structured medical advice
3. Comparative Analysis: Performance comparison between base and instruct models on structured output tasks

🔧 Technical Details

API Integration

• Uses DeepSeek API for hierarchical conversion
• OpenAI-compatible SDK interface
• No rate limiting (per DeepSeek API docs)
• Conservative concurrent requests (4 workers) for stability

Multiprocessing Strategy

• Process-level parallelism for concurrent API calls
• Independent worker processes to avoid GIL limitations
• Automatic progress tracking and checkpoint coordination

Data Pipeline

1. Download and cache dataset from HuggingFace
2. Parallel API conversion with retry logic and exponential backoff
3. Save checkpoints every 10 samples
4. Export in HuggingFace Dataset format (preserves original structure)
5. Failed conversions fall back to original responses

📝 Notes

Development Environment

• All processing uses the 11667 conda environment
• API key must be set in environment variables before running
• Dataset: 3,325 medical consultation dialogues (~6MB total)

Dataset Conversion

• DeepSeek API has no rate limiting (per official docs)
• 10 parallel workers for conversion (optimized)
• Full dataset conversion time: ~54 minutes (10 workers, 3,311 samples)
• Checkpoints saved every 10 samples for interruption recovery
• Arrow format for HuggingFace ecosystem compatibility
• Final success rate: 99.73% (3,311/3,320 valid samples)

Training Environment

• Tested on: AWS g5.2xlarge (24GB VRAM, 32GB RAM)
• GPU: NVIDIA A10G Tensor Core GPU
• Both models fit: Full fine-tuning possible for both Qwen2.5-0.5B and DeepSeek-R1-1.5B
• Dependencies: PyArrow automatically installed via datasets library (no manual setup needed)

👥 Team

Course: 11-667 Large Language Models: Methods and Applications
Carnegie Mellon University, Fall 2025

📄 License

This project is for educational purposes as part of CMU coursework.

🔗 References

• DeepSeek API Documentation
• HuggingFace Datasets
• Pyramid Principle
• MECE Framework