VirtualDoc

A medical AI assistant that combines Amazon Bedrock Knowledge Base retrieval with a fine-tuned Gemma-2-2B language model to answer health and medication questions through a ChatGPT-style web interface.

Unlike a standard RAG stack that pairs retrieval with an off-the-shelf API model, VirtualDoc's generation model was fine-tuned from scratch on curated medical Q&A data using QLoRA, then deployed as a real-time inference endpoint on Amazon SageMaker. Retrieval and generation are combined so that domain-specific facts come from retrieved documents while the fine-tuned model supplies medical reasoning and phrasing.

---

Project Motivation

General-purpose LLMs are not trained specifically on medical language and can produce fluent but unreliable answers to health questions. This project explores whether fine-tuning a small, efficient open model on real medical sources — combined with retrieval over the same sources — produces a more grounded and trustworthy medical assistant than either approach alone.

---

Skills Demonstrated

• Retrieval-Augmented Generation (RAG)
• LLM Fine-Tuning (QLoRA / LoRA with PEFT)
• 4-bit Quantization (bitsandbytes)
• Supervised Fine-Tuning (TRL SFTTrainer)
• Amazon Bedrock Knowledge Bases
• Amazon SageMaker Training Jobs
• Amazon SageMaker Real-Time Inference Endpoints
• FastAPI
• Next.js
• Docker Containerisation
• REST API Development
• Medical/Healthcare NLP

---

Features

• ChatGPT-style conversational interface
• Multiple conversation management with persistent, browser-based chat history
• Retrieval-Augmented Generation over Amazon Bedrock Knowledge Base
• Answer generation from a fine-tuned Gemma-2-2B model hosted on a SageMaker endpoint
• Inline citation panel showing retrieved source chunks and their similarity scores
• FastAPI REST backend
• Next.js frontend
• Docker containerisation for both services

---

Data Sources

The knowledge base and fine-tuning dataset are built from separate sets of public medical corpora.

Retrieval Knowledge Sources

Source	Content	Purpose
DailyMed	FDA prescription drug labeling (Structured Product Labels, SPL)	Retrieval
MedlinePlus	Consumer health information from the U.S. National Library of Medicine	Retrieval
PubMed Central (PMC)	Open-access biomedical research articles	Retrieval

Raw sources are downloaded, randomly sampled to 2,000 documents per source, and converted to plain text (data_pipeline/rag_data/), then uploaded to S3 and indexed into an Amazon Bedrock Knowledge Base for retrieval at query time.

Fine-Tuning Sources

Fine-tuning Datasets

• MedQuAD

  • Hugging Face: keivalya/MedQuad-MedicalQnADataset
  • Purpose: Fine-tuning

• MedDialog

  • Hugging Face: UCSD26/medical_dialog (processed.en)
  • Purpose: Fine-tuning

• PubMedQA

  • Hugging Face: pubmed_qa (pqa_labeled)
  • Purpose: Fine-tuning

• MedMCQA

  • Hugging Face: medmcqa
  • Purpose: Fine-tuning

• HealthSearchQA

  • Hugging Face: ruslanmv/ai-medical-chatbot
  • Purpose: Fine-tuning

Each source is downloaded in full and saved locally (data_pipeline/fine_tune_data/1_download.ipynb), then only 1,000 documents are randomly sampled from each source (data_pipeline/fine_tune_data/2_process.ipynb) and converted into instruction-style Q&A pairs, for a combined fine-tuning set of ~5,000 examples.

---

Model Fine-Tuning

The generation model is google/gemma-2-2b-it, fine-tuned with QLoRA on an Amazon SageMaker training job.

Fine-tuning Configuration

• Base Model: google/gemma-2-2b-it
• Method: QLoRA (4-bit NF4 quantization)
• LoRA Rank / Alpha: 16 / 32
• Target Modules: q_proj, k_proj, v_proj, o_proj
• Trainer: TRL SFTTrainer
• Training Data: ~5,000 instruction-formatted examples (1,000 sampled from each of five medical datasets)
• Epochs: 1

After training, the LoRA adapter is merged into the base model (data_pipeline/fine_tune_data/6_merge_and_upload.ipynb) and deployed to a SageMaker real-time endpoint (data_pipeline/fine_tune_data/7_deploy.ipynb) that the backend calls for inference.

---

Design Decisions

Retrieval

Amazon Bedrock Knowledge Bases provide managed semantic retrieval over the indexed medical corpora without requiring a self-managed vector database. Documents are chunked using a fixed-size strategy of 500 tokens per chunk with 10% overlap between consecutive chunks.

Fine-Tuned Generation

Rather than prompting a general-purpose API model, VirtualDoc uses a model fine-tuned on medical Q&A. The prompt instructs the model to prioritise retrieved context for medical-specific facts while allowing its fine-tuned knowledge to fill gaps when retrieval is insufficient.

