llm-semantic-router/mmbert-32k-yarn
112.5k
mmBERT-32K-YaRN
Modern Multilingual BERT with 32K context length - Extended from 8K to 32K tokens using YaRN RoPE scaling.
This model extends jhu-clsp/mmBERT-base (Modern Multilingual BERT supporting 1800+ languages) from 8,192 to 32,768 maximum context length using YaRN (Yet another RoPE extensioN) scaling method.
Model Description
Intended Use
This model is designed for:
- Long-document understanding in any of 1800+ languages
- Semantic routing for LLM request classification
- Document classification with extended context
- Information retrieval from long texts
- Multilingual NLP tasks requiring long context
Part of the vLLM Semantic Router Mixture-of-Models (MoM) family.
Evaluation Results
Distance-Based Retrieval (Key Metric for Long-Context)
Summary: Perfect retrieval up to 2048 tokens. Long-range capability improved from baseline ~33% to 50% (averaged across all distances ≥1024).
Multilingual MLM Accuracy
Overall: 100% (10/10 languages tested)
Perplexity by Context Length
Position-wise Accuracy (16K context)
Training Details
Training Configuration
base_model: jhu-clsp/mmBERT-base
rope_scaling_type: yarn
original_max_position_embeddings: 8192
target_max_position_embeddings: 32768
scaling_factor: 4.0
yarn_beta_fast: 32.0
yarn_beta_slow: 1.0
# Training hyperparameters
learning_rate: 1e-5
batch_size: 1 (effective: 16 with gradient accumulation)
gradient_accumulation_steps: 16
num_epochs: 1
warmup_steps: 100
lr_scheduler: constant_with_warmup
mlm_probability: 0.3
bf16: trueTraining Data
- Dataset: CC-100 (Common Crawl) multilingual corpus
- Samples: 30,774 sequences
- Sequence Length: 32,768 tokens each
- Total Tokens: ~1B tokens
Hardware
- GPU: AMD Instinct MI300X (192GB VRAM)
- Training Time: ~6.5 hours
- Framework: PyTorch 2.3 + ROCm 6.2
Usage
Basic Usage
from transformers import AutoModelForMaskedLM, AutoTokenizer
model = AutoModelForMaskedLM.from_pretrained("llm-semantic-router/mmbert-32k-yarn")
tokenizer = AutoTokenizer.from_pretrained("llm-semantic-router/mmbert-32k-yarn")
# Multilingual MLM example
text = "The capital of France is <mask>."
inputs = tokenizer(text, return_tensors="pt")
outputs = model(**inputs)
mask_idx = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero()[0, 1]
logits = outputs.logits[0, mask_idx]
top5 = tokenizer.decode(logits.topk(5).indices)
print(top5) # ['Paris', 'Strasbourg', 'Nice', 'Brussels', 'Lyon']Long Context Usage (32K tokens)
# Process long documents (up to 32K tokens)
long_document = "..." * 30000 # Your long text in any of 1800+ languages
inputs = tokenizer(
long_document,
return_tensors="pt",
max_length=32768,
truncation=True
)
outputs = model(**inputs)Feature Extraction
import torch
# Get embeddings for downstream tasks
with torch.no_grad():
outputs = model(**inputs, output_hidden_states=True)
# Use last hidden state or pooled output
embeddings = outputs.hidden_states[-1].mean(dim=1) # Mean poolingONNX Runtime Usage
An ONNX export is available for high-performance inference with ONNX Runtime.
Python
import onnxruntime as ort
from transformers import AutoTokenizer
import numpy as np
# Load tokenizer and ONNX model
tokenizer = AutoTokenizer.from_pretrained("llm-semantic-router/mmbert-32k-yarn")
sess = ort.InferenceSession(
"onnx/model.onnx", # or download from HF
providers=['ROCmExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
)
# Inference
text = "What is the weather like today?"
inputs = tokenizer(text, return_tensors="np", padding=True)
outputs = sess.run(None, {
'input_ids': inputs['input_ids'].astype(np.int64),
'attention_mask': inputs['attention_mask'].astype(np.int64)
})
embeddings = outputs[0].mean(axis=1) # Mean poolingRust (ort-binding)
use onnx_semantic_router::MmBertEmbeddingModel;
let model = MmBertEmbeddingModel::load("./mmbert-32k-yarn-onnx", false)?;
let embeddings = model.embed("What is the weather?")?;Latency Benchmarks (AMD MI300X)
Limitations
- Long-range retrieval: While the model handles 32K context, retrieval accuracy drops significantly beyond 2048 tokens distance
- Training data: Trained on CC-100 which may have biases from web crawl data
- Compute requirements: Full 32K context requires significant GPU memory (~180GB for batch size 1)
Citation
@misc{mmbert-32k-yarn,
title={mmBERT-32K-YaRN: Extended Context Modern Multilingual BERT},
author={vLLM Semantic Router Team},
year={2026},
publisher={Hugging Face},
url={https://huggingface.co/llm-semantic-router/mmbert-32k-yarn}
}References
- mmBERT - Modern Multilingual BERT (1800+ languages)
- ModernBERT - Base architecture
- YaRN - Yet another RoPE extensioN method
- vLLM Semantic Router - Mixture-of-Models routing
License
MIT License (same as mmBERT base model)