How one question travels through the agent loop

71class Family:
72    key: str
73    name: str
74    model_types: tuple[str, ...]  # config.json `model_type` values of this family
75
76    # Reading the model's output.
77    parse: Callable[..., dict]  # (raw output, {tool name: parameters schema}, language of its errors) -> thinking, tool_calls, content
78    thought_prefix: str  # a prompt that ends with this has the model start inside its thinking
79    channel_switch: dict[str, str]  # tokens that switch what the model is writing: thought, tool_call, content
80    stop_tokens: dict[str, dict[str, str]]  # tokens that end a generation -> language -> what stopping there means
81
82    # Writing the conversation back for the chat template.
83    reasoning_key: str  # the assistant-message field the chat template reads earlier thinking from
84    text_with_calls: bool  # whether text written before a tool call survives re-rendering the history
85
86    # Showing it on the page; the text is in both of the page's languages (i18n.py).
87    user_turn: str  # where a user message starts in a prompt; the step card shows the newest one
88    special_groups: dict[str, list]  # language -> (title, subtitle, ((token, name, what it is for), ...)) for each group of control tokens
89    thinking_help: dict[str, str]  # what the thinking toggle does to the prompt
90    tools_note: dict[str, str]  # how the tool schemas appear in the prompt
91    model_note: dict[str, str] = field(default_factory=dict)
92    sampling: dict | None = None  # recommended temperature / top_p / top_k when a checkpoint ships none
93    byte_level: bool = False  # the vocabulary spells bytes as characters, GPT-2 style (BYTE_OF); otherwise it is SentencePiece's
94    processor_kwargs: dict = field(default_factory=dict)  # passed to the processor along with images, audio and video
95    file_notes: dict[str, FileNote] = field(default_factory=dict)  # a file in the checkpoint's folder -> what it holds in this family
96    drafter_file_notes: dict[str, FileNote] = field(default_factory=dict)  # the same for the folder of its MTP drafter (DRAFTERS)
97
98    def show_piece(self, piece: str) -> str:
99        """A vocabulary entry as the text it stands for, which is what a token's popover shows: ␣ for a space,
100        ↵ for a newline, and a byte that is only part of a character as <0xE5>, the way SentencePiece writes one."""
101        if self.byte_level and all(c in BYTE_OF for c in piece):
102            # surrogateescape keeps each byte that is not whole UTF-8 as U+DC80–U+DCFF, so it can be written in hex.
103            text = bytes(BYTE_OF[c] for c in piece).decode("utf-8", errors="surrogateescape")
104            text = "".join(f"<0x{ord(c) - 0xDC00:02X}>" if "\udc80" <= c <= "\udcff" else c for c in text)
105        else:
106            text = piece.replace("▁", " ")  # SentencePiece marks a leading space with ▁
107        return text.replace(" ", "␣").replace("\n", "↵")
108
109    @cached_property
110    def special_tokens(self) -> dict[str, dict[str, tuple[str, str]]]:
111        """Control token -> language -> (name, what it is for)."""
112        out: dict[str, dict[str, tuple[str, str]]] = {}
113        for lang, groups in self.special_groups.items():
114            for _, _, rows in groups:
115                for piece, name, desc in rows:
116                    out.setdefault(piece, {})[lang] = (name, desc)
117        return out
118
119    @cached_property
120    def special_re(self) -> re.Pattern:
121        """Matches the control tokens in a prompt, longest first so <|image|> is not read as <|image>."""
122        pieces = sorted(self.special_tokens, key=len, reverse=True)
123        return re.compile("|".join(re.escape(p) for p in pieces))

144self.family = families.for_model_type(AutoConfig.from_pretrained(_download(model_id, "config.json")).model_type)
145files = self.folder = _download(model_id)
146# The processor holds the tokenizer and, for a model that reads more than text, what turns files into its inputs.
147self.processor = AutoProcessor.from_pretrained(files)
148self.tokenizer = getattr(self.processor, "tokenizer", self.processor)
149self.model = AutoModelForMultimodalLM.from_pretrained(files, dtype=self.dtype, device_map=self.device).eval()
150self.modalities = self._modalities()
151self.media_token_ids = {kind: getattr(self.model.config, f"{kind}_token_id", None) for kind in self.modalities}
152self.load_seconds = time.perf_counter() - start
153
154# The generation config's stop tokens, plus the family's: a checkpoint may list only some of them.
155# Qwen3.5 ships no generation_config.json, and its config.json names <|endoftext|> but not <|im_end|>,
156# which ends every chat turn.
157eos = self.model.generation_config.eos_token_id
158self.eos_ids = set(eos if isinstance(eos, list) else [eos]) | self.token_ids(self.family.stop_tokens)
159# Control tokens count as special even when the tokenizer does not mark them (Qwen's <think>).
160self.special_ids = set(self.tokenizer.all_special_ids) | self.eos_ids | self.token_ids(self.family.special_tokens)