Model Deployment

Hosting the fine-tuned model on a SageMaker real-time endpoint (rather than a third-party API) keeps the entire retrieval-and-generation loop inside AWS and demonstrates end-to-end ML infrastructure ownership, from training to serving.

Containerisation

The frontend and backend are containerised independently with Docker, allowing each service to be built, run, and scaled separately.

---

Architecture

User
 │
 ▼
Next.js Frontend
 │
 ▼
FastAPI Backend
 │
 ├──────────────┐
 ▼              ▼
Bedrock KB   SageMaker Endpoint
(retrieval)  (fine-tuned Gemma-2-2B)

---

Retrieval & Generation Pipeline

User Question
      │
      ▼
Next.js Frontend
      │
      ▼
FastAPI Backend
      │
      ▼
Bedrock KB Retrieve (top-k chunks)
      │
      ▼
Build Context + Gemma Chat Prompt
      │
      ▼
SageMaker Endpoint (fine-tuned Gemma-2-2B)
      │
      ▼
Answer + Cited Chunks

---

Fine-Tuning Pipeline

Raw Sources (MedQuAD, MedDialog, PubMedQA, MedMCQA, HealthSearchQA)
      │
      ▼
Download & Save Locally (1_download.ipynb)
      │
      ▼
Sample 1,000 Docs/Source & Convert to Instruction Q&A (2_process.ipynb)
      │
      ▼
Upload Training Set to S3 (3_upload_to_s3.ipynb)
      │
      ▼
Cache Base Model in S3 (4_cache_base_model.ipynb)
      │
      ▼
QLoRA Fine-Tuning, Local or SageMaker Training Job
(5_finetune_local.ipynb / 5_finetune_sagemaker_job.ipynb)
      │
      ▼
Merge LoRA Adapter into Base Model (6_merge_and_upload.ipynb)
      │
      ▼
Deploy to SageMaker Real-Time Endpoint (7_deploy.ipynb)

---

AWS Services Used

AWS Services

• Amazon Bedrock Knowledge Bases – Semantic retrieval over medical knowledge sources.
• Amazon SageMaker Training Jobs – QLoRA fine-tuning of the Gemma 2B model.
• Amazon SageMaker Endpoints – Real-time inference for the fine-tuned model.
• Amazon S3 – Storage for datasets, training data, and model artifacts.
• AWS IAM – Secure access control and permissions management.

---

Technology Stack

Frontend

• Next.js
• React
• TypeScript
• Tailwind CSS

Backend

• FastAPI
• Python
• Boto3

Machine Learning

• Hugging Face Transformers
• PEFT (LoRA)
• TRL (SFTTrainer)
• bitsandbytes (4-bit quantization)
• PyTorch

Cloud

• Amazon Bedrock
• Amazon SageMaker
• Amazon S3
• IAM

Containerisation

• Docker

---

Running Locally

Backend

cd backend
cp .env.example .env   # fill in AWS_REGION, KNOWLEDGE_BASE_ID, SAGEMAKER_ENDPOINT, AWS credentials
pip install -r requirements.txt
uvicorn main:app --reload

API available at http://127.0.0.1:8000/docs.

Frontend

cd frontend
cp .env.local.example .env.local   # set NEXT_PUBLIC_API_URL
npm install
npm run dev

Docker Compose

docker compose up --build

Frontend available at http://localhost:3000, backend at http://localhost:8000.

---

Example Queries

• I have a headache and a runny nose. What are the most common causes?
• What is the difference between cold and flu symptoms?
• When should I see a doctor for a persistent sore throat and cough?
• What are some safe, over-the-counter remedies for seasonal allergies?

---

Medical Disclaimer

VirtualDoc is a portfolio project demonstrating RAG and LLM fine-tuning techniques. It is not a medical device and does not provide medical advice. Always verify health-related guidance with a qualified healthcare professional.

---

Future Improvements

• Source citations with direct links back to DailyMed / MedlinePlus / PubMed Central entries.
• Streaming responses from the SageMaker endpoint.
• Larger and more diverse fine-tuning dataset with held-out evaluation.
• Automated evaluation harness for factual accuracy and hallucination rate.
• Persistent conversation history using a cloud database.
• User authentication and role-based access control.
• Cloud deployment of frontend/backend (ECS or App Runner) alongside the existing SageMaker/Bedrock infrastructure.
• Automated CI/CD deployment.

---

Acknowledgements

The frontend implementation was developed with the assistance of ChatGPT (OpenAI), including UI design guidance, code refinement, and debugging support.

The overall system architecture, model fine-tuning pipeline, Bedrock Knowledge Base integration, SageMaker training and deployment, backend implementation, and Docker containerisation were independently designed and implemented by the project author.

---

License

This project is licensed under the MIT License.