aman0419/VitalLM-12M-Medical

🏥 VitalLM-12M: A Medical Small Language Model

VitalLM-12M is a 12-million parameter language model specialized for biomedical text generation.

📊 Model Stats

• Parameters: 12M
• Architecture: Decoder-only Transformer (Custom SLM)
• Training Data: ~60M Tokens (PubMed Abstracts + Medical Dialogue)
• Vocabulary: 8,000 (Custom BPE Tokenizer)
• Context Window: 256 tokens
• Training Hardware: Single NVIDIA P100 GPU

🚀 Usage

Since this is a custom architecture, you need the model.py file included in this repository to run it.

import torch
from model import SLM, SLMConfig
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download

# 1. Download Model & Weights
repo_id = "aman0419/VitalLM-12M-Medical" 
model_path = hf_hub_download(repo_id, filename="vital_lm_12m.pt")
tokenizer_path = hf_hub_download(repo_id, filename="vital_tokenizer.json")

# 2. Initialize Architecture
config = SLMConfig(vocab_size=8000, block_size=256, n_embd=256, n_head=4, n_layer=12)
model = SLM(config)

# 3. Load Weights
state_dict = torch.load(model_path, map_location='cpu')
model.load_state_dict(state_dict)
model.eval()

# 4. Generate
tokenizer = Tokenizer.from_file(tokenizer_path)
input_text = "Hypertension is"
ids = torch.tensor(tokenizer.encode(input_text).ids).unsqueeze(0)

with torch.no_grad():
    output = model.generate(ids, max_new_tokens=50)
    print(tokenizer.decode(output[0].tolist()))

Remote API (Decentralized Inference)

You can use this model remotely (without downloading weights) by connecting to the hosted inference API. This demonstrates the "Edge-Client" capability of the VitalLM framework.

pip install gradio_client

from gradio_client import Client

# Connect to the live VitalLM Space
client = Client("aman0419/VitalLM-12M") 

print("Querying VitalLM...")
result = client.predict(
        "User: I have a severe headache and sensitivity to light. Assistant:", # Prompt
        100,  # Max Length
        0.7,  # Temperature
        fn_index=0
)

print(result)

Limitations

Hallucination: As a small model (12M), it may generate plausible-sounding but incorrect medical facts. It is intended for research and demonstration purposes only.

Context: Optimized for short medical definitions and simple Q&A.

Dialogue Bias: The model was trained on chat data and may occasionally default to conversational fillers (e.g., "Hi, thanks for asking") if not prompted specifically.