vllm-sr/mmbert32k-modality-router-lora

Modality Router LoRA - Smart Output Modality Selection

Part of the [MoM (Mixture of Models)](https://huggingface.co/vllm-sr) 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(
    "vllm-sr/mmbert-32k-yarn", num_labels=3
)
model = PeftModel.from_pretrained(base_model, "vllm-sr/mmbert32k-modality-router-lora")
tokenizer = AutoTokenizer.from_pretrained("vllm-sr/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: vllm-sr/mmbert32k-modality-router-merged

from transformers import pipeline

pipe = pipeline(
    "text-classification",
    model="vllm-sr/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: `vllm-sr/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/vllm-sr/mmbert32k-modality-router-lora}
}

Framework Versions

• —PEFT: 0.18.1
• —Transformers: 4.57.6
• —PyTorch: 2.9.1
• —Python: 3.12