Dr. AI v2 — Arabic / English Medical Assistant (NOT COMPLETE YET)
A self-contained 4B conversational medical assistant for Egyptian Arabic and English.
Two-stage fine-tune of google/medgemma-4b-it
(Gemma-3 4B), with both LoRA adapters merged into the base weights — load this repo directly,
no PEFT/base needed at inference.
Also available on GitHub: https://github.com/ehab215/Dr.-AI
How it was built
| Stage | Method | Data | Purpose |
|---|---|---|---|
1 — CPT (DR-AI-V1) | LoRA r=64 α=128 | ~1.1M Arabic/Egyptian + medical docs | Domain + dialect adaptation |
| 2 — SFT (this repo) | LoRA r=64 α=128, loss on assistant turn only | ~155K instruction/response pairs (Egyptian + MSA + English medical Q&A) | Instruction following as "Dr. AI" |
This build is the merged standalone from the best surviving Stage-2 checkpoint (step 12600).
Model facts
| Architecture | Gemma3ForConditionalGeneration (text used; vision tower unused) |
| Parameters | ~4B |
| Precision | bfloat16 |
| Tokenizer | Gemma (vocab 262208) |
| Fine-tuning context window | 1024 tokens (system + user + assistant) |
| Architecture max context | 131072 tokens (not fine-tuned that far — keep prompts/history within ~1024 for best behavior) |
| Chat format | Gemma-3 chat template (system / user / assistant) |
| Recommended generation | temperature=0.7, top_p=0.95, repetition_penalty=1.05, max_new_tokens=512 |
Quick start (Hugging Face)
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "ehab215/DR-AI-V2"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
messages = [
{"role": "system", "content": "You are Dr. AI, a helpful medical assistant. Answer in the language you are asked in."},
{"role": "user", "content": "ايه الفرق بين ضغط الدم الانقباضي و الانبساطي؟"},
]
inputs = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
inputs.pop("token_type_ids", None) # Gemma3 adds this for training; not needed for generation
out = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7, top_p=0.95, repetition_penalty=1.05)
print(tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))
Training details
- bf16 on A100 80GB; SFT loss computed only on the assistant response (prompt masked to -100).
- Stage 2: 3 epochs, LR 1e-4, cosine + 3% warmup, effective batch 32.
- Stage 1 validation perplexity: 13.0 → 3.29 (3.95×). Stage 2 best masked-val loss ≈ 1.92.
Intended use
Educational / informational medical Q&A and triage-style guidance in Egyptian Arabic and English.
⚠ Limitations & safety
Not a medical device. Do not use for diagnosis or treatment decisions. Outputs may be wrong or unsafe — always have a qualified clinician review. Arabic skews Egyptian; 1024-token fine-tuning context; instruction-tuned only (no RLHF). Always advises consulting a specialist.