Build a Self-Improving LLM Router

Fine-tune and deploy an adaptive LLM router with online feedback.

Introduction

30 Second Summary

Sending every question to the most capable option wastes resources. Sending a difficult question to a lighter option can produce a weak answer.

In this project, you will fine-tune Google BERT-Tiny as a cost-aware router for prompts. You will deploy the result as a FastAPI API on FastAPI Cloud with a contextual bandit that learns from feedback.

What You'll Build

You will watch a public endpoint's weak-versus-strong route probabilities shift after feedback for the same prompt.

By the end of this project, you'll have:

  • A fine-tuned transformer router that scores whether each prompt needs the strong route.
  • A quality-cost frontier that reveals how learned routing performs against four simple policies.
  • A live adaptive API where you can inspect the ONNX Runtime score for each prompt. You can watch Vowpal Wabbit update route probabilities after feedback.
  • Secret Mission: A controlled feedback experiment that repeatedly rewards the weak route until its probability distribution changes during one process lifetime.

Are there any prerequisites?

Comfort with Python and basic classification concepts will help. No previous fine-tuning or contextual-bandit experience is required.

Before We Start

Before the hands-on work begins, commit to a strictly free prototype whose fine-tuned BERT-Tiny classifier estimates when a strong model is needed. The deployed Vowpal Wabbit contextual bandit learns which mock route to prefer from feedback only for the lifetime of its running process.

Set Up the Free Training Environment

A self-improving router touches several Python workloads. Training needs a consistent stack from the first notebook cell.

Google Colab provides the hosted notebook plus free GPU access for this project. Pinning one past runtime keeps package conflicts from hiding the routing concepts.

In this step, get ready to:
  • Prepare a Colab notebook with the pinned past runtime.
  • Install the verified training packages.
  • Record the training plus deployment dependency specifications.
Select the pinned GPU runtime

Past runtime 2026.07 supplies Python 3.12.13. It also supplies PyTorch 2.11.0.

A CUDA availability check confirms whether PyTorch can use the selected T4 GPU. This gives the later fine-tuning step hardware acceleration.

  • Visit the official Google Colab product page in your browser.
  • Sign in with your Google account.
  • Open a blank Colab notebook.
  • Open the Runtime menu.
  • Select Change runtime type.
  • Choose 2026.07 from the past-runtime options.
  • Choose T4 GPU as the accelerator.
  • Apply the runtime selections.

Before you check the runtime, do you expect all three results to match the selected environment?

  • Add a code cell with the notebook's code-cell control.
  • Check the runtime versions plus GPU access by running this code:
import platform
import torch

print("Python:", platform.python_version())
print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())

What Does This Check Prove?

  • The platform.python_version() call reads the active Python interpreter version.
  • The torch.__version__ value identifies the PyTorch version supplied by the runtime.
  • The torch.cuda.is_available() result confirms whether PyTorch can access the GPU.

✔️ I see the required versions

Your notebook should report Python 3.12.13 plus PyTorch 2.11.0. It should also report CUDA available: True.

That is the runtime ready. Your notebook can now install the remaining project libraries.

ⓧ I see different versions

The notebook is connected to a different runtime. Switch it to the pinned past runtime before installing anything.

  • Return to the Runtime menu.
  • Select Change runtime type.
  • Choose 2026.07 from the past-runtime options.
  • Apply the runtime selection.
  • Run the runtime check again.

Still Seeing Different Versions?

Confirm that Colab finished reconnecting before you rerun the cell. A stale session can continue showing the previous environment.

Help me correct my Colab runtime versions.

ⓧ CUDA is unavailable

The Python runtime is active without GPU access. Reapply the accelerator selection before continuing.

  • Return to the Runtime menu.
  • Select Change runtime type.
  • Choose T4 GPU as the accelerator.
  • Apply the accelerator selection.
  • Run the runtime check again.

Still Seeing a False Result?

Confirm that the runtime dialog retained the T4 GPU selection. Wait for Colab to reconnect before running the cell again.

Help me enable the T4 GPU in Colab.

Install the pinned project packages

The past runtime already contains Python plus PyTorch. The remaining libraries support dataset loading, transformer fine-tuning, model export, API development, online learning, and HTTP testing.

This installation can take a few minutes while Colab downloads the pinned packages. A busy cell during that time means the environment is still being prepared.

  • Add a code cell below the runtime check.
  • Install the verified packages by running this command:
!pip install datasets==5.1.0 transformers==5.20.1 onnx==1.23.2 onnxscript==0.7.2 onnxruntime==1.31.0 "fastapi[standard]==0.148.0" vowpalwabbit==9.11.9 httpx==0.28.1

What Does This Install?

  • Hugging Face Datasets loads the RouteLLM dataset used for training.
  • Hugging Face Transformers provides the tokenizer plus BERT-Tiny classification model.
  • ONNX stores the exported transformer graph.
  • ONNX Script supports the PyTorch exporter.
  • ONNX Runtime runs the exported model inside the later API.
  • FastAPI provides the API framework.
  • Vowpal Wabbit provides the contextual-bandit policy.
  • HTTPX provides the HTTP client used by the project tests.
  • Wait for the installation cell to finish.
  • Confirm that the cell completes without a package installation failure.

The completed command installs every explicit pin, including ONNX Script 0.7.2. The final environment check confirms the runtime plus the main imported packages.

Did the Installation Stop?

Check that the notebook is still connected to past runtime 2026.07. Restore the connection before rerunning the installation cell.

Help me troubleshoot the pinned Colab package installation.

Define the dependency files and verify the environment

A dependency specification preserves the versions behind a reproducible run. This project keeps separate specifications for Colab training and cloud inference.

  • Expand the notebook's file pane from the left sidebar.
  • Create requirements-colab.txt with the file pane's new-file control.
  • Paste these training dependencies into requirements-colab.txt:
datasets==5.1.0
transformers==5.20.1
onnx==1.23.2
onnxscript==0.7.2
onnxruntime==1.31.0
fastapi[standard]==0.148.0
vowpalwabbit==9.11.9
httpx==0.28.1

Why Keep a Colab Specification?

This file records every package installed on top of the pinned Colab runtime. PyTorch stays outside the file because runtime 2026.07 already supplies version 2.11.0.

  • Save requirements-colab.txt.
  • Confirm that requirements-colab.txt appears in the notebook's file pane.

Missing the Colab Requirements File?

Check that the filename ends with .txt. Compare the saved contents with the reference below.

Help me check my Colab dependency file.

  • Create .python-version with the file pane's new-file control.
  • Paste this Python version into .python-version:
3.12

What Does This Version File Do?

This file pins the later cloud build to Python 3.12. That version supports every package in the deployment specification.

  • Save .python-version.
  • Confirm that .python-version appears in the notebook's file pane.

Missing the Python Version File?

Check that the filename begins with a full stop. Confirm that its only value is 3.12.

Help me check my Python version file.

  • Create requirements.txt with the file pane's new-file control.
  • Paste these deployment dependencies into requirements.txt:
fastapi[standard]==0.148.0
httpx==0.28.1
numpy==2.5.4
onnxruntime==1.31.0
tokenizers==0.23.3
vowpalwabbit==9.11.9

Why Is This Specification Smaller?

The deployed API runs an exported ONNX model. It does not train the transformer.

The deployment specification therefore leaves PyTorch plus the training libraries out of the cloud build. This keeps the runtime dependency set focused on inference plus feedback learning.

  • Save requirements.txt.
  • Confirm that requirements.txt appears in the notebook's file pane.

Do the Runtime Dependencies Look Incomplete?

Check the filename for an accidental extra extension. Compare each saved line with the full reference below.

Help me compare my dependency files.

Your three dependency files should now match these complete references.

✔️ Awesome, I've got everything!

Your training plus deployment dependency specifications are ready. Save each file before continuing.

ⓧ I'd like to double check the full code

3.12

datasets==5.1.0
transformers==5.20.1
onnx==1.23.2
onnxscript==0.7.2
onnxruntime==1.31.0
fastapi[standard]==0.148.0
vowpalwabbit==9.11.9
httpx==0.28.1

fastapi[standard]==0.148.0
httpx==0.28.1
numpy==2.5.4
onnxruntime==1.31.0
tokenizers==0.23.3
vowpalwabbit==9.11.9