268QWEN3_5 = Family(
269    key="qwen3_5",
270    name="Qwen3.5",
271    model_types=("qwen3_5",),
272    parse=parse_output,
273    thought_prefix="<think>\n",
274    channel_switch={"<think>": "thought", "</think>": "content", "<tool_call>": "tool_call", "</tool_call>": "content"},
275    stop_tokens={
276        "<|im_end|>": {"zh": "訊息結束；回覆和工具呼叫都停在這裡", "en": "end of message; replies and tool calls both stop here"},
277        "<|endoftext|>": {"zh": "文字結束", "en": "end of text"},
278    },
279    reasoning_key="reasoning_content",
280    text_with_calls=True,  # the template keeps a message's text in front of its <tool_call>
281    user_turn="<|im_start|>user",
282    special_groups=GROUPS,
283    thinking_help={
284        "zh": "chat template 在 assistant 訊息開頭先寫好 `<think>`，模型先寫推理；關掉時寫一段空的 `<think></think>`。"
285        "0.8B 官方預設不開，開了容易一直想下去停不下來。",
286        "en": "The chat template opens the assistant message with `<think>`, so the model reasons first; off, it writes an empty `<think></think>`. "
287        "Off by default for the 0.8B model, which once it starts thinking often cannot stop.",
288    },
289    tools_note={
290        "zh": "chat template 把每個工具的 JSON Schema 原樣放進 system turn 的 `<tools>…</tools>`，後面附上 XML 呼叫格式的說明。",
291        "en": "The chat template puts each tool's JSON Schema as it is in `<tools>…</tools>` in the system turn, followed by how to write a call in XML.",
292    },
293    model_note={
294        "zh": "混合架構：24 層裡每 4 層有 3 層是 Gated DeltaNet（線性 attention，狀態大小固定，不隨脈絡變長），1 層是 gated attention。"
295        "在 transformers 裡沒裝 flash-linear-attention 時，它們用 PyTorch 參考實作，prefill 比較慢。",
296        "en": "A hybrid: of every 4 of its 24 layers, 3 are Gated DeltaNet (linear attention, whose state has a fixed size and does not grow with the context) and 1 is gated attention. "
297        "In transformers, without flash-linear-attention installed, they run the PyTorch reference implementation, so prefill is slower.",
298    },
299    # The checkpoint ships no generation_config.json; these are the model card's non-thinking text settings.
300    sampling={"temperature": 1.0, "top_p": 1.0, "top_k": 20},
301    byte_level=True,
302    # Cap each video frame at the pixels the reference implementation (qwen-vl-utils) allows; transformers asks
303    # for this to be said outright until it becomes the default.
304    processor_kwargs={"cap_pixels_per_frame": True},
305    file_notes=FILES,
306)

50def _download(model_id: str, only: str | None = None) -> str:
51    """The model's folder in MODELS_DIR, after fetching the files it lacks (only `only`, if given), at the revision
52    families.REVISIONS pins it to.
53
54    from_pretrained(local_files_only=True) is no test of the cache: once config.json is cached, a tokenizer whose
55    own files are missing raises a ValueError (Gemma 4) or loads empty (Qwen3.5) rather than an OSError.
56    snapshot_download lists the model's files on the Hub and fetches only those the cache lacks; when the Hub
57    can't be reached it returns the cached folder, so a downloaded model still starts offline."""
58    if os.path.isdir(model_id):  # a checkpoint on disk rather than a Hugging Face id
59        return model_id
60    return snapshot_download(model_id, revision=families.REVISIONS.get(model_id), allow_patterns=only, cache_dir=MODELS_DIR)

62class FileNote:
63    """What one file in a checkpoint's folder holds in this family. The Model tab shows the note next to the file,
64    with the values in this checkpoint of `fields`, dotted paths into a JSON file (text_config.sliding_window)."""
65
66    text: dict[str, str]  # language -> the note
67    fields: tuple[str, ...] = ()

133@st.cache_resource(max_entries=1, show_spinner=False)  # one model in memory: switching models frees the last one
134def load_llm(model_id: str, url: str | None) -> Backend:
135    llm = vllm_backend.VLLMBackend(model_id, url) if url else LocalLLM(model_id)
136    llm.warmup()
137    return llm
138
139
140def release_model():
141    """Let the model in memory go now, rather than once the next one has loaded beside it."""
142    load_llm.clear()
143    gc.collect()
144    empty_cache()

188def current_llm() -> Backend:
189    model_id = st.session_state.model_id
190    if st.session_state.pop("stop_vllm", False):  # transformers was just picked again: vLLM's memory goes back first
191        with st.spinner(tr("停止 vLLM 的容器…", "Stopping vLLM's container…")):
192            vllm_docker.stop()
193    st.session_state.pop("engine_picked", None)  # vLLM is stopped or about to start: from now on this tab follows it too
194    url = vllm_server(model_id)["url"] if docker_vllm() else VLLM_URL
195    if url:
196        loading = tr(f"連線到 {url} 的 vLLM 伺服器，載入 {model_id} 的 tokenizer…", f"Connecting to the vLLM server at {url} and loading {model_id}'s tokenizer…")
197    else:
198        loading = tr(f"載入 {model_id}（第一次要下載權重；之後約 10–20 秒）…這時候切換的設定，載入完才會生效。",
199                     f"Loading {model_id} (the first time downloads the weights; after that it takes about 10–20 s)… What you switch meanwhile takes effect once it has loaded.")
200    with st.spinner(loading):
201        llm = load_llm(model_id, url)
202    if st.session_state.sampling_for != model_id:  # a model just chosen starts from its own recommended decoding, MTP off
203        # but on vLLM in Docker, where MTP is how the server was just started
204        for key, value in {**llm.info["defaults"], **({} if docker_vllm() else {"mtp": False})}.items():
205            st.session_state[key] = value
206        st.session_state.sampling_for = model_id
207    if st.session_state.mtp and llm.drafter_id and llm.drafter is None:
208        with st.spinner(tr(f"載入 MTP 草稿模型 {llm.drafter_id}（第一次要下載約 190 MB）…", f"Loading the MTP drafter {llm.drafter_id} (the first time downloads about 190 MB)…")):
209            llm.load_drafter()
210    return llm

