Felipe97/llama-cpp-compiled

1#include "sampling.h"2#include "log.h"3 4#ifdef LLAMA_USE_LLGUIDANCE5 6#    include "llguidance.h"7#    include <cmath>8 9struct llama_sampler_llg {10    const llama_vocab * vocab;11    std::string         grammar_kind;12    std::string         grammar_data;13    LlgTokenizer *      tokenizer;14    LlgMatcher *        grammar;15};16 17static LlgMatcher * llama_sampler_llg_new(LlgTokenizer * tokenizer, const char * grammar_kind,18                                          const char * grammar_data) {19    LlgConstraintInit cinit;20    llg_constraint_init_set_defaults(&cinit, tokenizer);21    const char * log_level = getenv("LLGUIDANCE_LOG_LEVEL");22    if (log_level && *log_level) {23        cinit.log_stderr_level = atoi(log_level);24    }25    auto c = llg_new_matcher(&cinit, grammar_kind, grammar_data);26    if (llg_matcher_get_error(c)) {27        LOG_ERR("llg error: %s\n", llg_matcher_get_error(c));28        llg_free_matcher(c);29        return nullptr;30    }31 32    return c;33}34 35static const char * llama_sampler_llg_name(const llama_sampler * /*smpl*/) {36    return "llguidance";37}38 39static void llama_sampler_llg_accept_impl(llama_sampler * smpl, llama_token token) {40    auto * ctx = (llama_sampler_llg *) smpl->ctx;41    if (ctx->grammar) {42        llg_matcher_consume_token(ctx->grammar, token);43    }44}45 46static void llama_sampler_llg_apply(llama_sampler * smpl, llama_token_data_array * cur_p) {47    auto * ctx = (llama_sampler_llg *) smpl->ctx;48    if (ctx->grammar) {49        const uint32_t * mask = llg_matcher_get_mask(ctx->grammar);50        if (mask == nullptr) {51            if (llg_matcher_compute_mask(ctx->grammar) == 0) {52                mask = llg_matcher_get_mask(ctx->grammar);53            } else {54                LOG_ERR("llg error: %s\n", llg_matcher_get_error(ctx->grammar));55                llg_free_matcher(ctx->grammar);56                ctx->grammar = nullptr;57                return;58            }59        }60 61        for (size_t i = 0; i < cur_p->size; ++i) {62            auto token = cur_p->data[i].id;63            if ((mask[token / 32] & (1 << (token % 32))) == 0) {64                cur_p->data[i].logit = -INFINITY;65            }66        }67    }68}69 70static void llama_sampler_llg_reset(llama_sampler * smpl) {71    auto * ctx = (llama_sampler_llg *) smpl->ctx;72    if (ctx->grammar) {73        llg_matcher_reset(ctx->grammar);74    }75}76 77static llama_sampler * llama_sampler_llg_clone(const llama_sampler * smpl) {78    const auto * ctx = (const llama_sampler_llg *) smpl->ctx;79 80    auto * result = llama_sampler_init_llg(ctx->vocab, nullptr, nullptr);81 82    // copy the state83    {84        auto * result_ctx = (llama_sampler_llg *) result->ctx;85 86        if (ctx->grammar) {87            result_ctx->grammar_kind = ctx->grammar_kind;88            result_ctx->grammar_data = ctx->grammar_data;89            result_ctx->grammar      = llg_clone_matcher(ctx->grammar);90            result_ctx->tokenizer    = llg_clone_tokenizer(ctx->tokenizer);91        }92    }93 94    return result;95}96 97static void llama_sampler_llg_free(llama_sampler * smpl) {98    const auto * ctx = (llama_sampler_llg *) smpl->ctx;99 100    if (ctx->grammar) {101        llg_free_matcher(ctx->grammar);102        llg_free_tokenizer(ctx->tokenizer);103    }104 105    delete ctx;106}107 108static llama_sampler_i llama_sampler_llg_i = {109    /* .name              = */ llama_sampler_llg_name,110    /* .accept            = */ llama_sampler_llg_accept_impl,111    /* .apply             = */ llama_sampler_llg_apply,112    /* .reset             = */ llama_sampler_llg_reset,113    /* .clone             = */ llama_sampler_llg_clone,114    /* .free              = */ llama_sampler_llg_free,115    /* .backend_init      = */ NULL,116    /* .backend_accept    = */ NULL,117    /* .backend_apply     = */ NULL,118    /* .backend_set_input = */ NULL,119    /* .backend_reset     = */ NULL,120    /* .copy_state        = */ NULL,121};122 123static size_t llama_sampler_llg_tokenize_fn(const void * user_data, const uint8_t * bytes, size_t bytes_len,124                                            uint32_t * output_tokens, size_t output_tokens_len) {125    const llama_vocab * vocab = (const llama_vocab *) user_data;126    int                 r     = 0;127    try {128        r = llama_tokenize(vocab, (const char *) bytes, bytes_len, (int32_t *) output_tokens, output_tokens_len, false,129                           true);130    } catch (const std::exception & e) {131        GGML_ABORT("llama_tokenize failed: %s\n", e.what());132    }133    if (r < 0) {134        return -r;135    }136    return r;137}138 139static LlgTokenizer * llama_sampler_llg_new_tokenizer(const llama_vocab * vocab) {140    // TODO store the tokenizer in the vocab somehow141    static const llama_vocab * vocab_cache;142    static LlgTokenizer *      tokenizer_cache;143 144    if (vocab_cache == vocab) {145        return llg_clone_tokenizer(tokenizer_cache);146    }147 148    auto tok_eos = llama_vocab_eot(vocab);149    if (tok_eos == LLAMA_TOKEN_NULL) {150        tok_eos = llama_vocab_eos(vocab);151    }152 153    size_t vocab_size = llama_vocab_n_tokens(vocab);154 155    auto token_lens       = new uint32_t[vocab_size];156    // we typically have ~7 bytes per token; let's go on the safe side here157    auto token_bytes_size = vocab_size * 16 + 1024 * 1024;158    auto token_bytes      = new uint8_t[token_bytes_size];159 160    size_t offset = 0;161    for (size_t i = 0; i < vocab_size; i++) {162        size_t max_token = 1024;163        if (token_bytes_size - offset < max_token) {164            GGML_ABORT("token_bytes buffer too small\n");165        }166 167        llama_token token = i;168        auto        dp    = (char *) token_bytes + offset;169        auto        size  = llama_detokenize(vocab, &token, 1, dp, max_token, false, false);170        if (size < 0) {171            GGML_ABORT("llama_detokenize failed\n");172        }173        if (size == 0) {174            size = llama_detokenize(vocab, &token, 1, dp + 1, max_token - 1, false, true);175            if (size < 0) {176                GGML_ABORT("llama_detokenize failed\n");177            }178            if (size != 0) {179                *dp = '\xff';  // special token prefix marker180                size += 1;181            }182        }183 184        token_lens[i] = size;185        offset += size;186    }187 188    LlgTokenizerInit tinit = {189        /* .vocab_size                         = */ (uint32_t) vocab_size,190        /* .tok_eos                            = */ (uint32_t) tok_eos,191        /* .token_lens                         = */ token_lens,192        /* .token_bytes                        = */ token_bytes,193        /* .tokenizer_json                     = */ nullptr,194        /* .tokenize_assumes_string            = */ true,195        /* .tokenize_fn                        = */ llama_sampler_llg_tokenize_fn,196        /* .use_approximate_greedy_tokenize_fn = */ false,197        /* .tokenize_user_data                 = */ vocab,198        /* .slices                             = */ nullptr,199    };200 201    char           error_buffer[1024];202    LlgTokenizer * tokenizer = llg_new_tokenizer(&tinit, error_buffer, sizeof(error_buffer));203 204    delete[] token_bytes;205    delete[] token_lens;206 207    if (tokenizer == nullptr) {208        LOG_ERR("llg tokenizer error: %s\n", error_buffer);209        return tokenizer;210    }211 212    if (tokenizer_cache) {213        llg_free_tokenizer(tokenizer_cache);214    }215    vocab_cache     = vocab;216    tokenizer_cache = tokenizer;217 218    return llg_clone_tokenizer(tokenizer_cache);219}220 221llama_sampler * llama_sampler_init_llg(const llama_vocab * vocab, const char * grammar_kind,222                                       const char * grammar_data) {223    auto * ctx = new llama_sampler_llg;224 225    if (grammar_kind != nullptr && grammar_kind[0] != '\0') {226        auto tokenizer = llama_sampler_llg_new_tokenizer(vocab);227        *ctx           = {228            /* .vocab        = */ vocab,229            /* .grammar_kind = */ grammar_kind,230            /* .grammar_data = */ grammar_data,231            /* .tokenizer    = */ tokenizer,232            /* .grammar      = */ llama_sampler_llg_new(tokenizer, grammar_kind, grammar_data),233        };234        if (ctx->grammar) {235            GGML_ASSERT(((size_t) llama_vocab_n_tokens(vocab) + 31) / 32 * 4 ==236                        llg_matcher_get_mask_byte_size(ctx->grammar));237        }238    } else {239        *ctx = {240            /* .vocab        = */ vocab,241            /* .grammar_kind = */ {},242            /* .grammar_data = */ {},243            /* .tokenizer    = */ nullptr,244            /* .grammar      = */ nullptr,245        };246    }247 248    return llama_sampler_init(249        /* .iface = */ &llama_sampler_llg_i,250        /* .ctx   = */ ctx);251}252 253#else254 255llama_sampler * llama_sampler_init_llg(const llama_vocab *, const char *, const char *) {256    LOG_WRN("llguidance (cmake -DLLAMA_LLGUIDANCE=ON) is not enabled");257    return nullptr;258}259 260#endif  // LLAMA_USE_LLGUIDANCE261