Before you run the final check, do you expect every installed package to match its project pin?

  • Add a code cell below the installation cell.
  • Verify the complete Colab environment by running this code:
import platform
from importlib.metadata import version
import torch

print("Python:", platform.python_version())
print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
for package in [
    "datasets",
    "transformers",
    "onnx",
    "onnxruntime",
    "vowpalwabbit",
    "fastapi",
    "httpx",
]:
    print(f"{package}:", version(package))

What Does the Final Check Cover?

The first three results verify the pinned Colab runtime plus GPU access. The package loop reads the installed distribution versions directly from the environment.

✔️ I see every pinned version

The environment is ready when the output contains these results:

  • Python 3.12.13.
  • PyTorch 2.11.0.
  • CUDA availability of True.
  • Datasets 5.1.0.
  • Transformers 5.20.1.
  • ONNX 1.23.2.
  • ONNX Runtime 1.31.0.
  • Vowpal Wabbit 9.11.9.
  • FastAPI 0.148.0.
  • HTTPX 0.28.1.

That is the environment locked in. Your notebook now has the verified stack needed for training, export, API development, and bandit learning.

ⓧ A package version differs

A mismatch means the environment does not yet match the dependency specification. Reinstall the pins inside the selected past runtime.

  • Return to the package installation cell.
  • Run the installation cell again.
  • Wait for the installation cell to finish.
  • Run the final environment check again.

Still Seeing a Version Mismatch?

Confirm that the notebook still uses past runtime 2026.07. A different runtime can restore incompatible preinstalled packages.

Help me fix my Colab package versions.

ⓧ The verification cell stops

A stopped cell usually identifies an incomplete package installation. Use the last printed package to narrow down where the check failed.

  • Read the last package name printed by the verification cell.
  • Return to the package installation cell.
  • Run the installation cell again.
  • Run the final environment check again.

Does the Check Still Stop?

Keep the notebook connected to runtime 2026.07 while repeating the installation. A runtime reconnect can discard an unfinished install.

Help me diagnose the failed environment check.

Your training environment now has the pinned stack plus reproducible dependency files. Next, you will expose the unreliable decisions produced by an untrained routing head.

Expose the Untrained Router's Failure

Your verified Google Colab runtime is ready to load data and run a compact transformer. You can now test what the pretrained model knows before it receives any routing examples.

The pretrained encoder can represent a prompt. However, its classification head has not learned what makes the weak or strong route appropriate. This step makes that missing decision signal visible.

In this step, get ready to:
  • Create the training script with reproducible dataset and model settings.
  • Inspect how the script builds a balanced routing sample.
  • Run the unfitted classifier before any training occurs.
Create the training script

A pretrained Google BERT-Tiny encoder supplies language features for each prompt. BertForSequenceClassification adds the two-label head that this project trains as a router.

  • Switch back to the file browser in your Colab notebook.
  • Create an adaptive_router folder inside /content using the folder control.
  • Create train_router.py inside /content/adaptive_router using the file control.
  • Use the full-code tab below to populate train_router.py.

✔️ Awesome, I've got everything!

Your train_router.py file is ready. Confirm that Colab has saved it inside /content/adaptive_router.

ⓧ I'd like to double check the full code

from argparse import ArgumentParser
from pathlib import Path
import json
import random

from datasets import load_dataset
import numpy as np
import onnxruntime as ort
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from transformers import AutoTokenizer, BertForSequenceClassification

SEED = 42
DATASET_ID = "routellm/gpt4_dataset"
DATASET_REVISION = "7ef62d6"
MODEL_ID = "google/bert_uncased_L-2_H-128_A-2"
SAMPLES_PER_CLASS = 1_500
TRAIN_PER_CLASS = 1_200
MAX_LENGTH = 96
BATCH_SIZE = 32
EPOCHS = 1
LEARNING_RATE = 5e-5
COST_WEIGHT = 0.25
ROUTE_COSTS = {"weak": 0.1, "strong": 1.0}
PROJECT_DIR = Path(__file__).resolve().parent
MODEL_DIR = PROJECT_DIR / "model"
ARTIFACT_DIR = PROJECT_DIR / "artifacts"


def set_seed():
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)


def load_balanced_sample():
    dataset = load_dataset(DATASET_ID, split="train", revision=DATASET_REVISION)
    by_label = {0: [], 1: []}
    for prompt, score in zip(dataset["prompt"], dataset["mixtral_score"]):
        label = 1 if int(score) < 4 else 0
        by_label[label].append(
            {"text": prompt, "label": label, "score": int(score)}
        )

    rng = random.Random(SEED)
    train_rows = []
    validation_rows = []
    for label in [0, 1]:
        rng.shuffle(by_label[label])
        selected = by_label[label][:SAMPLES_PER_CLASS]
        train_rows.extend(selected[:TRAIN_PER_CLASS])
        validation_rows.extend(selected[TRAIN_PER_CLASS:])

    rng.shuffle(train_rows)
    rng.shuffle(validation_rows)
    print("Training rows:", len(train_rows))
    print("Validation rows:", len(validation_rows))
    print(
        "Training labels:",
        {
            "weak_sufficient": sum(row["label"] == 0 for row in train_rows),
            "strong_needed": sum(row["label"] == 1 for row in train_rows),
        },
    )
    return train_rows, validation_rows


def make_loader(rows, tokenizer, shuffle):
    encoded = tokenizer(
        [row["text"] for row in rows],
        truncation=True,
        padding="max_length",
        max_length=MAX_LENGTH,
        return_tensors="pt",
    )
    labels = torch.tensor([row["label"] for row in rows], dtype=torch.long)
    dataset = TensorDataset(encoded["input_ids"], encoded["attention_mask"], labels)
    return DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=shuffle)


def collect_probabilities(model, loader, device):
    model.eval()
    labels = []
    strong_probabilities = []
    losses = []
    with torch.no_grad():
        for input_ids, attention_mask, batch_labels in loader:
            input_ids = input_ids.to(device)
            attention_mask = attention_mask.to(device)
            batch_labels = batch_labels.to(device)
            output = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=batch_labels,
            )
            probabilities = torch.softmax(output.logits, dim=1)[:, 1]
            losses.append(output.loss.item())
            labels.extend(batch_labels.cpu().tolist())
            strong_probabilities.extend(probabilities.cpu().tolist())
    return labels, strong_probabilities, sum(losses) / max(len(losses), 1)


def classification_metrics(labels, strong_probabilities, threshold=0.5):
    predictions = [1 if probability >= threshold else 0 for probability in strong_probabilities]
    true_positive = sum(y == 1 and p == 1 for y, p in zip(labels, predictions))
    false_positive = sum(y == 0 and p == 1 for y, p in zip(labels, predictions))
    false_negative = sum(y == 1 and p == 0 for y, p in zip(labels, predictions))
    accuracy = sum(y == p for y, p in zip(labels, predictions)) / len(labels)
    precision = true_positive / max(true_positive + false_positive, 1)
    recall = true_positive / max(true_positive + false_negative, 1)
    f1 = 2 * precision * recall / max(precision + recall, 1e-12)
    return {
        "accuracy": accuracy,
        "strong_precision": precision,
        "strong_recall": recall,
        "strong_f1": f1,
        "strong_route_rate": sum(predictions) / len(predictions),
    }


def train_one_epoch(model, loader, device):
    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
    model.train()
    total_loss = 0.0
    batches = 0
    for input_ids, attention_mask, labels in loader:
        input_ids = input_ids.to(device)
        attention_mask = attention_mask.to(device)
        labels = labels.to(device)
        optimizer.zero_grad(set_to_none=True)
        output = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
        )
        output.loss.backward()
        optimizer.step()
        total_loss += output.loss.item()
        batches += 1
    return total_loss / max(batches, 1)


def policy_metrics(name, routes, scores):
    quality = [1.0 if route == "strong" else score / 5.0 for route, score in zip(routes, scores)]
    costs = [ROUTE_COSTS[route] for route in routes]
    average_quality = sum(quality) / len(quality)
    average_cost = sum(costs) / len(costs)
    return {
        "name": name,
        "average_quality": average_quality,
        "average_cost": average_cost,
        "utility": average_quality - COST_WEIGHT * average_cost,
        "strong_route_rate": sum(route == "strong" for route in routes) / len(routes),
    }


