README.md

---
license: apache-2.0
language:
- en
tags:
- text-classification
- distilbert
- thinking-budget
- reasoning-budget
- agent-routing
- llm-routing
- qwen3
- deepseek-r1
- extended-thinking
- ai-agents
pipeline_tag: text-classification
---

# ThinkingBudgetRouter

A fast, lightweight 3-class classifier that decides **how many thinking tokens a query needs** — before you spend them.

Built on DistilBERT (66M params), fine-tuned to classify any user message into one of three thinking budget tiers:

| Label | Budget | Meaning |
|---|---|---|
| `no_thinking` | 0 tokens | Direct lookup or trivial — answer immediately |
| `brief_thinking` | ~512 tokens | Structured reasoning needed, but not exhaustive |
| `deep_thinking` | 8192+ tokens | Full chain-of-thought required |

## Why This Exists

Modern reasoning models — **Qwen3**, **DeepSeek-R1**, **Claude 3.7 Sonnet** (extended thinking), **Gemini 2.0 Flash Thinking** — all support a configurable thinking budget. But most users either always use maximum thinking (slow, expensive) or never use it (misses hard problems).

**ThinkingBudgetRouter** makes this decision in **~10ms on CPU**, before any tokens are spent reasoning.

## Quick Start

```python
from transformers import pipeline

router = pipeline("text-classification", model="tripathyShaswata/ThinkingBudgetRouter")

# Single prediction
result = router("What is the capital of France?")
print(result)
# [{'label': 'no_thinking', 'score': 0.97}]

# Batch
queries = [
    "What is 15 + 27?",                                          # no_thinking
    "What does HTTP stand for?",                                  # no_thinking
    "Write a Python function to merge two sorted arrays.",        # brief_thinking
    "Explain how Dijkstra's algorithm works step by step.",       # brief_thinking
    "Design a distributed rate limiter for 1 billion users.",     # deep_thinking
    "Prove there are infinitely many prime numbers.",              # deep_thinking
    "Debug this race condition in async code: ...",               # deep_thinking
]
results = router(queries)
for q, r in zip(queries, results):
    print(f"  {r['label']:>16} ({r['score']:.2f}) — {q}")
```

## Use With Qwen3 (Thinking Budget Control)

```python
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

budget_router = pipeline("text-classification", model="tripathyShaswata/ThinkingBudgetRouter")

BUDGET_MAP = {
    "no_thinking":    0,
    "brief_thinking": 512,
    "deep_thinking":  8192,
}

def call_with_budget(user_message: str, model, tokenizer):
    label = budget_router(user_message)[0]["label"]
    budget = BUDGET_MAP[label]

    messages = [{"role": "user", "content": user_message}]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        # Qwen3 thinking budget parameter
        thinking_budget=budget,
    )
    inputs = tokenizer([text], return_tensors="pt")
    output = model.generate(**inputs, max_new_tokens=2048)
    return tokenizer.decode(output[0], skip_special_tokens=True)
```

## Full 3-Stage Agent Pipeline

Pair with [AgentIntentRouter](https://huggingface.co/tripathyShaswata/AgentIntentRouter) and [QueryComplexityRouter](https://huggingface.co/tripathyShaswata/QueryComplexityRouter) for a complete routing stack:

```python
from transformers import pipeline

intent_router     = pipeline("text-classification", model="tripathyShaswata/AgentIntentRouter")
complexity_router = pipeline("text-classification", model="tripathyShaswata/QueryComplexityRouter")
budget_router     = pipeline("text-classification", model="tripathyShaswata/ThinkingBudgetRouter")

BUDGET_MAP = {"no_thinking": 0, "brief_thinking": 512, "deep_thinking": 8192}

def route(user_message: str):
    intent     = intent_router(user_message)[0]
    complexity = complexity_router(user_message)[0]
    budget     = budget_router(user_message)[0]

    print(f"Intent:     {intent['label']} ({intent['score']:.2f})")
    print(f"Complexity: {complexity['label']} ({complexity['score']:.2f})")
    print(f"Budget:     {budget['label']} → {BUDGET_MAP[budget['label']]} tokens")

    if complexity["label"] == "no_llm":
        return handle_with_rules(user_message)

    thinking_tokens = BUDGET_MAP[budget["label"]]

    if complexity["label"] == "small_llm":
        return call_small_model(user_message, thinking_budget=thinking_tokens)
    else:
        return call_large_model(user_message, thinking_budget=thinking_tokens)
```

```
User Message
    │
    ▼
AgentIntentRouter          ← WHAT: code / search / chat / math / ...
    │
    ▼
QueryComplexityRouter       ← WHICH MODEL: no_llm / small_llm / large_llm
    │
    ▼
ThinkingBudgetRouter        ← HOW MUCH THINKING: 0 / 512 / 8192 tokens
    │
    ▼
Model call with exact budget
```

## Budget Labels

### `no_thinking` — 0 tokens
Answer immediately, no chain-of-thought needed:
- *"What is the capital of France?"*
- *"What does HTTP stand for?"*
- *"Is 17 divisible by 3?"*
- *"Convert 100 km to miles."*
- *"What HTTP status code means Not Found?"*

### `brief_thinking` — ~512 tokens
Structured reasoning, but not exhaustive:
- *"Write a Python function to binary search a sorted list."*
- *"What is the difference between REST and GraphQL?"*
- *"Fix this off-by-one error: for i in range(len(arr)): total += arr[i+1]"*
- *"How would you detect a cycle in a linked list?"*
- *"Prove that the sum of two even numbers is even."*

### `deep_thinking` — 8192+ tokens
Full extended chain-of-thought required:
- *"Design a distributed rate limiter for 1 billion users."*
- *"Prove there are infinitely many prime numbers."*
- *"Debug this race condition: two asyncio coroutines sharing a counter."*
- *"Solve: given a matrix of costs, find the minimum-cost path from top-left to bottom-right."*
- *"Compare PostgreSQL, MongoDB, and Cassandra across consistency, latency, and cost."*

## Model Details

| Property | Value |
|---|---|
| Base model | `distilbert-base-uncased` |
| Parameters | ~66M |
| Inference (CPU) | ~10ms |
| Inference (GPU) | ~2ms |
| Input max length | 128 tokens |
| Task | 3-class sequence classification |

## Compatible Reasoning Models

| Model | Thinking control param |
|---|---|
| Qwen3 (all sizes) | `thinking_budget` in chat template |
| DeepSeek-R1 | `/think` token budget |
| Claude 3.7 Sonnet | `thinking.budget_tokens` in API |
| Gemini 2.0 Flash Thinking | `thinkingConfig.thinkingBudget` |