325def start_run(llm: Backend, text: str, uploads: list[tuple[str, bytes]] = ()) -> bool:
326    """Ask the question in a new run, with `uploads` (file name, contents) saved in the conversation's folder first."""
327    slot = generation_slot()
328    if not slot.lock.acquire(blocking=False):
329        st.toast(tr("另一個 agent 正在執行（可能在別的分頁），請稍等。", "Another agent is running (maybe in another tab). Please wait."))
330        return False
331    try:
332        attachments = [media.save(chat_folder(), name, data, llm.modalities) for name, data in uploads]
333        run = LiveRun(llm, st.session_state.history, text, attachments, chat_folder(), current_settings(llm), slot)
334    except Exception:
335        slot.lock.release()
336        raise
337    st.session_state.runs.append(run)
338    st.session_state.view = len(st.session_state.runs) - 1
339    return True

292def current_settings(llm: Backend) -> RunSettings:
293    s = st.session_state
294    return RunSettings(
295        sampling=Sampling(s.temperature, s.top_p, s.top_k, s.max_new_tokens, mtp=s.mtp and (llm.drafter is not None or docker_vllm())),
296        max_steps=s.max_steps,
297        thinking=s.thinking,
298        use_tools=s.use_tools,
299        tools=offered_tools(),
300        approve_tools=s.approve_tools,
301        lang=s.lang,
302    )

227class LiveRun:
228    """One user turn: the harness thread and the events it has emitted so far."""
229
230    def __init__(self, llm: Backend, history: list[dict], text: str, attachments: list[media.Attachment], files: Path,
231                 settings: RunSettings, slot: Slot):
232        self.text = text
233        self.attachments = attachments
234        self.events: list[dict] = []
235        self.absorbed = False
236        self.run = Run(llm, history, text, settings, self.events.append, attachments, files)
237        self._slot = slot
238        slot.run = self
239        self.thread = threading.Thread(target=self._work, daemon=True)
240        self.thread.start()
241
242    def _work(self):
243        try:
244            self.run.run()
245        finally:
246            if self._slot.run is self:
247                self._slot.run = None  # do not keep the model alive after a switch to another one
248            self._slot.lock.release()
249
250    @property
251    def live(self) -> bool:
252        return self.thread.is_alive()
253
254    def data(self) -> dict:
255        return render.build(self.events[:])

66def __init__(self, llm: Backend, history: list[dict], user_text: str, settings: RunSettings, emit: Callable[[dict], None],
67             attachments: list[media.Attachment] = (), files: Path | None = None):
68    self.id = uuid.uuid4().hex[:12]
69    self.llm = llm
70    self.settings = settings
71    self.messages = [*history, {"role": "user", "content": media.user_content(user_text, list(attachments))}]
72    self.files = files  # the conversation's folder, which run_python's sandbox sees as /data
73    self.tools = {name: t for name, t in TOOLS.items() if name in settings.tools} if settings.use_tools else {}
74    self._emit = emit
75    self._start = time.perf_counter()
76    self._stop = threading.Event()
77    self._timed_out = False
78    self._decisions: dict[str, bool] = {}
79    self._decided = threading.Condition()

98def run(self):
99    s = self.settings
100    family = self.llm.family  # how this model writes its turns, thinking and tool calls
101    tools = [t.schema() for t in self.tools.values()] or None
102    schemas = {t.name: t.parameters for t in TOOLS.values()}
103    system = {"role": "system", "content": SYSTEM_WITH_TOOLS if tools else SYSTEM_NO_TOOLS}
104    self.emit(
105        "run_start", run_id=self.id, model=self.llm.model_id, family=family.key,
106        settings=_settings_dict(s), tools=list(self.tools),
107    )
108
109    seen: list[tuple[str, str]] = []
110    format_errors = 0
111    prev_sequence = ""
112    stop_reason, final = "max_steps", ""

83def schema(self) -> dict:
84    return {
85        "type": "function",
86        "function": {"name": self.name, "description": self.description, "parameters": self.parameters},
87    }

115for step in range(1, s.max_steps + 1):
116    self.emit("step_start", step=step)
117    prompt = self.llm.render([system, *self.messages], tools, s.thinking)
118    encoded = self.llm.encode(prompt, media.parts(self.messages))
119    # Media sit in user messages, all before what is new since the last step; so the new part is the same
120    # text in the prompt and in what the model reads, only further along by what the media expanded into.
121    added_from = len(encoded.text) - len(prompt) + common_prefix(prompt, prev_sequence) if prev_sequence else 0
122    self.emit(
123        "context",
124        step=step,
125        prompt=encoded.text,
126        n_tokens=encoded.n_tokens,
127        added_from=added_from,
128        media_tokens=encoded.media_tokens,
129        context_length=self.llm.context_length,
130        roles=_role_counts([system, *self.messages]),
131    )

248def render(self, messages: list[dict], tools: list[dict] | None, thinking: bool, generation_prompt: bool = True) -> str:
249    """The chat template turns messages and tool schemas into the one string the model reads.
250    With `generation_prompt` it also opens the model's turn, so the model writes next."""
251    return self.tokenizer.apply_chat_template(
252        messages, tools=tools or None, tokenize=False, add_generation_prompt=generation_prompt, enable_thinking=thinking
253    )

57def common_prefix(a: str, b: str) -> int:
58    n = min(len(a), len(b))
59    i = 0
60    while i < n and a[i] == b[i]:
61        i += 1
62    return i