def build_evaluation(rows, strong_probabilities):
    scores = [row["score"] for row in rows]
    texts = [row["text"] for row in rows]
    rng = random.Random(SEED)
    median_words = sorted(len(text.split()) for text in texts)[len(texts) // 2]
    baselines = [
        policy_metrics("always_weak", ["weak"] * len(rows), scores),
        policy_metrics("always_strong", ["strong"] * len(rows), scores),
        policy_metrics(
            "random",
            [rng.choice(["weak", "strong"]) for _ in rows],
            scores,
        ),
        policy_metrics(
            "prompt_length",
            ["strong" if len(text.split()) >= median_words else "weak" for text in texts],
            scores,
        ),
    ]
    frontier = []
    for threshold in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]:
        routes = [
            "strong" if probability >= threshold else "weak"
            for probability in strong_probabilities
        ]
        row = policy_metrics(f"transformer_{threshold:.1f}", routes, scores)
        row["threshold"] = threshold
        frontier.append(row)
    selected = max(frontier, key=lambda row: row["utility"])
    return baselines, frontier, selected


class RouterOnnx(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, input_ids, attention_mask):
        return self.model(input_ids=input_ids, attention_mask=attention_mask).logits


def export_and_verify(model, tokenizer):
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    tokenizer.save_pretrained(MODEL_DIR)
    tokenizer_file = MODEL_DIR / "tokenizer.json"
    if not tokenizer_file.exists():
        raise RuntimeError("Expected model/tokenizer.json after save_pretrained().")

    model = model.to("cpu").eval()
    wrapper = RouterOnnx(model).eval()
    dummy = tokenizer(
        "Explain why the sky appears blue.",
        truncation=True,
        padding="max_length",
        max_length=MAX_LENGTH,
        return_tensors="pt",
    )
    onnx_program = torch.onnx.export(
        wrapper,
        (dummy["input_ids"], dummy["attention_mask"]),
        input_names=["input_ids", "attention_mask"],
        output_names=["logits"],
        dynamo=True,
    )
    onnx_path = MODEL_DIR / "router.onnx"
    onnx_program.save(onnx_path)

    with torch.no_grad():
        torch_logits = wrapper(dummy["input_ids"], dummy["attention_mask"]).numpy()
    session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
    onnx_logits = session.run(
        ["logits"],
        {
            "input_ids": dummy["input_ids"].numpy(),
            "attention_mask": dummy["attention_mask"].numpy(),
        },
    )[0]
    difference = float(np.max(np.abs(torch_logits - onnx_logits)))
    print("ONNX maximum-logit difference:", difference)
    print("Created", onnx_path.relative_to(PROJECT_DIR))


def main():
    parser = ArgumentParser()
    parser.add_argument("--baseline-only", action="store_true")
    args = parser.parse_args()

    set_seed()
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Device:", device)

    train_rows, validation_rows = load_balanced_sample()
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    train_loader = make_loader(train_rows, tokenizer, shuffle=True)
    validation_loader = make_loader(validation_rows, tokenizer, shuffle=False)
    model = BertForSequenceClassification.from_pretrained(
        MODEL_ID,
        num_labels=2,
        id2label={0: "weak_sufficient", 1: "strong_needed"},
        label2id={"weak_sufficient": 0, "strong_needed": 1},
    ).to(device)

    labels, initial_probabilities, initial_loss = collect_probabilities(
        model, validation_loader, device
    )
    initial_metrics = classification_metrics(labels, initial_probabilities)
    print("Unfitted validation loss:", initial_loss)
    print("Unfitted validation metrics:", json.dumps(initial_metrics, indent=2))
    print(
        "Unfitted sample route:",
        {
            "prompt": validation_rows[0]["text"][:120],
            "strong_probability": initial_probabilities[0],
            "route": "strong" if initial_probabilities[0] >= 0.5 else "weak",
        },
    )
    if args.baseline_only:
        return

    for epoch in range(1, EPOCHS + 1):
        train_loss = train_one_epoch(model, train_loader, device)
        print(f"Epoch {epoch}/{EPOCHS} training loss: {train_loss:.4f}")

    labels, strong_probabilities, validation_loss = collect_probabilities(
        model, validation_loader, device
    )
    metrics = classification_metrics(labels, strong_probabilities)
    baselines, frontier, selected = build_evaluation(
        validation_rows, strong_probabilities
    )
    print("Fine-tuned validation loss:", validation_loss)
    print("Fine-tuned metrics:", json.dumps(metrics, indent=2))
    print("Baselines:")
    for row in baselines:
        print(json.dumps(row))
    print("Threshold frontier:")
    for row in frontier:
        print(json.dumps(row))
    print("Selected threshold policy:", json.dumps(selected, indent=2))

    evaluation = {
        "dataset_id": DATASET_ID,
        "dataset_revision": DATASET_REVISION,
        "model_id": MODEL_ID,
        "sample": {
            "train_rows": len(train_rows),
            "validation_rows": len(validation_rows),
            "balanced": True,
            "seed": SEED,
        },
        "unfitted": initial_metrics,
        "fine_tuned": metrics,
        "baselines": baselines,
        "frontier": frontier,
    }
    policy = {
        "selected_threshold": selected["threshold"],
        "cost_weight": COST_WEIGHT,
        "route_costs": ROUTE_COSTS,
        "max_length": MAX_LENGTH,
        "actions": {"1": "weak", "2": "strong"},
    }
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(evaluation, indent=2), encoding="utf-8"
    )
    (ARTIFACT_DIR / "policy.json").write_text(
        json.dumps(policy, indent=2), encoding="utf-8"
    )
    export_and_verify(model, tokenizer)


if __name__ == "__main__":
    main()

How is the script organized?

  • The constants pin the dataset revision and model checkpoint used by every run.
  • The data helpers create the reproducible balanced split used in this step.
  • The probability helpers evaluate the unfitted classification head without changing its weights.
  • The --baseline-only branch exits before the training and export sections run.
  • Confirm that train_router.py appears inside /content/adaptive_router in the file browser.
Understand the balanced baseline

The public RouteLLM data contains prompts plus a weak-model score. This script maps scores below 4 to the strong_needed label.

Balanced sampling gives both labels equal representation. Seed 42 makes the selected rows and their order reproducible.

  • Find load_balanced_sample() in train_router.py.
  • Confirm that int(score) < 4 maps a row to label 1.
  • Confirm that each label contributes 1,200 training rows.
  • Confirm that each label contributes 300 validation rows.

Why balance the sample?

A balanced sample prevents the larger class from dominating this short teaching run. It produces 2,400 training rows plus 600 validation rows.

These results describe this reproducible sample. They are not a population-level benchmark for every prompt in the full dataset.

Run the unfitted router

The model now has an English-language encoder plus a randomly initialized two-label head. No routing examples have updated that head yet.

Before you run it, do you expect the sample route to provide a trustworthy routing decision?

  • Run the baseline-only path in a new Colab code cell with this command:
!cd /content/adaptive_router && python train_router.py --baseline-only

What does this run prove?

  • The command loads routellm/gpt4_dataset at revision 7ef62d6.
  • The script tokenizes the balanced sample with the pretrained checkpoint.
  • The PyTorch model calculates validation probabilities without running an optimizer.
  • The --baseline-only flag stops execution before fine-tuning begins.

The first download can take a few minutes while Colab retrieves the dataset and model files. The cell may stay busy during that time.

You should see 2,400 training rows plus 600 validation rows.

The output then reports the unfitted validation loss.

You should also see the unfitted classification metrics.

The final baseline output contains one sample prompt with an unreliable route decision.

You have now exposed the exact weakness that fine-tuning must solve. The transformer can process prompts, but its routing head has not learned a useful boundary.

Baseline cell did not finish?

  • Check that the notebook still shows the verified T4 runtime from the previous step.
  • Confirm that train_router.py is stored directly inside /content/adaptive_router.
  • Rerun the package installation cell from the previous step if the output mentions a missing dependency.

Need help diagnosing the baseline run?

Your untrained router has revealed its missing decision signal. Next, you will fine-tune that head and measure whether its routing quality justifies its normalized cost.

Fine-Tune and Evaluate the Cost-Aware Router

