qwen3-4b-router

Qwen3-4B-Router

软件工程 Agent 的下一步路由模型。它在每个主 Agent round 开始前，根据当前可见轨迹预测尚未执行的下一步，输出该步的 phase、phase_detail 和 simple_type，用于决定这一步该走快速模型（simple）还是主力模型（non_simple）。

• 基座：Qwen/Qwen3-4B-Instruct-2507
• 训练：SFT，lr 5e-5 + 2 epoch
• 输出：单行 JSON，三字段 phase / phase_detail / simple_type
• 系统提示词：仓库根目录的 router_system_prompt.txt

模型能做什么

输入是「任务描述 + 早期历史摘要 + 最近若干完整主 round 的轨迹」，输出是对尚未执行的下一步的预测：

{"phase":"localization","phase_detail":"localization_evidence_acquired","simple_type":"non_simple"}

• simple_type = simple：下一步低风险、可直接完成，路由到快速模型；
• simple_type = non_simple：下一步需要较强推理/诊断/权衡，保持主力模型。

测试集指标（temperature=0，threshold=0.5）

快速开始

方式一：transformers 直接加载

pip install transformers torch accelerate

import json
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "zxa11/qwen3-4b-router-best"
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 仓库根目录自带系统提示词，下载后用它的完整内容
system_prompt = open("router_system_prompt.txt").read()

def build_user(task, history_summary, recent_trajectory):
    return (
        f"## Task\n{task}\n\n"
        f"## Earlier history summary\n{history_summary}\n\n"
        f"## Recent trajectory\n{recent_trajectory}"
    )

messages = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": build_user("你的任务描述", "历史摘要", "最近轨迹")},
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=128, do_sample=False)
reply = tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(json.loads(reply.strip()))

方式二：vLLM 加载

pip install vllm

import json
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

model_name = "zxa11/qwen3-4b-router-best"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = LLM(model=model_name, max_model_len=33792, dtype="bfloat16")
sampling_params = SamplingParams(temperature=0, max_tokens=128, stop=["\n"])

system_prompt = open("router_system_prompt.txt").read()
messages = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": build_user("你的任务描述", "历史摘要", "最近轨迹")},
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
out = model.generate([prompt], sampling_params)[0].outputs[0].text.strip()
print(json.loads(out))

部署为服务

vllm serve zxa11/qwen3-4b-router-best --max-model-len 33792 --dtype bfloat16 --port 8000

然后通过 OpenAI 兼容接口调用（base_url = "http://127.0.0.1:8000/v1"）。

推理要点（重要）

1. 温度锁 0：这是确定性路由决策，不要调高温度（会破坏 JSON 格式）。
2. stop=["\n"]：强制在单行 JSON 末尾停住，防止后接噪声。
3. max_tokens=128：输出只有几十 token 的 JSON，够用。
4. 系统提示词必须用仓库里的 router_system_prompt.txt，不要自行改写，否则字段分布会漂移。

阈值路由（可选增强）

默认 simple_type 是模型 argmax 的硬判决。如果希望「不确定就保守走主力模型」，可以取模型在 simple_type 位置对 simple / non 两个候选的 logits 算 P(simple)，再做阈值判断：

• P(simple) > 0.7 → 路由 simple（平衡点，non 召回 ~80%）；
• P(simple) > 0.8 → 路由 simple（保守点，dangerous 压到 4 个，但 simple 召回降到 ~26%）。

用 vLLM 的 logprobs 参数即可拿到这两个候选的 logprob（token：simple=[22944]，non=[6280]）。

输出协议

模型只输出一个完整的单行 JSON 对象，字段顺序固定：

{"phase":"<phase>","phase_detail":"<phase_detail>","simple_type":"<simple 或 non_simple>"}

合法取值见 router_system_prompt.txt（phase 5 类、phase_detail 17 类、simple_type 2 类）。