195def _modalities(self) -> tuple[str, ...]:
196    """What the checkpoint reads besides text: each kind needs an encoder in the config and a part of the processor
197    to feed it. It is asked of the checkpoint, not the family, since checkpoints of one family can differ."""
198    config, processor = self.model.config, self.processor
199    reads = {
200        "image": getattr(config, "vision_config", None) is not None and hasattr(processor, "image_processor"),
201        "audio": getattr(config, "audio_config", None) is not None and hasattr(processor, "feature_extractor"),
202        "video": getattr(config, "video_token_id", None) is not None and hasattr(processor, "video_processor"),
203    }
204    return tuple(kind for kind, ok in reads.items() if ok)

96def user_content(text: str, attachments: list[Attachment]) -> str | list[dict]:
97    """The user's message: the text alone, or the media first (as their models are trained to read them), then the
98    text, and a note on the files only run_python can open: where they are and how a text file begins. The note is
99    for the model, so it is in English."""
100    if not attachments:
101        return text
102    files = [a for a in attachments if a.kind == "file"]
103    described = [f"/data/{a.name} ({fmt_size(a.size)})" + (f", whose first line is: {head}" if (head := preview(a.path)) else "") for a in files]
104    note = "Attached for run_python, which reads them:\n" + "\n\n".join(described) if files else ""
105    words = "\n\n".join(t for t in (text, note) if t)
106    return [a.part() for a in attachments if a.kind != "file"] + ([{"type": "text", "text": words}] if words else [])

255def encode(self, prompt: str, parts: list[dict]) -> Encoded:
256    """The prompt as the model reads it. Without media that is the prompt's tokens. With media (`parts`, the
257    conversation's media in the order of their placeholders), the processor replaces each placeholder with the
258    tokens that image, clip or video takes, and makes the inputs of the encoders whose output goes there."""
259    if not parts:
260        ids = self.tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids
261        return Encoded(prompt, ids, ids.shape[-1])
262    files = {kind: [p["path"] for p in parts if p["type"] == kind] for kind in media.MEDIA}
263    inputs = {}
264    if files["image"]:
265        inputs["images"] = [media.load_image(path) for path in files["image"]]
266    if files["audio"]:
267        inputs["audio"] = [media.load_audio(path, self.processor.feature_extractor.sampling_rate) for path in files["audio"]]
268    if files["video"]:
269        videos = [media.load_video(path) for path in files["video"]]
270        inputs["videos"] = [frames for frames, _ in videos]
271        inputs["videos_kwargs"] = {"video_metadata": [metadata for _, metadata in videos]}
272        # A processor that takes a set number of frames (Gemma 4: 32) refuses a video with fewer; it gets all there are.
273        wanted, fewest = getattr(self.processor.video_processor, "num_frames", None), min(len(frames) for frames, _ in videos)
274        if wanted and fewest < wanted:
275            inputs["videos_kwargs"]["num_frames"] = fewest
276    out = self.processor(text=prompt, return_tensors="pt", **inputs, **self.family.processor_kwargs)
277    ids = out.pop("input_ids")
278    out.pop("attention_mask", None)  # one sequence, no padding: every position counts, as with text
279    counts = {kind: int((ids == token).sum()) for kind, token in self.media_token_ids.items() if token is not None}
280    text = self.tokenizer.decode(ids[0], skip_special_tokens=False)
281    return Encoded(text, ids, ids.shape[-1], dict(out.to(self.device, dtype=self.dtype)), {k: n for k, n in counts.items() if n})

360def _encoded_media(self, inputs: dict) -> dict:
361    """Run the vision encoder over the prompt's images and videos before the prefill, as generate() does, and pass
362    its output in their place. Leaving it to the forward pass breaks videos on Gemma 4: transformers 5.20's
363    Gemma4ForConditionalGeneration takes num_frames_per_video but does not hand it on to the video encoder."""
364    if not inputs or "mm_encoder_outputs" not in inspect.signature(self.model.forward).parameters:
365        return inputs
366    base, encoded = self.model.base_model, {}
367    for kind in ("image", "video"):
368        encoder = getattr(base, f"get_{kind}_features", None)
369        args = {k: inputs[k] for k in inspect.signature(encoder).parameters if k in inputs} if encoder else {}
370        if any(k.startswith("pixel_values") for k in args):
371            encoded[kind] = encoder(**args, return_dict=True)
372    if not encoded:
373        return inputs
374    # The pixels go; the rest stays, such as the grid sizes Qwen3.5's positions are computed from.
375    return {**{k: v for k, v in inputs.items() if not k.startswith("pixel_values")}, "mm_encoder_outputs": encoded}

133done = None
134for ev in self.llm.generate(encoded, s.sampling, self._stop.is_set):
135    kind = ev.pop("type")
136    if kind == "done":
137        done = ev
138    else:
139        self.emit(kind, step=step, **ev)
140raw = done.pop("raw")
141self.emit("model_end", step=step, raw=raw, **done)

283def generate(self, prompt: Encoded, sampling: Sampling, should_stop: Callable[[], bool]) -> Iterator[dict]:
284    """Yield a prefill event, one event per token, then a done event."""
285    ids = prompt.ids.to(self.device)
286    channel = "thought" if prompt.text.endswith(self.family.thought_prefix) else "content"
287    drafting = sampling.mtp and self.drafter is not None
288    mtp = {"drafted": 0, "accepted": 0, "forwards": 0} if drafting else None
289    generated: list[int] = []
290    text = ""
291    finish = "max_new_tokens"
292
293    with torch.inference_mode():
294        # Models with multimodal RoPE (Qwen3.5) keep the position offset of the last prompt's images on the model
295        # and add it in every later forward pass; each prompt starts from its own.
296        if hasattr(self.model.base_model, "rope_deltas"):
297            self.model.base_model.rope_deltas = None
298        start = time.perf_counter()
299        out = self._forward(ids, None, 1, drafting, **self._encoded_media(prompt.media))
300        prefill_ms = (time.perf_counter() - start) * 1000
301        yield {"type": "prefill", "n_tokens": ids.shape[-1], "ms": round(prefill_ms, 1)}