Your seeded RouteLLM sample has already exposed the unreliable decisions from Google BERT-Tiny's unfitted classification head. Now you will train Google BERT-Tiny to produce a useful routing signal.

The unfitted head cannot separate easy prompts from prompts that need the strong route. Accuracy alone also hides whether a policy buys enough quality to justify its normalized route cost.

In this step, get ready to:
  • Fine-tune Google BERT-Tiny for one epoch on the balanced training sample.
  • Measure routing policies across quality, cost, utility, and strong-route rate.
  • Produce a verified ONNX export with its runtime tokenizer.
Turn labels into a trained signal

Fine-tuning uses the labeled training split to update the pretrained encoder and its classification head. The validation split stays outside those gradient updates.

The script evaluates all 600 held-out prompts after the epoch. These metrics reveal how the learned routing signal behaves on prompts it did not train against.

What do the routing metrics show?

  • Accuracy measures the share of validation prompts assigned the correct routing label.
  • The strong_precision value measures how often a predicted strong route was actually labeled strong_needed.
  • The strong_recall value measures how many strong_needed prompts the classifier found.
  • The strong_f1 value balances precision with recall.
  • The strong_route_rate value shows how frequently the policy chooses the more expensive route.
  • Return to the Google Colab notebook from earlier.
  • Confirm the notebook still shows the verified T4 runtime.
Run training and compare routing policies

A cost-aware evaluation gives every routing policy the same test. The script calls its combined score utility.

How is policy utility measured?

Average quality rewards routes that preserve answer quality. Average cost applies the project-defined normalized costs of 0.1 for weak routes and 1.0 for strong routes.

Utility subtracts 0.25 times average cost from average quality. The highest-utility threshold becomes the selected offline policy.

  • The always_weak baseline shows the cheapest fixed policy.
  • The always_strong baseline shows the highest-cost fixed policy.
  • The random baseline chooses between the two routes with the seeded random generator.
  • The prompt_length baseline uses the median prompt length as its decision boundary.
  • The transformer frontier tests thresholds from 0.1 through 0.9.

This is the slowest notebook cell in the training workflow. Expect it to stay active for several minutes while the T4 processes the training batches.

  • Predict whether the most accurate routing policy will also earn the highest utility.
  • Train the router and trigger its evaluation pipeline by running this cell:
!cd /content/adaptive_router && python train_router.py

What does this cell do?

  • Running without --baseline-only lets the script continue past the unfitted evaluation.
  • The script completes one training epoch before calculating held-out classification metrics.
  • It compares four baseline policies with the transformer threshold frontier.
  • It writes the evaluation report to artifacts/evaluation.json.
  • It writes the selected threshold policy to artifacts/policy.json.
  • It exports the trained classifier and verifies the exported logits.
  • Confirm the output includes Epoch 1/1 training loss:.
  • Find Fine-tuned validation loss: after the completed epoch.
  • Find Fine-tuned metrics: with accuracy and the strong-route metrics.
  • Count four JSON rows beneath Baselines:.
  • Confirm Threshold frontier: lists nine threshold rows.
  • Find Selected threshold policy: below the frontier.

Did the training cell stop early?

  • Confirm the notebook is still connected to the verified T4 runtime.
  • Rerun the pinned installation cell from earlier if Colab reset the runtime.
  • Rerun the training cell after the package imports succeed.

Help me diagnose the failed training run.

That is the core modeling loop complete. Your router now has a trained signal and a cost-aware threshold selected from the validation frontier.

Export the router for runtime inference

ONNX packages the fine-tuned computation graph for deployment without the PyTorch training runtime. ONNX Runtime scores the same dummy input through the exported graph.

The script compares those logits with the PyTorch logits. A small maximum difference confirms that the export preserves the trained computation closely.

  • Predict which model artifacts the cloud API needs to score a prompt.
  • Scroll to the final lines of the completed training cell.
  • Confirm ONNX maximum-logit difference: is followed by a small numeric value.
  • Confirm the next line reads Created model/router.onnx.
  • Expand /content/adaptive_router/model in Colab's file browser.
  • Confirm router.onnx and tokenizer.json are listed.
  • Expand /content/adaptive_router/artifacts in Colab's file browser.
  • Confirm evaluation.json and policy.json are listed.

You should now see the verified ONNX graph beside its tokenizer. You should also see the evaluation report beside the selected threshold policy.

✔️ Awesome, I've got everything!

Great. Double-check that the completed cell is saved in your Colab notebook.

ⓧ I'd like to double check the full code

from argparse import ArgumentParser
from pathlib import Path
import json
import random

from datasets import load_dataset
import numpy as np
import onnxruntime as ort
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from transformers import AutoTokenizer, BertForSequenceClassification

SEED = 42
DATASET_ID = "routellm/gpt4_dataset"
DATASET_REVISION = "7ef62d6"
MODEL_ID = "google/bert_uncased_L-2_H-128_A-2"
SAMPLES_PER_CLASS = 1_500
TRAIN_PER_CLASS = 1_200
MAX_LENGTH = 96
BATCH_SIZE = 32
EPOCHS = 1
LEARNING_RATE = 5e-5
COST_WEIGHT = 0.25
ROUTE_COSTS = {"weak": 0.1, "strong": 1.0}
PROJECT_DIR = Path(__file__).resolve().parent
MODEL_DIR = PROJECT_DIR / "model"
ARTIFACT_DIR = PROJECT_DIR / "artifacts"


def set_seed():
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)


def load_balanced_sample():
    dataset = load_dataset(DATASET_ID, split="train", revision=DATASET_REVISION)
    by_label = {0: [], 1: []}
    for prompt, score in zip(dataset["prompt"], dataset["mixtral_score"]):
        label = 1 if int(score) < 4 else 0
        by_label[label].append(
            {"text": prompt, "label": label, "score": int(score)}
        )

    rng = random.Random(SEED)
    train_rows = []
    validation_rows = []
    for label in [0, 1]:
        rng.shuffle(by_label[label])
        selected = by_label[label][:SAMPLES_PER_CLASS]
        train_rows.extend(selected[:TRAIN_PER_CLASS])
        validation_rows.extend(selected[TRAIN_PER_CLASS:])

    rng.shuffle(train_rows)
    rng.shuffle(validation_rows)
    print("Training rows:", len(train_rows))
    print("Validation rows:", len(validation_rows))
    print(
        "Training labels:",
        {
            "weak_sufficient": sum(row["label"] == 0 for row in train_rows),
            "strong_needed": sum(row["label"] == 1 for row in train_rows),
        },
    )
    return train_rows, validation_rows


def make_loader(rows, tokenizer, shuffle):
    encoded = tokenizer(
        [row["text"] for row in rows],
        truncation=True,
        padding="max_length",
        max_length=MAX_LENGTH,
        return_tensors="pt",
    )
    labels = torch.tensor([row["label"] for row in rows], dtype=torch.long)
    dataset = TensorDataset(encoded["input_ids"], encoded["attention_mask"], labels)
    return DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=shuffle)


def collect_probabilities(model, loader, device):
    model.eval()
    labels = []
    strong_probabilities = []
    losses = []
    with torch.no_grad():
        for input_ids, attention_mask, batch_labels in loader:
            input_ids = input_ids.to(device)
            attention_mask = attention_mask.to(device)
            batch_labels = batch_labels.to(device)
            output = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=batch_labels,
            )
            probabilities = torch.softmax(output.logits, dim=1)[:, 1]
            losses.append(output.loss.item())
            labels.extend(batch_labels.cpu().tolist())
            strong_probabilities.extend(probabilities.cpu().tolist())
    return labels, strong_probabilities, sum(losses) / max(len(losses), 1)


def classification_metrics(labels, strong_probabilities, threshold=0.5):
    predictions = [1 if probability >= threshold else 0 for probability in strong_probabilities]
    true_positive = sum(y == 1 and p == 1 for y, p in zip(labels, predictions))
    false_positive = sum(y == 0 and p == 1 for y, p in zip(labels, predictions))
    false_negative = sum(y == 1 and p == 0 for y, p in zip(labels, predictions))
    accuracy = sum(y == p for y, p in zip(labels, predictions)) / len(labels)
    precision = true_positive / max(true_positive + false_positive, 1)
    recall = true_positive / max(true_positive + false_negative, 1)
    f1 = 2 * precision * recall / max(precision + recall, 1e-12)
    return {
        "accuracy": accuracy,
        "strong_precision": precision,
        "strong_recall": recall,
        "strong_f1": f1,
        "strong_route_rate": sum(predictions) / len(predictions),
    }


