char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Aldehir Rojascommon : add support for multiple end sequences in the reasoning budget sampler (#25544)910196f6b

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.3 KiB52 linesraw
1#pragma once
2
3#include "llama.h"
4
5#include "common.h"
6
7#include <cstdint>
8#include <vector>
9
10enum common_reasoning_budget_state {
11    REASONING_BUDGET_IDLE,         // waiting for start sequence
12    REASONING_BUDGET_COUNTING,     // counting down tokens
13    REASONING_BUDGET_FORCING,      // forcing budget message + end sequence
14    REASONING_BUDGET_WAITING_UTF8, // budget exhausted, waiting for UTF-8 completion
15    REASONING_BUDGET_DONE,         // passthrough forever
16};
17
18// Creates a reasoning budget sampler that limits token generation inside a
19// reasoning block (e.g. between <think> and </think>).
20//
21// State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE
22//   IDLE:         passthrough, watching for a start sequence
23//   COUNTING:     counting down remaining tokens, watching for a natural end sequence
24//   WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence
25//   FORCING:      forces forced_tokens token-by-token (all other logits -> -inf)
26//   DONE:         passthrough forever
27//
28// Parameters:
29//   vocab          - vocabulary (used for UTF-8 boundary detection; can be nullptr)
30//   start_seqs     - token sequences, any of which activates counting
31//   end_seqs       - token sequences, any of which naturally deactivates
32//   forced_tokens  - token sequence forced when budget expires
33//   budget         - max tokens allowed in the reasoning block
34//   initial_state  - initial state
35//
36struct llama_sampler * common_reasoning_budget_init(
37        const struct llama_vocab        * vocab,
38        const std::vector<llama_tokens> & start_seqs,
39        const std::vector<llama_tokens> & end_seqs,
40        const llama_tokens              & forced_tokens,
41        int32_t                           budget,
42        common_reasoning_budget_state     initial_state = REASONING_BUDGET_IDLE);
43
44common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl);
45
46// The end sequence that transitioned the sampler to DONE, or nullptr if none
47// was recorded. Cleared when a new start sequence re-arms the sampler.
48const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl);
49
50// Manually transition the reasoning budget sampler into the FORCING state.
51// Returns true if the transition occurred.
52bool common_reasoning_budget_force(struct llama_sampler * smpl);