---
language:
- en
license: apache-2.0
tags:
- medical
- healthcare
- text-generation
- jax
- pytorch
datasets:
- Mohammed-Altaf/medical-instruction-100k
pipeline_tag: text-generation
---

# 🩺 MedBrain-0.5B

**MedBrain-0.5B** is a highly efficient, custom-trained medical language model designed to provide accurate, structured, and context-aware responses to healthcare inquiries, clinical handoffs, and patient education instructions.

Originally trained meticulously in a pure **Google JAX / Flax** environment for extreme performance, the weights have been seamlessly merged and optimized into a standard PyTorch parameter format for universal compatibility and instant deployment.

## 🚀 Quick Start

You can run this model instantly using standard Hugging Face `transformers`:

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "suhailult777/MedBrain-0.5B"

# Load Tokenizer and Model
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)

# Format your medical prompt
prompt = "A patient presents with sudden shortness of breath and left-sided chest pain. What are the immediate triage steps?"
formatted_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"

# Run inference
inputs = tokenizer(formatted_prompt, return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=150)

print(tokenizer.decode(outputs[0], skip_special_tokens=False))
```

## 🧠 Architecture & Methodology
- **Parameter Count:** ~0.5 Billion parameters
- **Optimization Strategy:** Low-Rank Adaptation (LoRA) - Rank 16, Alpha 16
- **Training Infrastructure:** Custom JAX native dynamic loop utilizing `Optax` schedulers.
- **Base Architecture Mapping:** Transformer-based Causal LM.
- **Dataset:** Fine-tuned on the structured `Mohammed-Altaf/medical-instruction-100k` corpus, which provides vast arrays of physician-patient interactions.

## 🛠️ Intended Use
1. **Medical Triage Assistance:** Assisting clinicians in organizing thoughts around symptoms.
2. **Clinical Handoff Generators:** Structuring patient handoff notes quickly.
3. **Patient Education:** Formatting complex medical jargon into easy-to-understand explanations.

## ⚠️ Limitations & Clinical Warning
This model is built as an **experimental research artifact**. 
It should **never** be used for clinical decision-making, raw diagnostic purposes, or serving as a replacement for a licensed healthcare professional. LLMs can hallucinate confidently. Always consult a certified physician for medical advice.