303decode_start = time.perf_counter()
304tokens = self._drafted(out, sampling, mtp) if drafting else self._one_at_a_time(out, sampling)
305for token_id, details in tokens:
306    if should_stop():
307        finish = "user_stopped"
308        break
309    generated.append(token_id)
310
311    piece = self.piece(token_id)
312    special = token_id in self.special_ids
313    if piece in self.family.channel_switch:
314        channel = self.family.channel_switch[piece]
315        token_channel = "control"
316    else:
317        token_channel = "control" if special else channel
318
319    # Decode the whole sequence and diff it, so characters split across byte tokens come out whole.
320    full = self.tokenizer.decode(generated, skip_special_tokens=False)
321    delta = "" if full.endswith("�") else full[len(text):]
322    if delta:
323        text = full
324    yield {
325        "type": "token",
326        "index": len(generated) - 1,
327        "id": token_id,
328        "piece": piece,
329        "text": delta,
330        "special": special,
331        "channel": token_channel,
332        **details,
333    }

335if token_id in self.eos_ids:
336    finish = "stop_token"
337    break
338if len(generated) == sampling.max_new_tokens:
339    break

377def _one_at_a_time(self, out, s: Sampling) -> Iterator[tuple[int, dict]]:
378    """Plain decoding: sample a token, feed it back with the cache for the next position's logits, repeat."""
379    while True:
380        token_id, details = self._sample(out.logits[0, -1].float().cpu(), s)
381        yield token_id, details
382        out = self._forward(torch.tensor([[token_id]], device=self.device), out.past_key_values, 1)

341decode_s = time.perf_counter() - decode_start
342yield {
343    "type": "done",
344    "finish": finish,
345    "stop_piece": self.piece(generated[-1]) if finish == "stop_token" else None,
346    "n_tokens": len(generated),
347    "decode_ms": round(decode_s * 1000, 1),
348    "tok_per_s": round(len(generated) / decode_s, 1) if decode_s > 0 else None,
349    "mtp": mtp,
350    "raw": self.tokenizer.decode(generated, skip_special_tokens=False),
351}

431def _sample(self, logits: torch.Tensor, s: Sampling, draft: int | None = None) -> tuple[int, dict]:
432    """Pick the next token and report the model's distribution and what the settings did to it.
433    A `draft` (the MTP drafter's guess) is kept with the probability the settings give it, which in greedy
434    decoding means when it is the top token; otherwise the pick comes from the other candidates. Either way
435    each token comes out exactly as often as without a drafter."""
436    probs = torch.softmax(logits, -1)  # the model's own distribution (T = 1, nothing cut off)
437    # Equally likely tokens go lowest id first, the one argmax and so generate() takes. topk leaves their order open
438    # (on the CPU it can list the higher id first), and in bfloat16 two logits are equal more often than one would think.
439    top = sorted(zip(*probs.topk(TOP_N)), key=lambda c: (-float(c[0]), int(c[1])))
440
441    if s.temperature <= 0:  # T -> 0 is greedy decoding
442        chosen = int(top[0][1])
443        kept = 1
444        p_sampled = 1.0
445    else:
446        scaled = logits / s.temperature
447        if s.top_k > 0:
448            cand_logits, cand_ids = scaled.topk(min(s.top_k, scaled.numel()))
449        else:
450            cand_logits, cand_ids = scaled.sort(descending=True)
451        cand_probs = torch.softmax(cand_logits, -1)
452        # Nucleus: keep the smallest prefix whose probability reaches top_p (always at least one token).
453        kept = int(((cand_probs.cumsum(-1) - cand_probs) < s.top_p).sum())
454        cand_ids, cand_probs = cand_ids[:kept], cand_probs[:kept] / cand_probs[:kept].sum()
455        weights = cand_probs
456        at = (cand_ids == draft).nonzero().flatten() if draft is not None else []
457        if len(at) and float(torch.rand(())) < float(cand_probs[at[0]]):
458            pick = int(at[0])
459        else:
460            if len(at):  # the guess was turned down: draw from the other candidates
461                weights = cand_probs.clone()
462                weights[at[0]] = 0
463            pick = int(torch.multinomial(weights, 1))
464        chosen = int(cand_ids[pick])
465        p_sampled = float(cand_probs[pick])
466
467    p = float(probs[chosen])
468    # Candidates are a prefix of the tokens sorted by probability, so rank decides who survived.
469    return chosen, {
470        "p": p,
471        "logprob": float(torch.log(probs[chosen])) if p > 0 else None,
472        "rank": int((probs > probs[chosen]).sum()) + 1,
473        "p_sampled": p_sampled,
474        "kept": kept,
475        "top": [
476            {"id": int(i), "piece": self.piece(int(i)), "p": float(v), "kept": rank < kept}
477            for rank, (v, i) in enumerate(top)
478        ],
479    }

