ehab215/DR-AI-V2

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.

License

Gemma Terms of Use.