Felipe97/llama-cpp-compiled

1#include "reasoning-budget.h"2#include "common.h"3#include "trie.h"4#include "unicode.h"5 6#include "log.h"7 8#include <algorithm>9#include <cmath>10#include <cstdint>11#include <string>12#include <vector>13 14struct token_matcher {15    std::vector<llama_tokens> seqs;16    common_aho_corasick ac;17    size_t state = 0;18 19    token_matcher(const std::vector<llama_tokens> & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {}20 21    static std::vector<llama_tokens> collect(const std::vector<llama_tokens> & seqs) {22        std::vector<llama_tokens> res;23        for (const auto & seq : seqs) {24            if (!seq.empty() && std::find(res.begin(), res.end(), seq) == res.end()) {25                res.push_back(seq);26            }27        }28        return res;29    }30 31    static common_trie build_trie(const std::vector<llama_tokens> & seqs) {32        common_trie t;33        for (const auto & seq : seqs) {34            t.insert(std::vector<uint32_t>(seq.begin(), seq.end()));35        }36        return t;37    }38 39    // returns the index into seqs of the longest sequence ending at this token, or -140    int32_t advance(llama_token token) {41        state = ac.next(state, (uint32_t) token);42        const int32_t p = ac.match_pattern(state);43        if (p >= 0) {44            state = 0;45        }46        return p;47    }48 49    void reset() { state = 0; }50};51 52struct common_reasoning_budget_ctx {53    const llama_vocab * vocab;54 55    token_matcher start_matcher;56    token_matcher end_matcher;57    llama_tokens forced_tokens;58 59    int32_t budget;           // maximum tokens in reasoning block60    int32_t remaining;        // tokens remaining in budget61 62    common_reasoning_budget_state state;63 64    // for forcing65    size_t force_pos;         // next position in forced_tokens to force66 67    int32_t end_match;        // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none68};69 70static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) {71    return "reasoning-budget";72}73 74static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_token token) {75    auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;76 77    switch (ctx->state) {78        case REASONING_BUDGET_IDLE:79        {80            if (ctx->start_matcher.advance(token) >= 0) {81                ctx->state = REASONING_BUDGET_COUNTING;82                ctx->remaining = ctx->budget;83                COM_TRC("activated, budget=%d tokens\n", ctx->budget);84 85                if (ctx->remaining <= 0) {86                    ctx->state = REASONING_BUDGET_FORCING;87                    ctx->force_pos = 0;88                    COM_TRC("%s", "budget=0, forcing immediately\n");89                }90            }91            break;92        }93        case REASONING_BUDGET_COUNTING:94        case REASONING_BUDGET_WAITING_UTF8:95        {96            const int32_t match = ctx->end_matcher.advance(token);97            if (match >= 0) {98                ctx->state = REASONING_BUDGET_DONE;99                ctx->end_match = match;100                COM_TRC("%s", "deactivated (natural end)\n");101                break;102            }103 104            bool utf8_complete = true;105            if (ctx->vocab != nullptr) {106                const std::string piece = common_token_to_piece(ctx->vocab, token, false);107                utf8_complete = common_utf8_is_complete(piece);108            }109 110            if (ctx->state == REASONING_BUDGET_WAITING_UTF8) {111                if (utf8_complete) {112                    ctx->state = REASONING_BUDGET_FORCING;113                    ctx->force_pos = 0;114                    ctx->end_matcher.reset();115                    COM_TRC("%s", "UTF-8 complete, now forcing end sequence\n");116                }117            } else if (ctx->state == REASONING_BUDGET_COUNTING) {118                ctx->remaining--;119                if (ctx->remaining <= 0) {120                    if (utf8_complete) {121                        ctx->state = REASONING_BUDGET_FORCING;122                        ctx->force_pos = 0;123                        ctx->end_matcher.reset();124                        COM_TRC("%s", "budget exhausted, forcing end sequence\n");125                    } else {126                        ctx->state = REASONING_BUDGET_WAITING_UTF8;127                        ctx->end_matcher.reset();128                        COM_TRC("%s", "budget exhausted, waiting for UTF-8 completion\n");129                    }130                }131            }132            break;133        }134        case REASONING_BUDGET_FORCING:135        {136            // track the end sequence within forced_tokens so it is also reported on DONE137            const int32_t match = ctx->end_matcher.advance(token);138            ctx->force_pos++;139            if (ctx->force_pos >= ctx->forced_tokens.size()) {140                ctx->state = REASONING_BUDGET_DONE;141                ctx->end_match = match;142                COM_TRC("%s", "forced sequence complete, done\n");143            }144            break;145        }146        case REASONING_BUDGET_DONE:147            // Re-arm on a new start tag: some models emit multiple <think> blocks148            // per response, and each should get a fresh budget window.149            if (ctx->start_matcher.advance(token) >= 0) {150                ctx->state = REASONING_BUDGET_COUNTING;151                ctx->remaining = ctx->budget;152                ctx->end_matcher.reset();153                ctx->end_match = -1;154                COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget);155 156                if (ctx->remaining <= 0) {157                    ctx->state = REASONING_BUDGET_FORCING;158                    ctx->force_pos = 0;159                    COM_TRC("%s", "budget=0, forcing immediately\n");160                }161            }162            break;163    }164}165 166static void common_reasoning_budget_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {167    auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;168 169    if (ctx->state != REASONING_BUDGET_FORCING) {170        // passthrough — don't modify logits171        return;172    }173 174    if (ctx->force_pos >= ctx->forced_tokens.size()) {175        return;176    }177 178    const llama_token forced = ctx->forced_tokens[ctx->force_pos];179 180    // set all logits to -inf except the forced token181    for (size_t i = 0; i < cur_p->size; i++) {182        if (cur_p->data[i].id != forced) {183            cur_p->data[i].logit = -INFINITY;184        }185    }186}187 188static void common_reasoning_budget_reset(struct llama_sampler * smpl) {189    auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;190    ctx->state = REASONING_BUDGET_IDLE;191    ctx->remaining = ctx->budget;192    ctx->start_matcher.reset();193    ctx->end_matcher.reset();194    ctx->force_pos = 0;195    ctx->end_match = -1;196}197 198static struct llama_sampler * common_reasoning_budget_init_state(199        const struct llama_vocab * vocab, const std::vector<llama_tokens> & start_seqs,200        const std::vector<llama_tokens> & end_seqs, const llama_tokens & forced_tokens,201        int32_t budget, common_reasoning_budget_state initial_state);202 203static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl);204 205static void common_reasoning_budget_free(struct llama_sampler * smpl) {206    delete (common_reasoning_budget_ctx *) smpl->ctx;207}208 209static struct llama_sampler_i common_reasoning_budget_i = {210    /* .name              = */ common_reasoning_budget_name,211    /* .accept            = */ common_reasoning_budget_accept,212    /* .apply             = */ common_reasoning_budget_apply,213    /* .reset             = */ common_reasoning_budget_reset,214    /* .clone             = */ common_reasoning_budget_clone,215    /* .free              = */ common_reasoning_budget_free,216    /* .backend_init      = */ nullptr,217    /* .backend_accept    = */ nullptr,218    /* .backend_apply     = */ nullptr,219    /* .backend_set_input = */ nullptr,220    /* .backend_reset     = */ nullptr,221    /* .copy_state        = */ nullptr,222};223 224static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl) {225    const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx;226 227    return llama_sampler_init(228        /* .iface = */ &common_reasoning_budget_i,229        /* .ctx   = */ new common_reasoning_budget_ctx(*ctx)230    );231}232 233static struct llama_sampler * common_reasoning_budget_init_state(234        const struct llama_vocab        * vocab,235        const std::vector<llama_tokens> & start_seqs,236        const std::vector<llama_tokens> & end_seqs,237        const llama_tokens              & forced_tokens,238        int32_t                           budget,239        common_reasoning_budget_state     initial_state) {240    // promote COUNTING with budget <= 0 to FORCING241    if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) {242        initial_state = REASONING_BUDGET_FORCING;243    }244 245    return llama_sampler_init(246        /* .iface = */ &common_reasoning_budget_i,247        /* .ctx   = */ new common_reasoning_budget_ctx {248            /* .vocab         = */ vocab,249            /* .start_matcher = */ token_matcher(start_seqs),250            /* .end_matcher   = */ token_matcher(end_seqs),251            /* .forced_tokens = */ forced_tokens,252            /* .budget        = */ budget,253            /* .remaining     = */ budget,254            /* .state         = */ initial_state,255            /* .force_pos     = */ 0,256            /* .end_match     = */ -1,257        }258    );259}260 261struct llama_sampler * common_reasoning_budget_init(262        const struct llama_vocab        * vocab,263        const std::vector<llama_tokens> & start_seqs,264        const std::vector<llama_tokens> & end_seqs,265        const llama_tokens              & forced_tokens,266        int32_t                           budget,267        common_reasoning_budget_state     initial_state) {268    return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, budget, initial_state);269}270 271common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) {272    if (!smpl) {273        return REASONING_BUDGET_IDLE;274    }275    return ((const common_reasoning_budget_ctx *)smpl->ctx)->state;276}277 278const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl) {279    if (!smpl) {280        return nullptr;281    }282 283    const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx;284    if (ctx->end_match < 0) {285        return nullptr;286    }287 288    return &ctx->end_matcher.seqs[ctx->end_match];289}290 291bool common_reasoning_budget_force(struct llama_sampler * smpl) {292    if (!smpl) {293        return false;294    }295 296    auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;297 298    // only a sampler that is actively counting down the budget may be forced;299    // any other state (idle, already forcing/waiting, or done) is left untouched300    if (ctx->state != REASONING_BUDGET_COUNTING) {301        return false;302    }303 304    ctx->state = REASONING_BUDGET_FORCING;305    ctx->force_pos = 0;306    ctx->end_matcher.reset();307    COM_TRC("%s", "forced into forcing state (manual transition)\n");308 309    return true;310}311