Pick and Spin query-complexity classifier

DistilBERT (distilbert-base-uncased) fine-tuned to put a prompt into one of three tiers, SIMPLE, MEDIUM or COMPLEX. It is the second stage of the hybrid classifier in the Pick router of Pick and Spin: prompts that match none of Pick's keyword lists are classified by this model.

Paper: Pick and Spin: Cold-Start-Aware Routing for Self-Hosted LLM Serving, Bhanu Prakash Vangala and Tanu Malik, IEEE CLOUD 2026.

Labels

Each prompt is labelled with the smallest group of models that an LLM judge marks correct on it in the Pick and Spin static baseline:

Label Meaning
SIMPLE a 1B-3B model (Llama-3.2-1B/3B, Qwen2.5-1.5B, Gemma-2-2B) answers correctly
MEDIUM no 1B-3B model does, but a 7B-14B model (Qwen2.5-7B/14B, Llama-3.1-8B, Gemma-2-9B) does
COMPLEX only Gemma-3-27B, Llama-3-70B or Kimi-K2 answers correctly, or no model does

The prompts are 31,019 questions from HumanEval, MBPP, GSM8K, MATH, TruthfulQA, MMLU-Pro, ARC and HellaSwag, split 80/20 with a fixed seed into 24,815 training and 6,204 validation prompts (src/classifier/generate_labels.py).

Training

5 epochs, batch size 32, learning rate 3e-5, class-weighted cross-entropy, inputs truncated to 256 tokens (src/classifier/train_distilbert.py).

Use

from transformers import pipeline

clf = pipeline("text-classification", model="bhanuprakashvangala/pickspin-distilbert-complexity",
               truncation=True, max_length=256)
clf("Prove that the square root of 2 is irrational.")

In the Pick and Spin repository, download it into models/distilbert-complexity-classifier:

hf download bhanuprakashvangala/pickspin-distilbert-complexity --local-dir models/distilbert-complexity-classifier

python src/classifier/evaluate.py there reports the accuracy of the keyword lists, this model and the hybrid classifier on the validation split.

Citation

@INPROCEEDINGS{11662462,
  author={Vangala, Bhanu Prakash and Malik, Tanu},
  booktitle={2026 IEEE 19th International Conference on Cloud Computing (CLOUD)}, 
  title={Pick and Spin: Cold-Start-Aware Routing for Self-Hosted LLM Serving}, 
  year={2026},
  volume={},
  number={},
  pages={388-394},
  abstract={Self-hosting multiple large language models on private or hybrid cloud infrastructure is increasingly appealing for organizations seeking privacy, cost control, and customization, yet no existing system coordinates model selection with model lifecycle management. Serving frameworks such as KServe and Ray Serve deploy models but lack query-aware routing. Routing frameworks such as RouteLLM and FrugalGPT select models but ignore whether those models are warm or cold. We present Pick and Spin, a multi-model orchestration framework built on Kubernetes that unifies these concerns. Pick classifies query complexity using a hybrid keyword-DistilBERT approach and selects models via Thompson Sampling, incorporating lifecycle state from Spin. Spin tracks per-model warm and cold states through a Cold Start State Machine, factors I/O-bound weight-loading penalties into routing decisions, and manages warm pools for frequently accessed models. The two components form a closed feedback loop where routing decisions account for resource availability and orchestration anticipates routing patterns. Deployed as a Helm umbrella chart serving nine LLMs (1Bโ€“27B parameters) via vLLM on NVIDIA A100 GPUs, Pick and Spin achieves 61% GPU utilization versus 31% for static deployment, reduces GPU-hours by 46%, and maintains a 9.2% per-query cold start rate across 310,190 total inference runs-31,019 queries evaluated against all nine models for ground truth, then routed through Pick and Spin-spanning eight standard benchmarks.},
  keywords={Modeling;Spinning;Routing;Printing;Graphics processing units;Large language models;Accuracy;Clouds;Loading;Timing;Cloud Computing;Kubernetes;LLM Inference;Model Orchestration;Cold Start Management;GPU Resource Optimization},
  doi={10.1109/CLOUD72782.2026.00050},
  ISSN={2159-6190},
  month={July},}

The benchmark prompts used for training keep their original licenses.

Downloads last month
41
Safetensors
Model size
67M params
Tensor type
F32
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Model tree for bhanuprakashvangala/pickspin-distilbert-complexity

Finetuned
(12820)
this model