center-of-excellence/complexity_classifier

Complexity Classifier

Binary text classifier that predicts whether an LLM task prompt is easy or hard, used by the model-routing pipeline to select the cheapest model that will succeed.

Fine-tuned from jhu-clsp/mmBERT-base on an internal benchmark of agent task prompts with max_length=512.

Labels

ID	Label
0	EASY
1	HARD

Metrics (test set, seed 42)

Metric	Value
Accuracy	0.7754
Macro F1	0.7742
F1 (hard)	0.7577
Precision	0.7734
Recall	0.7427
AUROC	0.8573
MCC	0.5490

Training Details

Hyperparameter	Value
Base model	jhu-clsp/mmBERT-base
Max length	512
Learning rate	2e-5
Batch size	16
Seed	42

How to use

from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="center-of-excellence/complexity_classifier",
)

result = classifier("Write a Python function that reverses a string.")
# [{'label': 'EASY', 'score': 0.91}]

Load a pinned version:

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

repo = "center-of-excellence/complexity_classifier"

tokenizer = AutoTokenizer.from_pretrained(repo, revision="v1")
model = AutoModelForSequenceClassification.from_pretrained(repo, revision="v1")

inputs = tokenizer("Your prompt here", return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
    logits = model(**inputs).logits
label = model.config.id2label[logits.argmax().item()]
print(label)  # EASY or HARD