def train_one_epoch(model, loader, device):
    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
    model.train()
    total_loss = 0.0
    batches = 0
    for input_ids, attention_mask, labels in loader:
        input_ids = input_ids.to(device)
        attention_mask = attention_mask.to(device)
        labels = labels.to(device)
        optimizer.zero_grad(set_to_none=True)
        output = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
        )
        output.loss.backward()
        optimizer.step()
        total_loss += output.loss.item()
        batches += 1
    return total_loss / max(batches, 1)


def policy_metrics(name, routes, scores):
    quality = [1.0 if route == "strong" else score / 5.0 for route, score in zip(routes, scores)]
    costs = [ROUTE_COSTS[route] for route in routes]
    average_quality = sum(quality) / len(quality)
    average_cost = sum(costs) / len(costs)
    return {
        "name": name,
        "average_quality": average_quality,
        "average_cost": average_cost,
        "utility": average_quality - COST_WEIGHT * average_cost,
        "strong_route_rate": sum(route == "strong" for route in routes) / len(routes),
    }


def build_evaluation(rows, strong_probabilities):
    scores = [row["score"] for row in rows]
    texts = [row["text"] for row in rows]
    rng = random.Random(SEED)
    median_words = sorted(len(text.split()) for text in texts)[len(texts) // 2]
    baselines = [
        policy_metrics("always_weak", ["weak"] * len(rows), scores),
        policy_metrics("always_strong", ["strong"] * len(rows), scores),
        policy_metrics(
            "random",
            [rng.choice(["weak", "strong"]) for _ in rows],
            scores,
        ),
        policy_metrics(
            "prompt_length",
            ["strong" if len(text.split()) >= median_words else "weak" for text in texts],
            scores,
        ),
    ]
    frontier = []
    for threshold in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]:
        routes = [
            "strong" if probability >= threshold else "weak"
            for probability in strong_probabilities
        ]
        row = policy_metrics(f"transformer_{threshold:.1f}", routes, scores)
        row["threshold"] = threshold
        frontier.append(row)
    selected = max(frontier, key=lambda row: row["utility"])
    return baselines, frontier, selected


class RouterOnnx(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, input_ids, attention_mask):
        return self.model(input_ids=input_ids, attention_mask=attention_mask).logits


def export_and_verify(model, tokenizer):
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    tokenizer.save_pretrained(MODEL_DIR)
    tokenizer_file = MODEL_DIR / "tokenizer.json"
    if not tokenizer_file.exists():
        raise RuntimeError("Expected model/tokenizer.json after save_pretrained().")

    model = model.to("cpu").eval()
    wrapper = RouterOnnx(model).eval()
    dummy = tokenizer(
        "Explain why the sky appears blue.",
        truncation=True,
        padding="max_length",
        max_length=MAX_LENGTH,
        return_tensors="pt",
    )
    onnx_program = torch.onnx.export(
        wrapper,
        (dummy["input_ids"], dummy["attention_mask"]),
        input_names=["input_ids", "attention_mask"],
        output_names=["logits"],
        dynamo=True,
    )
    onnx_path = MODEL_DIR / "router.onnx"
    onnx_program.save(onnx_path)

    with torch.no_grad():
        torch_logits = wrapper(dummy["input_ids"], dummy["attention_mask"]).numpy()
    session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
    onnx_logits = session.run(
        ["logits"],
        {
            "input_ids": dummy["input_ids"].numpy(),
            "attention_mask": dummy["attention_mask"].numpy(),
        },
    )[0]
    difference = float(np.max(np.abs(torch_logits - onnx_logits)))
    print("ONNX maximum-logit difference:", difference)
    print("Created", onnx_path.relative_to(PROJECT_DIR))


def main():
    parser = ArgumentParser()
    parser.add_argument("--baseline-only", action="store_true")
    args = parser.parse_args()

    set_seed()
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Device:", device)

    train_rows, validation_rows = load_balanced_sample()
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    train_loader = make_loader(train_rows, tokenizer, shuffle=True)
    validation_loader = make_loader(validation_rows, tokenizer, shuffle=False)
    model = BertForSequenceClassification.from_pretrained(
        MODEL_ID,
        num_labels=2,
        id2label={0: "weak_sufficient", 1: "strong_needed"},
        label2id={"weak_sufficient": 0, "strong_needed": 1},
    ).to(device)

    labels, initial_probabilities, initial_loss = collect_probabilities(
        model, validation_loader, device
    )
    initial_metrics = classification_metrics(labels, initial_probabilities)
    print("Unfitted validation loss:", initial_loss)
    print("Unfitted validation metrics:", json.dumps(initial_metrics, indent=2))
    print(
        "Unfitted sample route:",
        {
            "prompt": validation_rows[0]["text"][:120],
            "strong_probability": initial_probabilities[0],
            "route": "strong" if initial_probabilities[0] >= 0.5 else "weak",
        },
    )
    if args.baseline_only:
        return

    for epoch in range(1, EPOCHS + 1):
        train_loss = train_one_epoch(model, train_loader, device)
        print(f"Epoch {epoch}/{EPOCHS} training loss: {train_loss:.4f}")

    labels, strong_probabilities, validation_loss = collect_probabilities(
        model, validation_loader, device
    )
    metrics = classification_metrics(labels, strong_probabilities)
    baselines, frontier, selected = build_evaluation(
        validation_rows, strong_probabilities
    )
    print("Fine-tuned validation loss:", validation_loss)
    print("Fine-tuned metrics:", json.dumps(metrics, indent=2))
    print("Baselines:")
    for row in baselines:
        print(json.dumps(row))
    print("Threshold frontier:")
    for row in frontier:
        print(json.dumps(row))
    print("Selected threshold policy:", json.dumps(selected, indent=2))

    evaluation = {
        "dataset_id": DATASET_ID,
        "dataset_revision": DATASET_REVISION,
        "model_id": MODEL_ID,
        "sample": {
            "train_rows": len(train_rows),
            "validation_rows": len(validation_rows),
            "balanced": True,
            "seed": SEED,
        },
        "unfitted": initial_metrics,
        "fine_tuned": metrics,
        "baselines": baselines,
        "frontier": frontier,
    }
    policy = {
        "selected_threshold": selected["threshold"],
        "cost_weight": COST_WEIGHT,
        "route_costs": ROUTE_COSTS,
        "max_length": MAX_LENGTH,
        "actions": {"1": "weak", "2": "strong"},
    }
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(evaluation, indent=2), encoding="utf-8"
    )
    (ARTIFACT_DIR / "policy.json").write_text(
        json.dumps(policy, indent=2), encoding="utf-8"
    )
    export_and_verify(model, tokenizer)


if __name__ == "__main__":
    main()

Your offline router is trained, evaluated, and packaged for inference. Next, you will connect these artifacts to a feedback-learning API.

Add the Feedback-Learning API

Your fine-tuned Google BERT-Tiny router can estimate whether a prompt needs the strong route. Its threshold policy stays fixed when users report that a route performed poorly.

In this step, you will connect ONNX Runtime scoring to a FastAPI service. A Vowpal Wabbit contextual bandit will record each decision before learning from its feedback.

In this step, get ready to:
  • Build an API that scores prompts with the exported ONNX model.
  • Add weak-versus-strong exploration with logged selection probabilities.
  • Test the feedback loop before packaging the deployment bundle.
Build the adaptive API

A contextual bandit chooses an action using the context available at decision time. Learning requires the service to retain the chosen action plus its selection probability until feedback arrives.

Why Use Two Fixed Actions?

The available routes always mean weak or strong. The fixed meanings make Workspace(cb_explore=2, epsilon=0.2, quiet=True) a direct fit.

The transformer supplies the prompt context. The bandit adapts the probability assigned to each route while preserving exploration.

  • Create main.py inside /content/adaptive_router using the Colab file pane.

You should now see main.py beside train_router.py in the project directory.

  • Copy the complete reference from the second tab below into main.py.

✔️ Awesome, I've got everything!

Your main.py file should define the shared router state plus all four API endpoints.

ⓧ I'd like to double check the full code

