Modality Router LoRA - Smart Output Modality Selection
Part of the MoM (Mixture of Models) family for vLLM Semantic Router.
A LoRA adapter fine-tuned on mmbert-32k-yarn (307M parameter ModernBERT with 32K context and 1800+ language support) that classifies user prompt intent into the appropriate response modality:
Usage
With PEFT (LoRA adapter)
from peft import PeftModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
# Load base model + LoRA adapter
base_model = AutoModelForSequenceClassification.from_pretrained(
"llm-semantic-router/mmbert-32k-yarn", num_labels=3
)
model = PeftModel.from_pretrained(base_model, "llm-semantic-router/mmbert32k-modality-router-lora")
tokenizer = AutoTokenizer.from_pretrained("llm-semantic-router/mmbert32k-modality-router-lora")
# Label mapping
labels = {0: "AR", 1: "DIFFUSION", 2: "BOTH"}
# Classify prompts
prompts = [
"What are the benefits of exercise?",
"A serene Japanese garden with cherry blossoms, watercolor style",
"Explain how neural networks work and generate a diagram showing the architecture",
]
model.eval()
for prompt in prompts:
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
outputs = model(**inputs)
pred = torch.argmax(outputs.logits, dim=-1).item()
probs = torch.softmax(outputs.logits, dim=-1)[0]
print(f"Prompt: {prompt[:60]}...")
print(f" -> {labels[pred]} (confidence: {probs[pred]:.3f})")
print()
Use the merged model (recommended for production)
For easier deployment without PEFT dependency, use the merged version:
llm-semantic-router/mmbert32k-modality-router-merged
from transformers import pipeline
pipe = pipeline(
"text-classification",
model="llm-semantic-router/mmbert32k-modality-router-merged",
)
result = pipe("Draw a picture of a sunset over mountains")
print(result) # [{'label': 'DIFFUSION', 'score': 0.97}]
Model Details
Architecture
- Base model:
llm-semantic-router/mmbert-32k-yarn (307M parameters, ModernBERT + YaRN RoPE)
- Context length: 32,768 tokens
- Languages: 1800+ (via Gemma 2 tokenizer with 256K vocab)
- Adaptation: LoRA (Low-Rank Adaptation) via PEFT
- LoRA rank: 16
- LoRA alpha: 32
- LoRA dropout: 0.1
- Target modules:
attn.Wqkv, attn.Wo, mlp.Wi, mlp.Wo
- Modules saved:
classifier, score
- Task type: Sequence Classification (3 classes)
- Trainable parameters: ~3.4M (1.09% of total 310M)
Training
- Epochs: 10
- Batch size: 32
- Learning rate: 2e-5
- Weight decay: 0.15 (adaptive)
- Loss function: Focal Loss (gamma=2.0) with inverse-frequency class weights
- Class imbalance handling: Focal Loss + sqrt-dampened class weights + minority oversampling
- Hardware: AMD Instinct MI300X GPU (192GB VRAM)
- Training time: ~2 minutes
Training Data
The model is trained on a curated combination of 10 public datasets plus seed examples:
DIFFUSION class (image generation intent)
AR class (text-only intent)
BOTH class (mixed modality intent)
Evaluation Results
Per-class Performance
Intended Use
This model is designed for routing LLM requests in multi-model serving systems like vLLM Semantic Router. It enables:
- Smart Output Modality Selection: Automatically determine whether a user query needs text, image, or both
- Automatic Paradigm Routing: Route requests to the right model backend (AR LLM vs Diffusion model)
- Cost Optimization: Avoid sending simple text queries to expensive image generation pipelines
Out-of-Scope Use
- Not suitable for content moderation or safety classification
- Not designed for multi-turn conversation context (single-turn prompt classification only)
- May have reduced accuracy for very short or ambiguous prompts
Citation
@misc{modality-router-2025,
title={Modality Router: Smart Output Modality Selection for Multi-Model Serving},
author={vLLM Semantic Router Team},
year={2025},
url={https://huggingface.co/llm-semantic-router/mmbert32k-modality-router-lora}
}
Framework Versions
- PEFT: 0.18.1
- Transformers: 4.57.6
- PyTorch: 2.9.1
- Python: 3.12
