abdullah693/mixtral-8x7b-lhw-imnci-lora

Mixtral-8x7B · Lady Health Worker IMNCI Case-Decision Assistant (LoRA)

A LoRA adapter for mistralai/Mixtral-8x7B-Instruct-v0.1 that turns the model into a point-of-care decision-support assistant for Pakistan's Lady Health Workers (LHWs).

A Lady Health Worker advising a family in rural Sindh — the point-of-care setting this adapter is built for: offline, in the field, no doctor nearby.

Given a household clinical scenario described in plain language, the model produces a structured, protocol-faithful decision following Pakistan's national IMNCI / iCCM / EPI / CMAM / mhGAP / MNCH / family-planning guidelines:

Reasoning:      brief decision steps
Classification: the single protocol classification
Treatment:      what to do, in order
Medicine:       drug dose (duration); ...   or: none
Referral:       required=<true|false>; urgency=<none|urgent>; reason=<...>

The two correctness rules that drive the whole design: doses must match the protocol exactly, and wherever the protocol says refer, the answer must say refer (over-referral is acceptable, under-referral never is).

⚠️ Decision support, not a clinician. This model assists trained health workers; it does not replace clinical judgement. Every dose and referral must be verified against the current national protocol before acting. See Limitations & Safety below.

---

What it does

Input	Output
"A 2-year-old has had diarrhoea for 2 days. She is lethargic, eyes sunken, unable to drink."	Classification: Severe dehydration · Treatment: Plan C rapid rehydration · Medicine: ORS + zinc · Referral: required=true, urgent

It covers six domains: child illness, immunization, nutrition, maternal & newborn health (MNCH), mental health, and family planning.

Training

• Method: supervised fine-tuning (SFT), LoRA, train_on_inputs=false (loss on completions only).
• Base model: mistralai/Mixtral-8x7B-Instruct-v0.1 (47B-param sparse MoE, Apache-2.0).
• LoRA: r=64, alpha=128, dropout=0, target modules q_proj, k_proj, v_proj, o_proj.
• Schedule: 2 epochs / 396 steps, cosine LR with warmup, peak LR 1e-4, max_grad_norm=1, weight_decay=0.01.
• Result: train loss 1.24 → 0.50, final eval loss 0.525.

Training curves

Loss	Learning rate

Gradient norm	Head-to-head win rate

In a head-to-head preference evaluation the adapted model wins 74% of comparisons vs. 26% for the base model — a decisive improvement on protocol-faithful case decisions.

Usage

The adapter is merged into the base at load time and served in 8-bit — this is exactly what the demo Space does (see the quantization note below for why 8-bit and not 4-bit).

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel

BASE = "mistralai/Mixtral-8x7B-Instruct-v0.1"
ADAPTER = "abdullah693/mixtral-8x7b-lhw-imnci-lora"

tok = AutoTokenizer.from_pretrained(BASE)
base = AutoModelForCausalLM.from_pretrained(
    BASE,
    quantization_config=BitsAndBytesConfig(load_in_8bit=True),  # ~47 GB, fits one ~70 GB GPU
    device_map="auto",
)
model = PeftModel.from_pretrained(base, ADAPTER)
model = model.merge_and_unload()   # bake the LoRA delta into the (8-bit) weights → merged model

SYSTEM = (
    "You are a Lady Health Worker in rural Pakistan following the national IMNCI / LHW "
    "protocols. Reply in EXACTLY this format and STOP after the Referral line:\n"
    "Reasoning: ...\nClassification: ...\nTreatment: ...\nMedicine: ...\n"
    "Referral: required=<true|false>; urgency=<none|urgent>; reason=<...>"
)
scenario = ("A mother brings her 9-month-old boy (8 kg) with cough for 3 days. "
            "You count 52 breaths per minute. No chest indrawing. He is alert and feeding well.")

msgs = [{"role": "user", "content": SYSTEM + "\n\nScenario: " + scenario}]
ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(model.device)
out = model.generate(ids, max_new_tokens=400, do_sample=False, repetition_penalty=1.2)
print(tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True))

⚠️ Quantization note — serve at 8-bit, never 4-bit

This LoRA is a small-magnitude update (the delta is a fraction of a percent of each weight norm) that nonetheless moves behaviour a lot (74% head-to-head win rate). That makes the serving precision a correctness decision, not just a memory one:

• 4-bit quantization introduces ~5–7% per-weight error — larger than the LoRA delta itself. Whether you merge-then-quantize or quantize-then-merge, the delta lands inside the same 4-bit bucket as the base weight and is rounded away: a 4-bit model behaves almost identically to the un-fine-tuned base. (We confirmed this failure mode on a sibling model.)
• 8-bit quantization introduces only ~0.4% per-weight error — ~15× gentler — so the delta survives and the fine-tune is preserved. Use 8-bit.
• Cardinal rule: never evaluate or serve this adapter at 4-bit. Use 8-bit (recommended) or bf16 (full fidelity, ~94 GB).

Hardware: Mixtral 8x7B is ~94 GB in bf16 / ~47 GB in 8-bit / ~25 GB in 4-bit (don't). The demo Space loads the base in 8-bit, merges this adapter, and serves the merged 8-bit model on an on-demand ZeroGPU (NVIDIA RTX Pro 6000 Blackwell, xlarge = 96 GB).

Intended use

• Decision support and training/reference for community health workers and program designers.
• Offline-capable edge deployment is the design goal (the dataset and format target small, quantized on-device models; this Mixtral adapter is the high-capacity reference variant).

Limitations & safety

• Not a medical device and not a substitute for a clinician. Outputs must be verified against the current national protocol; guidelines change.
• Doses can be wrong. The model sometimes emits a plausible-but-incorrect dose or a non-canonical classification label. Treat every Medicine: line as a draft to confirm.
• Referral bias: the model is tuned to err toward referral; expect some over-referral.
• Scope: English only; Pakistan IMNCI/LHW protocols. Out-of-distribution scenarios (trauma, adult medicine, etc.) are not covered.
• Inherits the biases and failure modes of the Mixtral-8x7B-Instruct-v0.1 base model.

License

Apache-2.0, matching the base model. You are responsible for compliant, safe deployment.

---

Built for the Adaption Labs AutoScientist Challenge (Healthcare track). Trained with SFT + LoRA on a protocol-grounded synthetic dataset of LHW case decisions.