Given a prompt and a list of available models, which model is the best fit? A lightweight classifier that routes prompts to the right LLM by task category โ built for agentic AI systems.
Prompt in, task category + confidence out โ then the best available model for that category.
The classifier is trained on task categories only โ never specific model names โ so the model registry can change without retraining it.
Trained on ~1,700 examples (340/category) โ see the dataset card for how it was built.
Load the artifact directly from the Hub, or use the llm_router Python package.
from huggingface_hub import hf_hub_download
from sentence_transformers import SentenceTransformer
import joblib
path = hf_hub_download(repo_id="ai-mitra/llm-router", filename="classifier.joblib")
payload = joblib.load(path)
classifier, classes = payload["classifier"], payload["classes"]
embedder = SentenceTransformer(payload["embedding_model"])
embedding = embedder.encode(["Explain what a Python decorator does"])
probabilities = classifier.predict_proba(embedding)[0]
print(dict(zip(classes, probabilities)))
# {'coding': 0.87, 'reasoning': 0.06, 'security': 0.02, 'simple': 0.03, 'summarization': 0.02}