inference_wrapper.cpp

#include "inference_wrapper.h"

#include <iostream>
#include <fstream>
#include <algorithm>

/*
  Senior rules:
  - steady_clock only
  - never allocate per token
  - stream immediately
*/

InferenceWrapper::InferenceWrapper(int threads, int max_context)
    : n_threads(threads), max_ctx(max_context) {
    process_start = std::chrono::steady_clock::now();
}

InferenceWrapper::~InferenceWrapper() {
    if (ctx) llama_free(ctx);
    if (model) llama_free_model(model);
}

void InferenceWrapper::load_model(const std::string& model_path) {

    auto load_start = std::chrono::steady_clock::now();

    llama_model_params mparams = llama_model_default_params();
    mparams.use_mmap = true;
    mparams.use_mlock = false;

    model = llama_load_model_from_file(model_path.c_str(), mparams);
    if (!model) throw std::runtime_error("model load failed");

    vocab = llama_model_get_vocab(model);

    llama_context_params cparams = llama_context_default_params();
    cparams.n_ctx = max_ctx;
    cparams.n_threads = n_threads;
    cparams.n_threads_batch = n_threads;

    ctx = llama_new_context_with_model(model, cparams);
    if (!ctx) throw std::runtime_error("context init failed");

    auto load_end = std::chrono::steady_clock::now();
    double load_ms =
        std::chrono::duration<double, std::milli>(load_end - load_start).count();

    std::cout << "[METRIC] model_load_ms=" << load_ms << std::endl;
}

void InferenceWrapper::generate(const std::string& prompt) {

    /* ---- tokenize prompt ---- */
    std::vector<llama_token> prompt_tokens(prompt.size() + 16);
    int n_prompt = llama_tokenize(
        vocab,
        prompt.c_str(),
        prompt.size(),
        prompt_tokens.data(),
        prompt_tokens.size(),
        true,
        false
    );
    prompt_tokens.resize(n_prompt);

    /* ---- evaluate prompt ---- */
    llama_batch batch = llama_batch_init(n_prompt, 0, 1);
    for (int i = 0; i < n_prompt; ++i) {
        batch.token[i] = prompt_tokens[i];
        batch.pos[i] = i;
        batch.n_seq_id[i] = 1;
        batch.seq_id[i][0] = 0;
        batch.logits[i] = false;
    }
    batch.n_tokens = n_prompt;

    llama_decode(ctx, batch);
    llama_batch_free(batch);

    /* ---- generation loop ---- */
    const int max_gen = 128;
    auto prev_ts = std::chrono::steady_clock::now();

    for (int i = 0; i < max_gen; ++i) {

        const float* logits = llama_get_logits(ctx);
        int n_vocab = llama_vocab_n_tokens(vocab);

        llama_token next =
            std::max_element(logits, logits + n_vocab) - logits;

        /* decode token */
        llama_batch b = llama_batch_init(1, 0, 1);
        b.token[0] = next;
        b.pos[0] = n_prompt + i;
        b.n_seq_id[0] = 1;
        b.seq_id[0][0] = 0;
        b.logits[0] = true;
        b.n_tokens = 1;

        auto t0 = std::chrono::steady_clock::now();
        llama_decode(ctx, b);
        auto t1 = std::chrono::steady_clock::now();
        llama_batch_free(b);

        if (!first_token_seen) {
            double ttfb =
                std::chrono::duration<double, std::milli>(t1 - process_start).count();
            std::cout << "[METRIC] time_to_first_token_ms=" << ttfb << std::endl;
            first_token_seen = true;
        }

        /* stream token */
        char piece[32];
        int len = llama_token_to_piece(vocab, next, piece, sizeof(piece), 0, false);
        if (len > 0) {
            std::cout.write(piece, len);
            std::cout.flush();
        }

        if (next == llama_vocab_eos(vocab)) break;

        prev_ts = t1;
    }

    std::cout << std::endl;
}

size_t InferenceWrapper::read_rss_kb() {
    std::ifstream f("/proc/self/status");
    std::string line;
    while (std::getline(f, line)) {
        if (line.rfind("VmRSS:", 0) == 0)
            return std::stoul(line.substr(6));
    }
    return 0;
}