Function generate

smallthinker/tools/run/run.cpp:988–1025

helper function to evaluate a prompt and generate a response

// helper function to evaluate a prompt and generate a response
static int generate(LlamaData & llama_data, const std::string & prompt, std::string & response) {
    const llama_vocab * vocab = llama_model_get_vocab(llama_data.model.get());

    std::vector<llama_token> tokens;
    if (tokenize_prompt(vocab, prompt, tokens, llama_data) < 0) {
        return 1;
    }

    // prepare a batch for the prompt
    llama_batch batch = llama_batch_get_one(tokens.data(), tokens.size());
    llama_token new_token_id;
    while (true) {
        check_context_size(llama_data.context, batch);
        if (llama_decode(llama_data.context.get(), batch)) {
            printe("failed to decode\n");
            return 1;
        }

        // sample the next token, check is it an end of generation?
        new_token_id = llama_sampler_sample(llama_data.sampler.get(), llama_data.context.get(), -1);
        if (llama_vocab_is_eog(vocab, new_token_id)) {
            break;
        }

        std::string piece;
        if (convert_token_to_string(vocab, new_token_id, piece)) {
            return 1;
        }

        print_word_and_concatenate_to_response(piece, response);

        // prepare the next batch with the sampled token
        batch = llama_batch_get_one(&new_token_id, 1);
    }

    printf(LOG_COL_DEFAULT);
    return 0;
}

static int read_user_input(std::string & user_input) {
    static const char * prompt_prefix_env = std::getenv("LLAMA_PROMPT_PREFIX");