224def load_drafter(self):
225    """Load the MTP drafter that goes with this model (families.DRAFTERS), if any, and warm it up."""
226    with self._drafter_lock:
227        if self.drafter is not None or self.drafter_id is None:
228            return
229        start = time.perf_counter()
230        folder = _download(self.drafter_id)
231        self.drafter = AutoModelForCausalLM.from_pretrained(folder, dtype=self.dtype, device_map=self.device).eval()
232        self.draft_tokens = self.drafter.generation_config.num_assistant_tokens  # guesses per forward pass
233        self.warmup(mtp=True)
234        cfg = self.drafter.config.get_text_config()
235        self.info["drafter"] = {
236            "model_id": self.drafter_id,
237            "parameters": self.drafter.num_parameters(),
238            "layers": cfg.num_hidden_layers,
239            "hidden_size": cfg.hidden_size,
240            "draft_tokens": self.draft_tokens,
241            "load_seconds": round(time.perf_counter() - start, 1),
242            "folder": describe_folder(folder, self.family.drafter_file_notes),
243        }

417def _draft(self, token_id: int, hidden: torch.Tensor, shared: dict, position: int) -> list[int]:
418    """The drafter's greedy guesses for the tokens after `token_id`. It has no cache of its own: it attends to the
419    model's keys and values, so every guess sits at the position of `token_id`, and what it has guessed so far
420    comes in through its input, the embedding of the last guess next to the drafter's last hidden state."""
421    embed = self.model.get_input_embeddings()
422    positions = torch.tensor([[position]], device=self.device)
423    guesses = []
424    while len(guesses) < self.draft_tokens and token_id not in self.eos_ids:
425        x = torch.cat([embed(torch.tensor([[token_id]], device=self.device)), hidden], -1)
426        out = self.drafter(inputs_embeds=x, position_ids=positions, shared_kv_states=shared)
427        token_id, hidden = int(out.logits[0, -1].argmax()), out.last_hidden_state
428        guesses.append(token_id)
429    return guesses

384def _drafted(self, out, s: Sampling, stats: dict) -> Iterator[tuple[int, dict]]:
385    """MTP decoding. After each token the drafter guesses the next few, and one forward pass over the token and the
386    guesses gives the model's distribution at every guessed position. _sample checks the guesses in order: each
387    one it keeps is a token for free, the first one it turns down is replaced by the model's own pick, and the
388    rest are dropped. When every guess holds, the last position gives one more token."""
389    past = out.past_key_values
390    past.activate_past_recording()  # so the sliding-window layers can still drop guesses that were turned down
391    hidden, shared = out.hidden_states[-1][:, -1:], out.shared_kv_states
392    token_id, details = self._sample(out.logits[0, -1].float().cpu(), s)
393    while True:
394        yield token_id, details
395        guesses = self._draft(token_id, hidden, shared, past.get_seq_length())
396        out = self._forward(torch.tensor([[token_id, *guesses]], device=self.device), past, len(guesses) + 1, True)
397        logits = out.logits[0].float().cpu()
398        stats["forwards"] += 1
399        stats["drafted"] += len(guesses)
400        for n, guess in enumerate(guesses):
401            token_id, details = self._sample(logits[n], s, guess)
402            if token_id != guess:
403                details["draft"] = {"piece": self.piece(guess), "accepted": False}
404                break
405            stats["accepted"] += 1
406            yield token_id, {**details, "draft": {"piece": self.piece(guess), "accepted": True}}
407        else:
408            n = len(guesses)
409            token_id, details = self._sample(logits[n], s)
410        # Keep the cache and the drafter's inputs up to the last kept guess (or the token, when none held).
411        dropped = len(guesses) - n
412        past.crop(-dropped)
413        hidden = out.hidden_states[-1][:, n : n + 1]
414        shared = {k: (key[:, :, : key.shape[2] - dropped], value[:, :, : value.shape[2] - dropped])
415                  for k, (key, value) in out.shared_kv_states.items()}

456at = (cand_ids == draft).nonzero().flatten() if draft is not None else []
457if len(at) and float(torch.rand(())) < float(cand_probs[at[0]]):
458    pick = int(at[0])
459else:
460    if len(at):  # the guess was turned down: draw from the other candidates
461        weights = cand_probs.clone()
462        weights[at[0]] = 0
463    pick = int(torch.multinomial(weights, 1))

63class Backend(Protocol):
64    """What runs the model. The page holds one at a time (app.page.current_llm) and runs one turn on it at a time."""
65
66    model_id: str
67    family: families.Family  # how the model writes its turns, thinking and tool calls (families/)
68    modalities: tuple[str, ...]  # what it reads besides text: "image", "audio", "video"
69    context_length: int  # the most tokens it reads at once
70    info: dict  # what the Model tab shows, and in "defaults" the recommended temperature, top_p and top_k
71    drafter_id: str | None  # the MTP drafter that goes with the model (families.DRAFTERS), if any
72    drafter: object | None  # that drafter once load_drafter() has loaded it; Sampling.mtp turns it on only then
73
74    def warmup(self) -> None:
75        """Pay what the first question would otherwise wait for, such as compiling kernels."""
76
77    def load_drafter(self) -> None:
78        """Load the MTP drafter, if the model has one and it is not loaded yet."""
79
80    def render(self, messages: list[dict], tools: list[dict] | None, thinking: bool, generation_prompt: bool = True) -> str:
81        """The chat template's string for the messages and tool schemas, which is what the model reads; with
82        `generation_prompt` it also opens the model's turn. The harness finds what is new in a step by comparing
83        this string with the last step's, so it has to be the prompt itself, not something like it."""
84
85    def encode(self, prompt: str, parts: list[dict]) -> Encoded:
86        """The prompt as the model reads it, with each image, clip or video in `parts` (media.parts) expanded into
87        the tokens it takes."""
88
89    def generate(self, prompt: Encoded, sampling: Sampling, should_stop: Callable[[], bool]) -> Iterator[dict]:
90        """Generate the model's turn, asking should_stop() before every token. It yields, in order:
91
92        - {"type": "prefill", "n_tokens", "ms"}: the prompt has been read, in `ms` milliseconds;
93        - {"type": "token", ...} for each token, with TOKEN_FIELDS and the TOKEN_DETAILS the backend knows;
94        - {"type": "done", "finish", "stop_piece", "n_tokens", "decode_ms", "tok_per_s", "mtp", "raw"}: `finish`
95          is "stop_token", "max_new_tokens" or "user_stopped"; `stop_piece` is the stop token's piece; `mtp` is
96          {"drafted", "accepted", "forwards"} with MTP on (and "kept_at", "turned_down_at" from a backend that
97          learns the guesses only then), else None; and `raw` is the whole output, special tokens included, which
98          the family parses."""

