llm-semantic-router / mmbert32k-modality-router-merged

Modality Router (Merged) - Smart Output Modality Selection

Part of the MoM (Mixture of Models) family for vLLM Semantic Router .

This is the merged (ready-to-use) version of mmbert32k-modality-router-lora . LoRA weights have been merged into the mmbert-32k-yarn base model for easy deployment without the PEFT dependency.

A text classifier based on ModernBERT (307M params, 32K context, 1800+ languages) that determines the appropriate response modality for user prompts:

Quick Start

Pipeline API (simplest)

from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="llm-semantic-router/mmbert32k-modality-router-merged",
)

results = classifier([
    "What are the benefits of exercise?",
    "A serene Japanese garden with cherry blossoms, watercolor style",
    "Explain how neural networks work and generate a diagram",
])

for r in results:
    print(f"{r['label']}: {r['score']:.3f}")
# AR: 0.995
# DIFFUSION: 0.717
# BOTH: 0.978

Direct Model Usage

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

model = AutoModelForSequenceClassification.from_pretrained(
    "llm-semantic-router/mmbert32k-modality-router-merged"
)
tokenizer = AutoTokenizer.from_pretrained(
    "llm-semantic-router/mmbert32k-modality-router-merged"
)

prompts = [
    "Summarize the key points of quantum computing",
    "portrait of a woman in renaissance style, oil painting, dramatic lighting",
    "Write a blog post about climate change and include relevant charts",
]

model.eval()
inputs = tokenizer(prompts, return_tensors="pt", truncation=True, padding=True, max_length=512)
with torch.no_grad():
    outputs = model(**inputs)

predictions = torch.argmax(outputs.logits, dim=-1)
labels = model.config.id2label
for prompt, pred_id in zip(prompts, predictions):
    print(f"{labels[pred_id.item()]}: {prompt[:60]}...")
# AR: Summarize the key points of quantum computing...
# DIFFUSION: portrait of a woman in renaissance style, oil painting, d...
# BOTH: Write a blog post about climate change and include releva...

Integration with vLLM Semantic Router

# Example: Route requests to different model backends
def route_request(prompt: str, classifier) -> str:
    """Route a user prompt to the appropriate model backend."""
    result = classifier(prompt)[0]
    modality = result["label"]
    confidence = result["score"]

    if modality == "AR":
        return call_llm_backend(prompt)        # e.g., Llama, Qwen
    elif modality == "DIFFUSION":
        return call_diffusion_backend(prompt)   # e.g., Flux, SDXL
    else:  # BOTH
        text = call_llm_backend(prompt)
        image = call_diffusion_backend(prompt)
        return combine_response(text, image)

ONNX Runtime (for production latency)

The base model ( mmbert-32k-yarn ) supports ONNX export for sub-5ms inference on AMD MI300X GPUs.

Model Details

Training Configuration

Training Data

Trained on a curated combination of 10 public datasets covering diverse prompt styles:

DIFFUSION class

• Gustavosta/Stable-Diffusion-Prompts - 80K curated SD prompts
• FredZhang7/stable-diffusion-prompts-2.47M - 2.47M SD prompts
• nateraw/parti-prompts - Google Parti benchmark
• fal/image-generation-prompts - Diverse image prompts
• allenai/WildChat (mined) - Real user image requests

AR class

• OpenAssistant/oasst2 - 135K instruction conversations
• tatsu-lab/alpaca - 52K Stanford instructions
• databricks/databricks-dolly-15k - 15K categorized instructions
• stingning/ultrachat - 1.5M multi-turn conversations
• allenai/WildChat (mined) - Real user text prompts

BOTH class

• mqliu/InterleavedBench - Gold-standard interleaved text+image (EMNLP 2024)
• allenai/WildChat (mined) - Real user multimodal prompts
• Curated seed examples (40+ across diverse domains)

Evaluation Results

Per-class Performance

Example Classifications

Intended Use

This model is designed for routing LLM requests in multi-model serving systems:

• Smart Output Modality Selection : Automatically determine whether a user query needs text, image, or both
• Automatic Paradigm Routing : Route requests to the right backend (AR LLM vs Diffusion model vs both)
• Cost Optimization : Avoid sending simple text queries to expensive image generation pipelines
• Latency Reduction : Skip unnecessary model invocations by predicting the needed output type upfront

Limitations

• Single-turn prompt classification only (no conversation context)
• Primarily trained on English data (multilingual capability inherited from base model)
• Not designed for content moderation or safety classification

Related Models

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-merged}
}

Framework Versions

• Transformers: 4.57.6
• PyTorch: 2.9.1
• Safetensors: 0.5.x
• Python: 3.12