62 lines
2.6 KiB
Markdown
62 lines
2.6 KiB
Markdown
---
|
|
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.
|