48TOKEN_FIELDS = ("index", "id", "piece", "text", "special", "channel")
49# What a token event carries only when the backend knows it:
50#   p, logprob  the token's probability in the model's own distribution (temperature 1, nothing cut off), and its log
51#   rank        the token's place in that distribution, 1 for the most likely
52#   p_sampled   its probability in the distribution it was drawn from, after temperature, top-k and top-p
53#   kept        how many candidates top-k and top-p left
54#   top         the 5 most likely tokens, each with id, piece, p, and kept (whether top-k and top-p left it) when known
55#   draft       with MTP, what the drafter guessed here and whether the model kept it; absent where it guessed nothing.
56#               A backend that learns only at the end of a step which tokens were guesses (vllm_backend.py) says so in
57#               done's mtp instead (kept_at, turned_down_at), and render.py marks those tokens; it may not know what a
58#               turned-down guess was (piece None).
59TOKEN_DETAILS = ("p", "logprob", "rank", "p_sampled", "kept", "top", "draft")

538def _sampling_params(self, prompt: Encoded, s: Sampling) -> dict:
539    """Each sidebar setting, and every other sampling setting set to do nothing, so that a default the server took
540    from the checkpoint (a repetition penalty, say) cannot change the distribution the page shows."""
541    return {
542        "max_tokens": max(1, min(s.max_new_tokens, self.context_length - prompt.n_tokens)),  # vLLM refuses a request that cannot fit
543        "temperature": max(s.temperature, 0.0),  # 0 is greedy there as here
544        "top_p": s.top_p,
545        "top_k": s.top_k,  # 0 is no limit there as here
546        "min_p": 0.0,
547        "repetition_penalty": 1.0,
548        "presence_penalty": 0.0,
549        "frequency_penalty": 0.0,
550        "stop_token_ids": sorted(self.eos_ids),
551    }

598def _details(self, token_id: int, logprobs: dict, n: int, s: Sampling) -> dict:
599    """What /v1/completions said about the n-th token of a chunk: its log probability, and those of the 5 most likely
600    tokens, keyed "token_id:<id>". The token's rank is known only when it is one of those 5."""
601    token_logprobs, top_logprobs = logprobs.get("token_logprobs") or [], logprobs.get("top_logprobs") or []
602    if n >= len(token_logprobs) or token_logprobs[n] is None:
603        return {}
604    logprob = token_logprobs[n]
605    listed = top_logprobs[n] if n < len(top_logprobs) and top_logprobs[n] else {}
606    # The chosen token is listed too, even when it is not among the 5.
607    listed = {int(k.removeprefix("token_id:")): v for k, v in listed.items() if k.startswith("token_id:")}
608    top = sorted(listed.items(), key=lambda c: (-c[1], c[0]))[:TOP_N]
609    rank = 1 + sum(v > logprob for _, v in top) if any(i == token_id for i, _ in top) else None
610    return self._token_details(token_id, logprob, listed, rank, None, s)

664def _token_details(self, token_id: int, logprob: float, listed: dict[int, float], rank: int | None,
665                   mask: list[int] | None, s: Sampling) -> dict:
666    """The details of a token event, from what the server said about the token: its log probability, those of the
667    tokens listed with it, its rank and its sampling mask when it sent them. What the numbers may be shown as
668    depends on the server's --logprobs-mode. Only the model's own log probabilities are p, logprob and the top 5.
669    Processed ones are the distribution the token was drawn from, so the token's is its sampling probability, and
670    the tokens top-k and top-p cut come back as UNREPORTED. Logits are neither."""
671    details = {"rank": rank} if rank else {}
672    if self.logprobs_mode == "raw_logprobs":
673        # Equally likely tokens go lowest id first, as in LocalLLM.
674        ranked = sorted(listed.items(), key=lambda c: (-c[1], c[0]))
675        if rank is not None and rank > TOP_N:  # listed only because it was chosen
676            ranked = [c for c in ranked if c[0] != token_id]
677        top = [{"id": i, "piece": self.piece(i), "p": math.exp(v)} for i, v in ranked[:TOP_N]]
678        details.update(p=math.exp(logprob), logprob=logprob, top=top)
679    elif self.logprobs_mode == "processed_logprobs":
680        details["p_sampled"] = math.exp(logprob)
681    if mask is not None:
682        details["kept"] = len(set(mask))
683    elif self.logprobs_mode.startswith("processed") and any(v <= UNREPORTED for v in listed.values()):
684        # Some of the most likely tokens were cut, and those left are more likely still, so all of them are listed.
685        details["kept"] = sum(v > UNREPORTED for v in listed.values())
686    return self._greedy(details, token_id, s)

