mirror of
https://github.com/ggerganov/llama.cpp.git
synced 2024-12-26 03:14:35 +00:00
0cc63754b8
It's like simple-chat but it uses smart pointers to avoid manual memory cleanups. Less memory leaks in the code now. Avoid printing multiple dots. Split code into smaller functions. Uses no exception handling. Signed-off-by: Eric Curtin <ecurtin@redhat.com>
410 lines
13 KiB
C++
410 lines
13 KiB
C++
#if defined(_WIN32)
|
|
#include <windows.h>
|
|
#else
|
|
#include <unistd.h>
|
|
#endif
|
|
|
|
#include <climits>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <iostream>
|
|
#include <sstream>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
#include "llama-cpp.h"
|
|
|
|
typedef std::unique_ptr<char[]> char_array_ptr;
|
|
|
|
struct Argument {
|
|
std::string flag;
|
|
std::string help_text;
|
|
};
|
|
|
|
struct Options {
|
|
std::string model_path, prompt_non_interactive;
|
|
int ngl = 99;
|
|
int n_ctx = 2048;
|
|
};
|
|
|
|
class ArgumentParser {
|
|
public:
|
|
ArgumentParser(const char * program_name) : program_name(program_name) {}
|
|
|
|
void add_argument(const std::string & flag, std::string & var, const std::string & help_text = "") {
|
|
string_args[flag] = &var;
|
|
arguments.push_back({flag, help_text});
|
|
}
|
|
|
|
void add_argument(const std::string & flag, int & var, const std::string & help_text = "") {
|
|
int_args[flag] = &var;
|
|
arguments.push_back({flag, help_text});
|
|
}
|
|
|
|
int parse(int argc, const char ** argv) {
|
|
for (int i = 1; i < argc; ++i) {
|
|
std::string arg = argv[i];
|
|
if (string_args.count(arg)) {
|
|
if (i + 1 < argc) {
|
|
*string_args[arg] = argv[++i];
|
|
} else {
|
|
fprintf(stderr, "error: missing value for %s\n", arg.c_str());
|
|
print_usage();
|
|
return 1;
|
|
}
|
|
} else if (int_args.count(arg)) {
|
|
if (i + 1 < argc) {
|
|
if (parse_int_arg(argv[++i], *int_args[arg]) != 0) {
|
|
fprintf(stderr, "error: invalid value for %s: %s\n", arg.c_str(), argv[i]);
|
|
print_usage();
|
|
return 1;
|
|
}
|
|
} else {
|
|
fprintf(stderr, "error: missing value for %s\n", arg.c_str());
|
|
print_usage();
|
|
return 1;
|
|
}
|
|
} else {
|
|
fprintf(stderr, "error: unrecognized argument %s\n", arg.c_str());
|
|
print_usage();
|
|
return 1;
|
|
}
|
|
}
|
|
|
|
if (string_args["-m"]->empty()) {
|
|
fprintf(stderr, "error: -m is required\n");
|
|
print_usage();
|
|
return 1;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
private:
|
|
const char * program_name;
|
|
std::unordered_map<std::string, std::string *> string_args;
|
|
std::unordered_map<std::string, int *> int_args;
|
|
std::vector<Argument> arguments;
|
|
|
|
int parse_int_arg(const char * arg, int & value) {
|
|
char * end;
|
|
const long val = std::strtol(arg, &end, 10);
|
|
if (*end == '\0' && val >= INT_MIN && val <= INT_MAX) {
|
|
value = static_cast<int>(val);
|
|
return 0;
|
|
}
|
|
return 1;
|
|
}
|
|
|
|
void print_usage() const {
|
|
printf("\nUsage:\n");
|
|
printf(" %s [OPTIONS]\n\n", program_name);
|
|
printf("Options:\n");
|
|
for (const auto & arg : arguments) {
|
|
printf(" %-10s %s\n", arg.flag.c_str(), arg.help_text.c_str());
|
|
}
|
|
|
|
printf("\n");
|
|
}
|
|
};
|
|
|
|
class LlamaData {
|
|
public:
|
|
llama_model_ptr model;
|
|
llama_sampler_ptr sampler;
|
|
llama_context_ptr context;
|
|
std::vector<llama_chat_message> messages;
|
|
|
|
int init(const Options & opt) {
|
|
model = initialize_model(opt.model_path, opt.ngl);
|
|
if (!model) {
|
|
return 1;
|
|
}
|
|
|
|
context = initialize_context(model, opt.n_ctx);
|
|
if (!context) {
|
|
return 1;
|
|
}
|
|
|
|
sampler = initialize_sampler();
|
|
return 0;
|
|
}
|
|
|
|
private:
|
|
// Initializes the model and returns a unique pointer to it
|
|
llama_model_ptr initialize_model(const std::string & model_path, const int ngl) {
|
|
llama_model_params model_params = llama_model_default_params();
|
|
model_params.n_gpu_layers = ngl;
|
|
|
|
llama_model_ptr model(llama_load_model_from_file(model_path.c_str(), model_params));
|
|
if (!model) {
|
|
fprintf(stderr, "%s: error: unable to load model\n", __func__);
|
|
}
|
|
|
|
return model;
|
|
}
|
|
|
|
// Initializes the context with the specified parameters
|
|
llama_context_ptr initialize_context(const llama_model_ptr & model, const int n_ctx) {
|
|
llama_context_params ctx_params = llama_context_default_params();
|
|
ctx_params.n_ctx = n_ctx;
|
|
ctx_params.n_batch = n_ctx;
|
|
|
|
llama_context_ptr context(llama_new_context_with_model(model.get(), ctx_params));
|
|
if (!context) {
|
|
fprintf(stderr, "%s: error: failed to create the llama_context\n", __func__);
|
|
}
|
|
|
|
return context;
|
|
}
|
|
|
|
// Initializes and configures the sampler
|
|
llama_sampler_ptr initialize_sampler() {
|
|
llama_sampler_ptr sampler(llama_sampler_chain_init(llama_sampler_chain_default_params()));
|
|
llama_sampler_chain_add(sampler.get(), llama_sampler_init_min_p(0.05f, 1));
|
|
llama_sampler_chain_add(sampler.get(), llama_sampler_init_temp(0.8f));
|
|
llama_sampler_chain_add(sampler.get(), llama_sampler_init_dist(LLAMA_DEFAULT_SEED));
|
|
|
|
return sampler;
|
|
}
|
|
};
|
|
|
|
// Add a message to `messages` and store its content in `owned_content`
|
|
static void add_message(const char * role, const std::string & text, LlamaData & llama_data,
|
|
std::vector<char_array_ptr> & owned_content) {
|
|
char_array_ptr content(new char[text.size() + 1]);
|
|
std::strcpy(content.get(), text.c_str());
|
|
llama_data.messages.push_back({role, content.get()});
|
|
owned_content.push_back(std::move(content));
|
|
}
|
|
|
|
// Function to apply the chat template and resize `formatted` if needed
|
|
static int apply_chat_template(const LlamaData & llama_data, std::vector<char> & formatted, const bool append) {
|
|
int result = llama_chat_apply_template(llama_data.model.get(), nullptr, llama_data.messages.data(),
|
|
llama_data.messages.size(), append, formatted.data(), formatted.size());
|
|
if (result > static_cast<int>(formatted.size())) {
|
|
formatted.resize(result);
|
|
result = llama_chat_apply_template(llama_data.model.get(), nullptr, llama_data.messages.data(),
|
|
llama_data.messages.size(), append, formatted.data(), formatted.size());
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
// Function to tokenize the prompt
|
|
static int tokenize_prompt(const llama_model_ptr & model, const std::string & prompt,
|
|
std::vector<llama_token> & prompt_tokens) {
|
|
const int n_prompt_tokens = -llama_tokenize(model.get(), prompt.c_str(), prompt.size(), NULL, 0, true, true);
|
|
prompt_tokens.resize(n_prompt_tokens);
|
|
if (llama_tokenize(model.get(), prompt.c_str(), prompt.size(), prompt_tokens.data(), prompt_tokens.size(), true,
|
|
true) < 0) {
|
|
GGML_ABORT("failed to tokenize the prompt\n");
|
|
}
|
|
|
|
return n_prompt_tokens;
|
|
}
|
|
|
|
// Check if we have enough space in the context to evaluate this batch
|
|
static int check_context_size(const llama_context_ptr & ctx, const llama_batch & batch) {
|
|
const int n_ctx = llama_n_ctx(ctx.get());
|
|
const int n_ctx_used = llama_get_kv_cache_used_cells(ctx.get());
|
|
if (n_ctx_used + batch.n_tokens > n_ctx) {
|
|
printf("\033[0m\n");
|
|
fprintf(stderr, "context size exceeded\n");
|
|
return 1;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
// convert the token to a string
|
|
static int convert_token_to_string(const llama_model_ptr & model, const llama_token token_id, std::string & piece) {
|
|
char buf[256];
|
|
int n = llama_token_to_piece(model.get(), token_id, buf, sizeof(buf), 0, true);
|
|
if (n < 0) {
|
|
GGML_ABORT("failed to convert token to piece\n");
|
|
}
|
|
|
|
piece = std::string(buf, n);
|
|
return 0;
|
|
}
|
|
|
|
static void print_word_and_concatenate_to_response(const std::string & piece, std::string & response) {
|
|
printf("%s", piece.c_str());
|
|
fflush(stdout);
|
|
response += piece;
|
|
}
|
|
|
|
// helper function to evaluate a prompt and generate a response
|
|
static int generate(LlamaData & llama_data, const std::string & prompt, std::string & response) {
|
|
std::vector<llama_token> prompt_tokens;
|
|
const int n_prompt_tokens = tokenize_prompt(llama_data.model, prompt, prompt_tokens);
|
|
if (n_prompt_tokens < 0) {
|
|
return 1;
|
|
}
|
|
|
|
// prepare a batch for the prompt
|
|
llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size());
|
|
llama_token new_token_id;
|
|
while (true) {
|
|
check_context_size(llama_data.context, batch);
|
|
if (llama_decode(llama_data.context.get(), batch)) {
|
|
GGML_ABORT("failed to decode\n");
|
|
}
|
|
|
|
// 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_token_is_eog(llama_data.model.get(), new_token_id)) {
|
|
break;
|
|
}
|
|
|
|
std::string piece;
|
|
if (convert_token_to_string(llama_data.model, 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);
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static int parse_arguments(const int argc, const char ** argv, Options & opt) {
|
|
ArgumentParser parser(argv[0]);
|
|
parser.add_argument("-m", opt.model_path, "model");
|
|
parser.add_argument("-p", opt.prompt_non_interactive, "prompt");
|
|
parser.add_argument("-c", opt.n_ctx, "context_size");
|
|
parser.add_argument("-ngl", opt.ngl, "n_gpu_layers");
|
|
if (parser.parse(argc, argv)) {
|
|
return 1;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static int read_user_input(std::string & user) {
|
|
std::getline(std::cin, user);
|
|
return user.empty(); // Indicate an error or empty input
|
|
}
|
|
|
|
// Function to generate a response based on the prompt
|
|
static int generate_response(LlamaData & llama_data, const std::string & prompt, std::string & response) {
|
|
// Set response color
|
|
printf("\033[33m");
|
|
if (generate(llama_data, prompt, response)) {
|
|
fprintf(stderr, "failed to generate response\n");
|
|
return 1;
|
|
}
|
|
|
|
// End response with color reset and newline
|
|
printf("\n\033[0m");
|
|
return 0;
|
|
}
|
|
|
|
// Helper function to apply the chat template and handle errors
|
|
static int apply_chat_template_with_error_handling(const LlamaData & llama_data, std::vector<char> & formatted,
|
|
const bool is_user_input, int & output_length) {
|
|
const int new_len = apply_chat_template(llama_data, formatted, is_user_input);
|
|
if (new_len < 0) {
|
|
fprintf(stderr, "failed to apply the chat template\n");
|
|
return -1;
|
|
}
|
|
|
|
output_length = new_len;
|
|
return 0;
|
|
}
|
|
|
|
// Helper function to handle user input
|
|
static bool handle_user_input(std::string & user_input, const std::string & prompt_non_interactive) {
|
|
if (!prompt_non_interactive.empty()) {
|
|
user_input = prompt_non_interactive;
|
|
return true; // No need for interactive input
|
|
}
|
|
|
|
printf("\033[32m> \033[0m");
|
|
return !read_user_input(user_input); // Returns false if input ends the loop
|
|
}
|
|
|
|
// Function to tokenize the prompt
|
|
static int chat_loop(LlamaData & llama_data, std::string & prompt_non_interactive) {
|
|
std::vector<char_array_ptr> owned_content;
|
|
std::vector<char> fmtted(llama_n_ctx(llama_data.context.get()));
|
|
int prev_len = 0;
|
|
|
|
while (true) {
|
|
// Get user input
|
|
std::string user_input;
|
|
if (!handle_user_input(user_input, prompt_non_interactive)) {
|
|
break;
|
|
}
|
|
|
|
add_message("user", prompt_non_interactive.empty() ? user_input : prompt_non_interactive, llama_data,
|
|
owned_content);
|
|
|
|
int new_len;
|
|
if (apply_chat_template_with_error_handling(llama_data, fmtted, true, new_len) < 0) {
|
|
return 1;
|
|
}
|
|
|
|
std::string prompt(fmtted.begin() + prev_len, fmtted.begin() + new_len);
|
|
std::string response;
|
|
if (generate_response(llama_data, prompt, response)) {
|
|
return 1;
|
|
}
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
static void log_callback(const enum ggml_log_level level, const char * text, void *) {
|
|
if (level == GGML_LOG_LEVEL_ERROR) {
|
|
fprintf(stderr, "%s", text);
|
|
}
|
|
}
|
|
|
|
static bool is_stdin_a_terminal() {
|
|
#if defined(_WIN32)
|
|
HANDLE hStdin = GetStdHandle(STD_INPUT_HANDLE);
|
|
DWORD mode;
|
|
return GetConsoleMode(hStdin, &mode);
|
|
#else
|
|
return isatty(STDIN_FILENO);
|
|
#endif
|
|
}
|
|
|
|
static std::string read_pipe_data() {
|
|
std::ostringstream result;
|
|
result << std::cin.rdbuf(); // Read all data from std::cin
|
|
return result.str();
|
|
}
|
|
|
|
int main(int argc, const char ** argv) {
|
|
Options opt;
|
|
if (parse_arguments(argc, argv, opt)) {
|
|
return 1;
|
|
}
|
|
|
|
if (!is_stdin_a_terminal()) {
|
|
if (!opt.prompt_non_interactive.empty()) {
|
|
opt.prompt_non_interactive += "\n\n";
|
|
}
|
|
|
|
opt.prompt_non_interactive += read_pipe_data();
|
|
}
|
|
|
|
llama_log_set(log_callback, nullptr);
|
|
LlamaData llama_data;
|
|
if (llama_data.init(opt)) {
|
|
return 1;
|
|
}
|
|
|
|
if (chat_loop(llama_data, opt.prompt_non_interactive)) {
|
|
return 1;
|
|
}
|
|
|
|
return 0;
|
|
}
|