Views
No views yet
Qwen3.5-0.8B (LoRA) → last-token hidden h[1024]
├── routing_head Linear(1024→2) → sigmoid → (p_haiku, p_opus) # BCE
└── token_head Linear(1024→2) → z-scored log1p(output tokens) # MSEadapter_model.safetensors) plus the two
head weights and the token-target normalization stats in heads.pt
(routing_head, token_head, hidden, tok_mean, tok_std).1import torch, torch.nn as nn
2from transformers import AutoModel, AutoTokenizer
3from peft import PeftModel
4from huggingface_hub import hf_hub_download
5
6REPO = "youngryankim/qwen3.5-0.8b-cost-aware-router"
7tok = AutoTokenizer.from_pretrained(REPO)
8backbone = AutoModel.from_pretrained("Qwen/Qwen3.5-0.8B", dtype=torch.bfloat16)
9backbone = PeftModel.from_pretrained(backbone, REPO).eval().cuda()
10
11heads = torch.load(hf_hub_download(REPO, "heads.pt"), map_location="cuda")
12rh = nn.Linear(heads["hidden"], 2).bfloat16().cuda(); rh.load_state_dict(heads["routing_head"]); rh.eval()
13th = nn.Linear(heads["hidden"], 2).bfloat16().cuda(); th.load_state_dict(heads["token_head"]); th.eval()
14mean, std = heads["tok_mean"], heads["tok_std"]
15
16SYS = ("You are a routing model. Read the user query and assess which model can "
17 "answer it and how long each answer will be.")
18
19@torch.no_grad()
20def route(query, input_tokens=200):
21 enc = tok.apply_chat_template([{"role":"system","content":SYS},
22 {"role":"user","content":query}],
23 add_generation_prompt=True, tokenize=True,
24 return_dict=True, return_tensors="pt").to("cuda")
25 h = backbone(**enc).last_hidden_state[:, -1, :]
26 p_h, p_o = torch.sigmoid(rh(h).float())[0].tolist()
27 t = th(h).float()[0].tolist()
28 out_h = torch.expm1(torch.tensor(t[0]*std[0]+mean[0])).item()
29 out_o = torch.expm1(torch.tensor(t[1]*std[1]+mean[1])).item()
30 cost_h = 5e-6*out_h + 1e-6*input_tokens # haiku $5/$1 per Mtok
31 cost_o = 25e-6*out_o + 5e-6*input_tokens # opus $25/$5 per Mtok
32 score = p_o - p_h # routing score
33 cost_aware = score / max(cost_o - cost_h, 1e-6)
34 return dict(p_haiku=p_h, p_opus=p_o, pred_out_h=out_h, pred_out_o=out_o,
35 pred_cost_h=cost_h, pred_cost_o=cost_o, score=score, cost_aware=cost_aware)
36
37print(route("What is 17 * 23?"))score (or cost_aware, under a budget) exceeds a threshold
swept on your validation set.