from pathlib import Path
import json
import random
import threading
from uuid import uuid4

from fastapi import FastAPI
import numpy as np
import onnxruntime as ort
from pydantic import BaseModel
from tokenizers import Tokenizer
from vowpalwabbit import Workspace

PROJECT_DIR = Path(__file__).resolve().parent
MODEL_DIR = PROJECT_DIR / "model"
POLICY_PATH = PROJECT_DIR / "artifacts" / "policy.json"
ACTIONS = {1: "weak", 2: "strong"}
MOCK_RESPONSES = {
    "weak": "Mock weak backend: concise low-cost response selected.",
    "strong": "Mock strong backend: detailed high-capability response selected.",
}


class RouteRequest(BaseModel):
    prompt: str


class FeedbackRequest(BaseModel):
    decision_id: str
    quality: float


class RouterState:
    def __init__(self):
        self.policy = json.loads(POLICY_PATH.read_text(encoding="utf-8"))
        self.max_length = int(self.policy["max_length"])
        self.cost_weight = float(self.policy["cost_weight"])
        self.route_costs = {
            route: float(cost) for route, cost in self.policy["route_costs"].items()
        }
        self.tokenizer = Tokenizer.from_file(str(MODEL_DIR / "tokenizer.json"))
        pad_id = self.tokenizer.token_to_id("[PAD]")
        if pad_id is None:
            raise RuntimeError("The saved tokenizer has no [PAD] token.")
        self.tokenizer.enable_truncation(max_length=self.max_length)
        self.tokenizer.enable_padding(
            length=self.max_length,
            pad_id=pad_id,
            pad_token="[PAD]",
        )
        self.session = ort.InferenceSession(
            str(MODEL_DIR / "router.onnx"),
            providers=["CPUExecutionProvider"],
        )
        self.vw = Workspace(cb_explore=2, epsilon=0.2, quiet=True)
        self.random = random.Random(42)
        self.lock = threading.Lock()
        self.pending = {}
        self.route_counts = {"weak": 0, "strong": 0}
        self.feedback_count = 0
        self.total_feedback_cost = 0.0
        self.closed = False
        self._warm_start()

    def _features(self, strong_probability, prompt_length):
        return (
            "|context "
            f"strong_probability:{strong_probability:.6f} "
            f"prompt_length:{prompt_length}"
        )

    def _warm_start(self):
        threshold = float(self.policy["selected_threshold"])
        for _ in range(5):
            for integer_score in range(11):
                score = integer_score / 10
                action = 2 if score >= threshold else 1
                features = self._features(score, 50)
                self.vw.learn(f"{action}:0.0:1.0 {features}")

    def _transformer_score(self, prompt):
        encoding = self.tokenizer.encode(prompt)
        inputs = {
            "input_ids": np.asarray([encoding.ids], dtype=np.int64),
            "attention_mask": np.asarray(
                [encoding.attention_mask], dtype=np.int64
            ),
        }
        logits = self.session.run(["logits"], inputs)[0]
        shifted = logits - logits.max(axis=1, keepdims=True)
        probabilities = np.exp(shifted) / np.exp(shifted).sum(
            axis=1, keepdims=True
        )
        return float(probabilities[0, 1])

    def route(self, prompt):
        prompt = prompt.strip()
        if not prompt:
            return {"error": "prompt must not be empty"}
        prompt_length = len(prompt.split())
        with self.lock:
            strong_probability = self._transformer_score(prompt)
            features = self._features(strong_probability, prompt_length)
            raw_probabilities = [
                max(float(value), 0.0) for value in self.vw.predict(features)
            ]
            if len(raw_probabilities) != 2:
                raise RuntimeError("Expected exactly two VW action probabilities.")
            total = sum(raw_probabilities)
            probabilities = (
                [0.5, 0.5]
                if total <= 0
                else [value / total for value in raw_probabilities]
            )
            action = self.random.choices(
                [1, 2], weights=probabilities, k=1
            )[0]
            route = ACTIONS[action]
            selection_probability = probabilities[action - 1]
            decision_id = str(uuid4())
            self.pending[decision_id] = {
                "action": action,
                "route": route,
                "probability": selection_probability,
                "features": features,
            }
            self.route_counts[route] += 1
        return {
            "decision_id": decision_id,
            "route": route,
            "mock_response": MOCK_RESPONSES[route],
            "transformer_strong_probability": strong_probability,
            "action_probabilities": {
                "weak": probabilities[0],
                "strong": probabilities[1],
            },
            "selection_probability": selection_probability,
            "persistence": "in-memory",
        }

    def feedback(self, decision_id, quality):
        if quality < 0.0 or quality > 1.0:
            return {"error": "quality must be between 0.0 and 1.0"}
        with self.lock:
            decision = self.pending.pop(decision_id, None)
            if decision is None:
                return {"error": "unknown or already-used decision_id"}
            combined_cost = (
                1.0 - quality
                + self.cost_weight * self.route_costs[decision["route"]]
            )
            learning_example = (
                f"{decision['action']}:{combined_cost:.6f}:"
                f"{decision['probability']:.6f} {decision['features']}"
            )
            self.vw.learn(learning_example)
            self.feedback_count += 1
            self.total_feedback_cost += combined_cost
        return {
            "status": "learned",
            "decision_id": decision_id,
            "route": decision["route"],
            "quality": quality,
            "combined_cost": combined_cost,
            "logged_probability": decision["probability"],
        }

    def stats(self):
        with self.lock:
            average_cost = (
                self.total_feedback_cost / self.feedback_count
                if self.feedback_count
                else 0.0
            )
            return {
                "route_counts": dict(self.route_counts),
                "feedback_count": self.feedback_count,
                "average_feedback_cost": average_cost,
                "pending_decisions": len(self.pending),
                "exploration": "epsilon-greedy",
                "epsilon": 0.2,
                "persistence": "in-memory",
            }

    def close(self):
        if not self.closed:
            self.vw.finish()
            self.closed = True


STATE = RouterState()
app = FastAPI(title="Adaptive LLM Router Prototype", version="1.0.0")


@app.get("/health")
def health():
    return {
        "status": "ok",
        "model": "BERT-Tiny ONNX router",
        "bandit": "Vowpal Wabbit cb_explore",
        "persistence": "in-memory",
    }


@app.post("/route")
def route(request: RouteRequest):
    return STATE.route(request.prompt)


@app.post("/feedback")
def feedback(request: FeedbackRequest):
    return STATE.feedback(request.decision_id, request.quality)


@app.get("/policy-stats")
def policy_stats():
    return STATE.stats()

What Does This File Do?

  • The service loads artifacts/policy.json to reuse the selected threshold plus the project-defined route costs.
  • The tokenizer prepares fixed-length inputs for model/router.onnx.
  • The bandit predicts one probability for each fixed action.
  • The feedback method learns from the selected action plus its logged probability.
  • Save main.py.
  • Confirm the bottom of main.py contains the four endpoint decorators.

You should see request models followed by one shared RouterState instance. The endpoint functions connect that state to routing plus feedback requests.

How Does the Learning Loop Work?

  • The _transformer_score() method converts ONNX logits into a strong-route probability.
  • The route() method combines that score with prompt length before requesting two action probabilities.
  • Each pending decision stores the selected action plus its probability.
  • The feedback() method converts reported quality into a learning cost.

Expect Ephemeral Learning

The policy workspace plus pending decisions live in process memory. Restarting the process clears the feedback history.

The lock protects this shared state from concurrent changes inside one process. It does not create durable storage.

Does the File Look Incomplete?

  • Check that main.py starts with the import from pathlib.
  • Check that RouterState contains methods for scoring plus feedback.
  • Confirm the file is stored directly inside /content/adaptive_router.

Compare my adaptive-router API with the complete reference.

Test one feedback cycle

A smoke test checks the complete scoring plus exploration plus feedback loop. It also proves that the route response exposes probabilities for both actions.

  • Create smoke_test.py inside /content/adaptive_router using the Colab file pane.

You should now see the empty smoke_test.py file beside main.py.

  • Add one complete route plus feedback check by pasting this into smoke_test.py:
import json
from main import STATE

health = {
    "status": "ok",
    "model": "BERT-Tiny ONNX router",
    "persistence": "in-memory",
}
decision = STATE.route(
    "Prove that every finite subgroup of the multiplicative group of a field is cyclic."
)
feedback = STATE.feedback(decision["decision_id"], quality=0.8)
stats = STATE.stats()