449def encode(self, prompt: str, parts: list[dict]) -> Encoded:
450    """The prompt's token ids, which are what the server reads. With media (`parts`, in the order of their
451    placeholders), the server expands each placeholder into the tokens its file takes, so it is asked how first, with
452    a one-token request; the text and the count are of the prompt as the server reads it, and the ids stay the ones
453    the chat template wrote, which generate() sends again with the files."""
454    ids = self.tokenizer(prompt, add_special_tokens=False).input_ids
455    if not parts:
456        return Encoded(prompt, ids, len(ids))
457    if unread := {p["type"] for p in parts} - set(self.modalities):  # media.save makes those plain files, so this is a bug
458        raise ValueError(f"The vLLM server takes no {', '.join(sorted(unread))} here, only {', '.join(self.modalities) or 'text'}.")
459    files = [content_part(p, self.audio_rate) for p in parts]
460    expanded = self._generate_once(ids, files, return_token_ids=True)["prompt_token_ids"]
461    counts = Counter(kind for token in expanded for kind, media_id in self.media_token_ids.items() if media_id is not None and token == media_id)
462    return Encoded(self.tokenizer.decode(expanded, skip_special_tokens=False), ids, len(expanded),
463                   {"content_parts": files, "prompt_token_ids": expanded}, dict(counts))

243def drafted_tokens(accepted: list[int], drafted: list[int], n_tokens: int) -> dict | None:
244    """Which of a step's tokens were guesses the model kept (kept_at) and which replace a guess it turned down
245    (turned_down_at), from vLLM's per-step counts (--per-request-spec-decode-metrics detailed). The first token comes
246    from the prefill; after it, each verify step gives the guesses it kept and then one token of the model's own,
247    which replaces the next guess when it kept fewer than were drafted. vLLM does not count a step without guesses,
248    so when the steps do not account for every token, nobody can tell which tokens were guessed: None."""
249    kept, turned_down, at = [], [], 1
250    for n_kept, n_drafted in zip(accepted, drafted):
251        kept += range(at, at + n_kept)
252        if n_kept < n_drafted:
253            turned_down.append(at + n_kept)
254        at += n_kept + 1
255    last = accepted[-1] + 1 if accepted else 1
256    if not at - last < n_tokens <= at:  # only a stop token or the token limit may cut the last step short
257        return None
258    return {"kept_at": [i for i in kept if i < n_tokens], "turned_down_at": [i for i in turned_down if i < n_tokens]}

144prefix = family.thought_prefix if prompt.endswith(family.thought_prefix) else ""
145parsed = family.parse(prefix + raw, schemas, s.lang)
146calls = parsed["tool_calls"] if tools else []
147self.emit("parsed", step=step, thinking=parsed["thinking"], content=parsed["content"], tool_calls=calls)

32def parse_output(raw: str, schemas: dict | None = None, lang: str = "zh") -> dict:
33    """Split one model turn into thinking, tool calls and visible content. Gemma marks each value's
34    type itself, so the tool schemas are not needed. Errors are worded in `lang`."""
35    text = raw
36    thinking = None
37    if THOUGHT_OPEN in text:
38        before, _, rest = text.partition(THOUGHT_OPEN)
39        thought, closed, after = rest.partition(THOUGHT_CLOSE)
40        thinking = thought.strip()
41        text = before + (after if closed else "")
42
43    calls = []
44    for m in CALL_RE.finditer(text):
45        call = {"name": m.group(1), "raw_args": m.group(2), "arguments": None, "error": None}
46        try:
47            call["arguments"] = parse_args(m.group(2))
48        except ParseError as e:
49            call["error"] = e.text[lang]
50        calls.append(call)
51
52    rest = CALL_RE.sub("", text)
53    if CALL_OPEN in rest:
54        head, _, tail = rest.partition(CALL_OPEN)
55        name = re.match(r"call:(\w+)", tail)
56        calls.append({
57            "name": name.group(1) if name else "?",
58            "raw_args": tail,
59            "arguments": None,
60            "error": pick(lang, "工具呼叫沒有結束標記 <tool_call|>", "The tool call has no closing <tool_call|>"),
61        })
62        rest = head
63
64    for marker in END_MARKERS:
65        rest = rest.replace(marker, "")
66    return {"thinking": thinking, "tool_calls": calls, "content": rest.strip()}

99def value(self):
100    self.skip_ws()
101    if self.text.startswith(STR, self.pos):
102        return self.delimited()
103    ch = self.peek()
104    if ch == "{":
105        return self.obj()
106    if ch == "[":
107        return self.arr()
108    if ch == '"':
109        return self.json_string()
110    return self.scalar()

31def parse_output(raw: str, schemas: dict | None = None, lang: str = "zh") -> dict:
32    """Split one model turn into thinking, tool calls and visible content. `schemas` maps a tool
33    name to its parameters schema, which says what type each plain-text value is. Errors are worded in `lang`."""
34    text = raw
35    thinking = None
36    if THINK_CLOSE in text:
37        thought, _, text = text.partition(THINK_CLOSE)
38        thinking = thought.split(THINK_OPEN)[-1].strip()
39    elif THINK_OPEN in text:  # cut off while still thinking
40        text, _, thought = text.partition(THINK_OPEN)
41        thinking = thought.strip()
42
43    calls = [parse_call(m.group(1), schemas or {}, lang) for m in CALL_RE.finditer(text)]
44    rest = CALL_RE.sub("", text)
45    if CALL_OPEN in rest:
46        rest, _, tail = rest.partition(CALL_OPEN)
47        call = parse_call(tail, schemas or {}, lang)
48        call["arguments"], call["error"] = None, pick(lang, "工具呼叫沒有結束標記 </tool_call>", "The tool call has no closing </tool_call>")
49        calls.append(call)
50
51    for marker in END_MARKERS:
52        rest = rest.replace(marker, "")
53    return {"thinking": thinking, "tool_calls": calls, "content": rest.strip()}