🩺 MediQwen: Advanced Medical Chatbot Fine-Tuning
MediQwen is a high-performance medical assistant developed by fine-tuning the Qwen-3-0.6B large language model. This project leverages Parameter-Efficient Fine-Tuning (PEFT) techniques, specifically LoRA, to adapt the model for specialized medical dialogue tasks using the AI Medical Chatbot dataset.
🚀 Key Features
- Base Model: Qwen/Qwen3-0.6B - A compact yet powerful foundation.
- Efficient Fine-Tuning: Utilizes LoRA (Low-Rank Adaptation) to minimize VRAM usage while maintaining high performance.
- Dataset: Trained on a high-quality subset of the ruslanmv/ai-medical-chatbot dataset (250k+ Q&A pairs).
- Chat Templates: Automatically applies role-based chat templates for seamless multi-turn conversations.
- Optimized for T4: Designed to run and train efficiently on Google Colab's T4 GPUs.
🏗️ Technical Architecture
The following diagram illustrates the fine-tuning workflow:
graph LR
A[Base Model: Qwen-3-0.6B] --> B{PEFT Wrapper}
B --> C[LoRA Adapters]
D[Medical Dataset] --> E[Chat Template Processing]
E --> F[SFT Training]
C --> F
F --> G[Fine-Tuned MediQwen]
🛠️ Installation & Setup
Ensure you have a modern Python environment and a CUDA-enabled GPU.
pip install -q torch transformers peft trl datasets
📖 Fine-Tuning Process
1. Data Preparation
We use the apply_chat_template method to structure the doctor-patient dialogues correctly.
def format_chat(example):
messages = [
{"role": "user", "content": example["Patient"]},
{"role": "assistant", "content": example["Doctor"]},
]
text = tokenizer.apply_chat_template(messages, tokenize=False)
return {"input_ids": tokenizer(text, truncation=True, max_length=512)["input_ids"]}
2. LoRA Configuration
The model is optimized using specific target modules to focus on the self-attention mechanism.
| Hyperparameter | Value |
|---|---|
| Rank (r) | 16 |
| Alpha | 32 |
| Dropout | 0.05 |
| Target Modules | q_proj, k_proj, v_proj, o_proj |
3. Training
Leveraging the SFTTrainer from the trl library for high-level abstraction.
trainer = SFTTrainer(
model=model,
train_dataset=dataset,
args=SFTConfig(
output_dir="./medical-qwen",
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
learning_rate=2e-4,
num_train_epochs=1,
fp16=True
)
)
💬 Usage & Inference
After fine-tuning, you can interact with the model using the standard Transformers pipeline:
from transformers import pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
prompt = "Patient: I've been feeling dizzy lately and have a metallic taste in my mouth. What could it be?"
# Note: Ensure the prompt is formatted using the chat template for best results.
📊 Dataset Overview
The ruslanmv/ai-medical-chatbot dataset provides a rich source of medical interactions:
- Patients: Descriptive symptoms and health concerns.
- Doctors: Professional, empathetic, and informative responses.
📜 License
Distributed under the Apache 2.0 License. See LICENSE for more information.
Developed with ❤️ using Qwen and Hugging Face