print("Health:", json.dumps(health, indent=2))
print("Route:", json.dumps(decision, indent=2))
print("Feedback:", json.dumps(feedback, indent=2))
print("Policy statistics:", json.dumps(stats, indent=2))

if stats["feedback_count"] != 1:
    raise RuntimeError("The feedback update was not recorded.")
if set(decision["action_probabilities"]) != {"weak", "strong"}:
    raise RuntimeError("The route response is missing action probabilities.")
STATE.close()

What Does This Test Check?

  • The field-theory prompt creates one decision through the shared router state.
  • The test submits a quality score of 0.8 for that exact decision.
  • The first assertion confirms that one feedback event was recorded.
  • The second assertion confirms that the response contains both fixed actions.
  • Save smoke_test.py.

Before you run the test, predict whether the final feedback count will be zero or one.

  • Run the complete feedback cycle with this Colab command:
!cd /content/adaptive_router && python smoke_test.py

What Should You See?

The output starts with a healthy state. The route section contains probability values for weak plus strong.

The feedback section reports learned. The policy statistics show feedback_count equal to 1.

That is the full online-learning loop working. Your router recorded the decision probability before learning from the reported outcome.

Smoke Test Did Not Finish?

  • Confirm artifacts/policy.json still exists inside /content/adaptive_router.
  • Confirm model/router.onnx remains inside the model directory.
  • Confirm model/tokenizer.json remains inside the model directory.

Help me debug my adaptive-router smoke test.

Package the deployment bundle

The cloud deployment needs the API source plus every model artifact used at runtime. An archive keeps the working project together for the next step.

  • Create bundle_project.py inside /content/adaptive_router using the Colab file pane.

You should now see bundle_project.py beside the API plus smoke test files.

  • Add the archive logic by pasting this into bundle_project.py:
from pathlib import Path
import shutil

project_dir = Path(__file__).resolve().parent
archive_base = project_dir.parent / "adaptive_router_bundle"
archive_path = shutil.make_archive(
    str(archive_base),
    "zip",
    root_dir=project_dir,
)
print("Created", archive_path)

What Does This Script Do?

  • The script resolves the current adaptive_router directory.
  • The archive base points one level above that directory.
  • The zip format packages everything under the project directory.
  • The final print exposes the completed archive path.
  • Save bundle_project.py.

Before you run the script, predict the path it will print for the finished archive.

  • Create the deployment archive by running this Colab command:
!cd /content/adaptive_router && python bundle_project.py

What Should You See?

The output reports Created /content/adaptive_router_bundle.zip. The archive packages every file under /content/adaptive_router.

That includes the API plus smoke test. It also includes the requirements plus the tokenizer plus the ONNX model plus the saved policy artifacts.

  • Confirm adaptive_router_bundle.zip appears under /content in the Colab file pane.

You have packaged the complete adaptive router for deployment. The bundle now carries the model plus the feedback-learning API into the next step.

Archive Missing From the File Pane?

  • Refresh the Colab file pane.
  • Confirm the command printed /content/adaptive_router_bundle.zip.
  • Confirm bundle_project.py is stored inside /content/adaptive_router.

Help me find my generated deployment bundle.

Your API can now score prompts plus learn from route feedback. Next, you will deploy the bundle before testing the same behavior through a public URL.

Deploy the Adaptive Router to FastAPI Cloud

Your adaptive router now passes its smoke test in Google Colab. The next goal is a public URL that anyone can call.

FastAPI Cloud gives your API a public home on its free Hobby plan. The deployment keeps ONNX Runtime inference plus the mock backends you already built.

The Vowpal Wabbit policy stores its learning state in process memory. A one-replica maximum prevents separate processes from holding divergent policies.

Cold starts can still clear the policy state. This makes the live service a demo-ready prototype.

In this step, get ready to:
  • Prepare the extracted router bundle in a Python environment on your Mac.
  • Confirm FastAPI discovers the application through the local documentation page.
  • Publish the router with one process-local policy on FastAPI Cloud.
Test the bundle on your Mac

A cloud build is easier to debug after the extracted bundle works on your Mac. A virtual environment keeps its runtime packages separate from your other Python projects.

Finding the generated archive is a little fiddly because it sits outside the project folder in Colab. You will move it into your Mac's Downloads folder before using macOS Terminal.

  • In Colab's left sidebar, select the folder icon.
  • Find adaptive_router_bundle.zip in /content.
  • Download adaptive_router_bundle.zip to your Mac.
  • Open your browser's downloads list.
  • Reveal adaptive_router_bundle.zip in Finder.
  • Double-click adaptive_router_bundle.zip to extract it.
  • Rename the extracted adaptive_router_bundle folder to adaptive_router.
  • Press Cmd+Space to open Spotlight.
  • Type Terminal.
  • Press Enter to open Terminal.
  • Enter the extracted adaptive_router folder by running this command:
cd ~/Downloads/adaptive_router

Why move into this folder?

The cd command makes ~/Downloads/adaptive_router the current directory. Every setup command now applies to the extracted deployment bundle.

The project pins Python 3.12 through .python-version. Check your Mac before creating the deployment environment.

  • Check the available Python version by running this command:
python --version

What does this check prove?

The command prints the Python interpreter selected by your shell. The next commands require Python 3.12.

✔️ I see the required version

Your Mac is using Python 3.12. That interpreter can create the matching deployment environment.

ⓧ I see an older version

Your current Python cannot create the required 3.12 environment. Install the official Python 3.12.10 macOS build before continuing.

  • Open the official Python 3.12.10 release page.
  • Choose Download macOS 64-bit universal2 installer.
  • Double-click the downloaded .pkg file.
  • Complete the installer with its default options.
  • Close Terminal after the installation finishes.
  • Use Spotlight to reopen Terminal.
  • Return to ~/Downloads/adaptive_router.
  • Confirm the updated version by running this command:
python --version

What should the command confirm?

The reopened shell should now report Python 3.12. The installer adds the Python executable to your shell path.

ⓧ Command not found

Your shell cannot find Python. Install the official Python 3.12.10 macOS build before creating the virtual environment.

  • Open the official Python 3.12.10 release page.
  • Choose Download macOS 64-bit universal2 installer.
  • Double-click the downloaded .pkg file.
  • Complete the installer with its default options.
  • Close Terminal after the installation finishes.
  • Use Spotlight to reopen Terminal.
  • Return to ~/Downloads/adaptive_router.
  • Confirm Python is available by running this command:
python --version

What should the command confirm?

The reopened shell should report Python 3.12. This confirms the installer made Python available to Terminal.

The dependency install can take a few minutes while your Mac prepares the runtime packages. A quiet terminal during part of the installation does not mean it has stalled.

  • Create the deployment environment by running these commands:
python -m venv .venv
source .venv/bin/activate
python -m pip install "fastapi[standard]==0.148.0" httpx==0.28.1 numpy==2.5.4 onnxruntime==1.31.0 tokenizers==0.23.3 vowpalwabbit==9.11.9

What do these commands prepare?

  • The first command creates an isolated .venv environment with your Python interpreter.
  • The second command makes the environment's python executable active in the current terminal.
  • The final command installs the pinned packages required by main.py.

The local development command discovers the root main.py file automatically. It keeps Terminal busy until you stop the server with Ctrl+C.

Before you run it, do you expect FastAPI to find the existing app instance without any configuration changes?

  • Start the local API by running this command:
fastapi dev

What does this command start?

The command imports app from the root main.py file. It serves the interactive documentation at http://127.0.0.1:8000/docs.

You should see output confirming that the development server is listening on http://127.0.0.1:8000. This proves FastAPI discovered the application.

  • Switch back to the browser you used for Colab.
  • Open a new browser tab.
  • Enter http://127.0.0.1:8000/docs in the address bar.
  • Press Enter.
  • Expand the GET /health operation.
  • Send the health request using the operation controls.

You should receive a response with status set to ok. The response also names the BERT-Tiny ONNX router plus the in-memory policy.

  • Return to Terminal.
  • Press Ctrl+C to stop the development server.

Local documentation not loading?

  • Confirm the terminal prompt begins with (.venv).
  • Confirm main.py sits directly inside ~/Downloads/adaptive_router.
  • Confirm the package installation completed before starting the server.

Help me troubleshoot the local adaptive router.

Create the FastAPI Cloud deployment

Direct deployment sends the extracted directory to FastAPI Cloud. The platform reads .python-version plus requirements.txt before discovering the root main.py application.

Cloud signup can feel like a billing checkpoint. This app uses the $0/month Hobby plan with no credit card required.

The first deployment opens a browser for login or account creation. The flow also asks you to select a team before creating the app.

The upload can sit quietly while FastAPI Cloud transfers the ONNX model and builds the runtime. Keep Terminal open until the deployment result appears.

  • Send the extracted adaptive_router directory to FastAPI Cloud by running this command:
fastapi deploy . --large-file-threshold 25

What does this deployment command do?

The command deploys the current directory. The --large-file-threshold 25 flag raises the warning threshold to 25 MB for the model bundle.

  • Complete the FastAPI Cloud login or account-creation flow in the browser.
  • Select a team for the deployment.
  • Create the Hobby app when prompted.
  • Return to Terminal after the browser flow completes.
  • Save the public app URL here: https://your-app.fastapicloud.dev.

FastAPI Cloud now reports a successful deployment with a public app URL. Your extracted project is linked to that cloud app.

Deployment did not finish?

  • Confirm the local development server was stopped before deploying.
  • Confirm requirements.txt remains inside ~/Downloads/adaptive_router.
  • Confirm your Mac stays connected while the model bundle uploads.

Help me troubleshoot the cloud deployment.

Limit the policy to one replica and verify it

Each running process owns its own policy plus pending decisions. A maximum of one replica keeps live requests on one Vowpal Wabbit workspace.

  • Return to the FastAPI Cloud browser tab opened during deployment.
  • Select your deployed app.
  • Open the app's Scaling page.
  • Set the Replicas maximum to 1.
  • Click Update.

That is the cloud constraint handled. Your app now avoids multiple live processes learning separate policies.

What can still reset?

One replica keeps current requests on one process. Scale-to-zero or a process restart creates a fresh process that clears the in-memory policy.

  • Set ROUTER_URL by replacing https://your-app.fastapicloud.dev with https://your-app.fastapicloud.dev before running this command:
export ROUTER_URL="https://your-app.fastapicloud.dev"

What does this command store?

The command stores your public app address in ROUTER_URL for the current terminal session. The next request can reuse that address.

FastAPI Cloud may need to wake the Hobby app after it scales to zero. The first public request can take longer while that process starts.

Before you send the request, do you expect the public health response to match the local one?

  • Check the public health endpoint by running this command:
curl "$ROUTER_URL/health"

What does this request prove?

The curl request calls the deployed /health endpoint. A successful response proves the public process loaded the API plus its model artifacts.

You should see "status":"ok", "model":"BERT-Tiny ONNX router", "bandit":"Vowpal Wabbit cb_explore", and "persistence":"in-memory" in the response.

  • Return to the browser tab for your public app.
  • Append /docs to the public app URL.
  • Expand the POST /route operation.
  • Enable the request controls.
  • Replace the prompt value with Prove that every finite subgroup of the multiplicative group of a field is cyclic..
  • Send the route request.

You should receive a decision_id plus weak and strong action probabilities. Your adaptive router is now responding through its public FastAPI Cloud URL.

Public request not responding?

  • Wait for the scale-to-zero app to finish starting before repeating the health request.
  • Confirm ROUTER_URL contains the public app URL without /docs.
  • Confirm the Scaling page still shows a maximum of 1 replica.

Help me troubleshoot the public router.

Secret mission

Make Feedback Reverse a Route Preference

One feedback update proves that the learning path executes. This experiment repeats one fixed prompt for 40 rounds to show the live weak-route probability move during one process lifetime.

Clean Up Your Resources

Clean Up Your Resources

Your FastAPI Cloud Hobby app currently costs $0/month. Your Colab runtime and local files have no ongoing financial cost.

Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

Resources you used:

  • A FastAPI Cloud Hobby deployment with a one-replica maximum and process-local Vowpal Wabbit policy state.
  • A Google Colab notebook with an active T4 runtime and a cached RouteLLM dataset download.
  • The ~/Downloads/adaptive_router folder containing your virtual environment, API, scripts, reports, policy, tokenizer, and ONNX model.
  • The ~/Downloads/adaptive_router_bundle.zip deployment archive.

Keep everything running

No action needed. Choose this if you are still testing the adaptive router or demonstrating its feedback loop.

  • Keep the Hobby deployment configured with a one-replica maximum.
  • Keep the Colab notebook for future training runs.
  • Retain the ~/Downloads/adaptive_router folder for local development.
  • Retain ~/Downloads/adaptive_router_bundle.zip as a deployment backup.
  • Check the FastAPI Cloud pricing page before future testing.

The deployed app automatically scales to zero while idle. A cold start or restart can reset its in-memory learning state.

Pause - I'll come back to this later

Pause active use to release the Colab GPU while preserving your cloud deployment and project artifacts.

  • Stop sending requests to the public app through ROUTER_URL.
  • Return to the Colab notebook from earlier.
  • End the active T4 runtime using the notebook runtime controls.
  • Keep the ~/Downloads/adaptive_router folder in your Downloads directory.
  • Keep ~/Downloads/adaptive_router_bundle.zip in your Downloads directory.

FastAPI Cloud scales the idle Hobby app to zero. The next request can trigger a cold start that resets the in-memory policy.

Ending the Colab runtime releases its GPU session. It also clears the runtime's cached dataset files.

Delete - I don't want to use this again

This option permanently removes the public deployment and your project files. The adaptive-router API stops responding after the cloud deployment is cancelled.

  • Return to the FastAPI Cloud account dashboard from earlier.
  • Select the Hobby deployment created for this project.
  • Cancel the active deployment using the available deployment control.
  • Return to the public /health endpoint in your browser.

You should no longer receive a response containing "status":"ok" from the deployment.

Can't Find the Deployment Control?

FastAPI Cloud does not currently publish the exact single-app deletion label. Contact support@fastapicloud.com if the account dashboard does not show a cancellation control.

Help me write a FastAPI Cloud deployment cancellation request.

  • Remove the local FastAPI Cloud link from ~/Downloads/adaptive_router by running these commands:
cd ~/Downloads/adaptive_router
fastapi cloud unlink

What Do These Commands Do?

The first command returns to your extracted project folder. The second command removes the local .fastapicloud directory.

Unlinking leaves the cloud deployment active. Cancelling the deployment first prevents the public app from remaining online.

Deleting the local files is permanent. The next commands target only the extracted project folder and its deployment archive.

  • Remove the extracted project folder and deployment archive by running these commands:
cd ~/Downloads
rm -rf adaptive_router
rm -f adaptive_router_bundle.zip

What Gets Deleted?

The first removal command deletes the adaptive_router folder. The second removal command deletes adaptive_router_bundle.zip.

This removes your local virtual environment, model, tokenizer, evaluation report, policy file, scripts, and dependency files.

  • Return to the Colab notebook from earlier.
  • End the active T4 runtime using the notebook runtime controls.
  • Confirm the notebook no longer shows an attached runtime.
  • Move the notebook to trash using its file controls.
  • Permanently remove the notebook using your Google account's trash controls.

Cleanup complete. Your deployment, Colab runtime, cached dataset, notebook, local project folder, and deployment archive are removed.

Nice Work!

Nice Work!

You did it! Your adaptive LLM router is live on FastAPI Cloud. It combines a fine-tuned transformer scorer with feedback-driven routing in a public API.

You've learned how to:

  • Fine-tune Google BERT-Tiny on a balanced RouteLLM sample after exposing the unfitted classification head's unreliable decisions.
  • Compare learned routing with four baselines through a quality-cost frontier. Export the classifier to ONNX. Verify its logits against PyTorch.
  • Deploy a public FastAPI API that scores prompts through ONNX Runtime. Add a Vowpal Wabbit contextual bandit that learns from logged feedback. Constrain the process-local policy to one cloud replica.
  • Secret Mission: Run a controlled online-learning experiment that shifts weak-versus-strong route probabilities across 40 feedback rounds. Explain why the change lasts only for the current process lifetime.

Ready to quiz yourself?