diff options
| -rw-r--r-- | Makefile | 9 | ||||
| -rw-r--r-- | llm.c | 200 | ||||
| -rw-r--r-- | llm.h | 13 | ||||
| -rw-r--r-- | main.c | 294 |
4 files changed, 330 insertions, 186 deletions
@@ -2,7 +2,7 @@ CC = cc CXX = c++ CFLAGS = -march=native -O2 -std=c11 -Wall -Wextra CXXFLAGS = -march=native -O2 -O2 -Wall -Wextra -INCLUDES = -Ideps/include +INCLUDES = -I. -Ideps/include LLAMA_TAG = b10107 LLAMA_URL = https://github.com/ggml-org/llama.cpp/archive/refs/tags/$(LLAMA_TAG).tar.gz @@ -18,8 +18,8 @@ STATIC_LIBS = deps/lib/libllama.a \ SYS_LIBS = -static -pthread -lm -lc++ -lc++abi TARGET = lex -SRC = main.c -OBJS = main.o +SRC = llm.c main.c +OBJS = $(SRC:.c=.o) .PHONY: all clean distclean @@ -40,6 +40,9 @@ $(TARGET): deps/lib/libllama.a $(OBJS) .c.o: $(CC) $(CFLAGS) $(INCLUDES) -c $< -o $@ +# Rebuild objects if the shared header changes +$(OBJS): llm.h + # Download llama source and build statically deps/lib/libllama.a deps/lib/libggml.a deps/lib/libggml-cpu.a: @mkdir -p deps/lib deps/include @@ -0,0 +1,200 @@ +#include <err.h> + +#include "llm.h" +#include "mem.h" + +#define MAX_TOKENS 300 + +#define SYSTEM_PROMPT \ + "You are a dictionary assistant. Define the target word in the context provided.\n" \ + "DO NOT repeat target word in definition. Choose one or two of Formality levels.\n\n" \ + "FOLLOW the output format:\n" \ + "<WORD> (<part of speech>) — <Formality: Conversational|Formal|Literary|Figurative|Archaic>\n\n" \ + "DEFINITION:\n" \ + "<terse, highly accurate definition in context>\n\n" \ + "GENERAL TONE:\n" \ + "<One terse sentence of describing general tone. One terse sentence of standard conversational synonyms.>\n\n" \ + "SYNONYMS:\n" \ + "<Three to four standard conversational synonyms.>\n\n" \ + "EXAMPLES:\n" \ + "1. <example 1>\n" \ + "2. <example 2>\n\n" + +struct llama_model *llm_init(const char *model_path) +{ + struct llama_model *model; + struct llama_model_params mparams; + + llama_backend_init(); + + mparams = llama_model_default_params(); + mparams.n_gpu_layers = 0; /* force all layers onto CPU */ + mparams.load_mode = LLAMA_LOAD_MODE_MMAP; + + model = llama_model_load_from_file(model_path, mparams); + return model; +} + +void llm_free(struct llama_model *model) +{ + llama_model_free(model); + llama_backend_free(); +} + +void llm_process_request(struct llama_model *model, + const char *user_prompt, FILE *out) +{ + int n_prompt_tokens, i; + char *prompt; + int prompt_len; + + llama_token *prompt_tokens; + llama_token new_token_id; + + struct llama_batch batch; + struct llama_context *ctx; + struct llama_context_params cparams; + struct llama_sampler *smpl; + struct llama_sampler_chain_params sparams; + + struct llama_chat_message messages[] = { + { "system", SYSTEM_PROMPT }, + { "user", user_prompt } + }; + + const char *tmpl = llama_model_chat_template(model, NULL); + if (!tmpl) + fprintf(stderr, "Warning: model has no embedded chat template\n"); + + prompt_len = llama_chat_apply_template(NULL, messages, 2, true, NULL, 0); + if (prompt_len <= 0) { + fprintf(stderr, "Error: failed to calculate chat template size\n"); + return; + } + + prompt = MALLOC((size_t)prompt_len + 1); + if (llama_chat_apply_template(NULL, messages, 2, true, prompt, prompt_len + 1) < 0) { + fprintf(stderr, "Error: failed to apply chat template\n"); + free(prompt); + return; + } + + cparams = llama_context_default_params(); + cparams.n_ctx = 1024; /* context size in tokens */ + cparams.n_threads = 4; + cparams.n_threads_batch = 4; + + ctx = llama_init_from_model(model, cparams); + if (!ctx) { + fprintf(stderr, "Error: failed to create context\n"); + free(prompt); + return; + } + + const struct llama_vocab *vocab = llama_model_get_vocab(model); + if (!vocab) { + fprintf(stderr, "Error: failed to obtain model vocabulary\n"); + free(prompt); + llama_free(ctx); + return; + } + + n_prompt_tokens = -llama_tokenize(vocab, + prompt, prompt_len, NULL, 0, false, true); + + if (n_prompt_tokens <= 0) { + fprintf(stderr, "Error: tokenization sizing failed\n"); + free(prompt); + llama_free(ctx); + return; + } + + if (n_prompt_tokens + MAX_TOKENS > (int)cparams.n_ctx) { + fprintf(stderr, "Error: token count exceeds context size\n"); + free(prompt); + llama_free(ctx); + return; + } + + prompt_tokens = MALLOC((size_t)n_prompt_tokens * sizeof(llama_token)); + + if (llama_tokenize(vocab, prompt, prompt_len, prompt_tokens, + n_prompt_tokens, false, true) < 0) { + fprintf(stderr, "Error: Tokenization failed\n"); + free(prompt); + free(prompt_tokens); + llama_free(ctx); + return; + } + + free(prompt); /* Prompt buffer is fully tokenized and no longer needed */ + + /* Ingest prompt tokens in one batch (parallelizes matrix + * multiplications across tokens in the batch) */ + batch = llama_batch_get_one(prompt_tokens, n_prompt_tokens); + + if (llama_decode(ctx, batch) != 0) { + fprintf(stderr, "Error: Prompt evaluation failed\n"); + free(prompt_tokens); + llama_free(ctx); + return; + } + + free(prompt_tokens); + + /* Generation loop */ + sparams = llama_sampler_chain_default_params(); + smpl = llama_sampler_chain_init(sparams); + + llama_sampler_chain_add(smpl, llama_sampler_init_penalties( + 64, /* last_n: lookback window (64 is standard) */ + 1.1f, /* repeat_penalty */ + 0.0f, /* frequency_penalty */ + 0.0f /* presence_penalty */ + )); + + /* Pick the top token */ + llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); + + for (i = 0; i < MAX_TOKENS; i++) { + /* Model outputs next tokens for every token in the prompt. + * We need the one after the last token in the prompt */ + new_token_id = llama_sampler_sample(smpl, ctx, -1); + + /* Tell the sampler chain which token was chosen */ + llama_sampler_accept(smpl, new_token_id); + + /* Check for end-of-generation (EOG) tokens: + * EOS: end-of-sequence + * EOT: end-of-turn */ + if (llama_vocab_is_eog(vocab, new_token_id)) + break; + + /* Convert numeric token id to printable text */ + char buf[128]; + int n = llama_token_to_piece(vocab, new_token_id, buf, + sizeof(buf), 0, false); + + if (n > 0) { + fwrite(buf, 1, (size_t)n, out); + fflush(out); + } + + /* Create batch with 1 token for the next forward pass */ + batch = llama_batch_get_one(&new_token_id, 1); + + if (llama_decode(ctx, batch) != 0) { + fprintf(stderr, "llama_decode failed!\n"); + break; + } + } + + fprintf(out, "\n\n"); + fflush(out); + + //llama_perf_context_print(ctx); + + llama_sampler_free(smpl); + llama_free(ctx); +} + @@ -0,0 +1,13 @@ +#ifndef LLM_H +#define LLM_H + +#include "llama.h" + +struct llama_model *llm_init(const char *model_path); + +void llm_process_request(struct llama_model *model, + const char *user_prompt, FILE *out); + +void llm_free(struct llama_model *model); + +#endif /* LLM_H */ @@ -2,229 +2,157 @@ #include <stdbool.h> #include <ctype.h> #include <err.h> +#include <errno.h> +#include <signal.h> +#include <string.h> #include <unistd.h> +#include <sys/un.h> +#include <sys/stat.h> +#include <sys/socket.h> -#include "mem.h" -#include "llama.h" +#include "llm.h" +#define SOCK_PATH "/var/run/lex.sock" #define MODEL_PATH "qwen2.5-3b-instruct-q4_k_m.gguf" -#define SYSTEM_PROMPT \ - "You are a dictionary assistant. Define the target word in the context provided.\n" \ - "DO NOT repeat target word in definition. Choose one or two of Formality levels.\n\n" \ - "FOLLOW the output format:\n" \ - "<WORD> (<part of speech>) — <Formality: Conversational|Formal|Literary|Figurative|Archaic>\n\n" \ - "DEFINITION:\n" \ - "<terse, highly accurate definition in context>\n\n" \ - "GENERAL TONE:\n" \ - "<One terse sentence of describing general tone. One terse sentence of standard conversational synonyms.>\n\n" \ - "SYNONYMS:\n" \ - "<Three to four standard conversational synonyms.>\n\n" \ - "EXAMPLES:\n" \ - "1. <example 1>\n" \ - "2. <example 2>\n\n" - -#define MAX_TOKENS 300 - -static void process_request(struct llama_model *model, const char *user_prompt) +static void usage(const char *bin) { - int n_prompt_tokens, i; - char *prompt; - int prompt_len; - - llama_token *prompt_tokens; - llama_token new_token_id; - - struct llama_batch batch; - struct llama_context *ctx; - struct llama_context_params cparams; - struct llama_sampler *smpl; - struct llama_sampler_chain_params sparams; - - struct llama_chat_message messages[] = { - { "system", SYSTEM_PROMPT }, - { "user", user_prompt } - }; - - const char *tmpl = llama_model_chat_template(model, NULL); - if (!tmpl) - fprintf(stderr, "Warning: model has no embedded chat template\n"); - - prompt_len = llama_chat_apply_template(NULL, messages, 2, true, NULL, 0); - if (prompt_len <= 0) { - fprintf(stderr, "Error: failed to calculate chat template size\n"); - return; - } - - prompt = MALLOC((size_t)prompt_len + 1); - if (llama_chat_apply_template(NULL, messages, 2, true, prompt, prompt_len + 1) < 0) { - fprintf(stderr, "Error: failed to apply chat template\n"); - free(prompt); - return; - } - - cparams = llama_context_default_params(); - cparams.n_ctx = 1024; /* context size in tokens */ - cparams.n_threads = 4; - cparams.n_threads_batch = 4; - - ctx = llama_init_from_model(model, cparams); - if (!ctx) { - fprintf(stderr, "Error: failed to create context\n"); - free(prompt); - return; - } - - const struct llama_vocab *vocab = llama_model_get_vocab(model); - if (!vocab) { - fprintf(stderr, "Error: failed to obtain model vocabulary\n"); - free(prompt); - llama_free(ctx); - return; - } - - n_prompt_tokens = -llama_tokenize(vocab, - prompt, prompt_len, NULL, 0, false, true); + errx(1, "usage: %s [-m model_path] [-d] [prompt]", bin); +} - if (n_prompt_tokens <= 0) { - fprintf(stderr, "Error: tokenization sizing failed\n"); - free(prompt); - llama_free(ctx); - return; - } +static void loop(struct llama_model *model, int lsock) +{ + int csock, c; + FILE *cfp; + size_t len; + char line[4096]; + + for (;;) { + csock = accept(lsock, NULL, NULL); + if (csock == -1) { + if (errno != EINTR) + warn("accept()"); + continue; + } - if (n_prompt_tokens + MAX_TOKENS > (int)cparams.n_ctx) { - fprintf(stderr, "Error: token count exceeds context size\n"); - free(prompt); - llama_free(ctx); - return; - } + if (!(cfp = fdopen(csock, "r+"))) { + warn("cfp fdopen()"); + close(csock); + continue; + } - prompt_tokens = MALLOC((size_t)n_prompt_tokens * sizeof(llama_token)); + if (fgets(line, sizeof(line), cfp) != NULL) { + len = strcspn(line, "\n"); + if (line[len] == '\0' && len == sizeof(line) - 1) { + warnx("request too large"); + // drain and skip + while ((c = fgetc(cfp)) != '\n' && c != EOF) + ; + fclose(cfp); + continue; + } + + line[len] = '\0'; + if (line[0] != '\0') + llm_process_request(model, line, cfp); + } - if (llama_tokenize(vocab, prompt, prompt_len, prompt_tokens, - n_prompt_tokens, false, true) < 0) { - fprintf(stderr, "Error: Tokenization failed\n"); - free(prompt); - free(prompt_tokens); - llama_free(ctx); - return; + fclose(cfp); } +} - free(prompt); /* Prompt buffer is fully tokenized and no longer needed */ +int sock_init(void) +{ + int sock; + struct sockaddr_un addr; - /* Ingest prompt tokens in one batch (parallelizes matrix - * multiplications across tokens in the batch) */ - batch = llama_batch_get_one(prompt_tokens, n_prompt_tokens); + sock = socket(AF_UNIX, SOCK_STREAM, 0); + if (sock == -1) + err(1, "socket()"); - if (llama_decode(ctx, batch) != 0) { - fprintf(stderr, "Error: Prompt evaluation failed\n"); - free(prompt_tokens); - llama_free(ctx); - return; - } + memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + strlcpy(addr.sun_path, SOCK_PATH, sizeof(addr.sun_path)); - free(prompt_tokens); + unlink(SOCK_PATH); /* remove stale socket from a previous run */ + + if (bind(sock, (struct sockaddr *)&addr, sizeof(addr)) == -1) + err(1, "bind()"); - /* Generation loop */ - sparams = llama_sampler_chain_default_params(); - smpl = llama_sampler_chain_init(sparams); + if (listen(sock, 1) == -1) /* backlog = 1 */ + err(1, "listen()"); - llama_sampler_chain_add(smpl, llama_sampler_init_penalties( - 64, /* last_n: lookback window (64 is standard) */ - 1.1f, /* repeat_penalty */ - 0.0f, /* frequency_penalty */ - 0.0f /* presence_penalty */ - )); + return sock; +} - /* Pick the top token */ - llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); +int main(int argc , char *argv[]) +{ + int daemon_mode, sock, opt; + const char *prog; + const char *prompt; + const char *model_path; + struct llama_model *model; - for (i = 0; i < MAX_TOKENS; i++) { - /* Model outputs next tokens for every token in the prompt. - * We need the one after the last token in the prompt */ - new_token_id = llama_sampler_sample(smpl, ctx, -1); + if (pledge("stdio rpath wpath cpath unix", NULL) == -1) + err(1, "pledge failed"); - /* Tell the sampler chain which token was chosen */ - llama_sampler_accept(smpl, new_token_id); + sock = -1; + daemon_mode = 0; + prompt = NULL; + prog = argv[0]; + model_path = MODEL_PATH; - /* Check for end-of-generation (EOG) tokens: - * EOS: end-of-sequence - * EOT: end-of-turn */ - if (llama_vocab_is_eog(vocab, new_token_id)) + while ((opt = getopt(argc, argv, "dm:")) != -1) { + switch (opt) { + case 'd': + daemon_mode = 1; break; - - /* Convert numeric token id to printable text */ - char buf[128]; - int n = llama_token_to_piece(vocab, new_token_id, buf, - sizeof(buf), 0, false); - - if (n > 0) { - fwrite(buf, 1, (size_t)n, stdout); - fflush(stdout); - } - - /* Create batch with 1 token for the next forward pass */ - batch = llama_batch_get_one(&new_token_id, 1); - - if (llama_decode(ctx, batch) != 0) { - fprintf(stderr, "llama_decode failed!\n"); + case 'm': + model_path = optarg; break; + default: + usage(prog); } } - printf("\n\n"); - fflush(stdout); - - //llama_perf_context_print(ctx); + argc -= optind; + argv += optind; - llama_sampler_free(smpl); - llama_free(ctx); -} - -int main(int argc , char *argv[]) -{ - const char *prompt; - const char *model_path; - struct llama_model *model; - struct llama_model_params mparams; - - if (argc == 2) { - model_path = MODEL_PATH; - prompt = argv[1]; - } else if (argc == 3) { - model_path = argv[1]; - prompt = argv[2]; + if (daemon_mode) { + if (argc != 0) // -d takes no positional args + usage(prog); } else { - errx(1, "usage: %s [model_path] <prompt>", argv[0]); + if (argc != 1) // normal mode requires one prompt + usage(prog); + + prompt = argv[0]; } if (unveil(model_path, "r") == -1) err(1, "unveil %s failed", model_path); - + if (daemon_mode && unveil(SOCK_PATH, "rwc") == -1) + err(1, "unveil %s failed", SOCK_PATH); if (unveil(NULL, NULL) == -1) err(1, "unveil lock failed"); - if (pledge("stdio rpath", NULL) == -1) - err(1, "initial pledge failed"); - - llama_backend_init(); - - mparams = llama_model_default_params(); - mparams.n_gpu_layers = 0; /* force all layers onto CPU */ - mparams.load_mode = LLAMA_LOAD_MODE_MMAP; + if (daemon_mode) { + sock = sock_init(); + if (pledge("stdio rpath unix", NULL) == -1) + err(1, "Failed to narrow pledge after sock_init()"); + } - model = llama_model_load_from_file(model_path, mparams); + model = llm_init(model_path); if (!model) - errx(1, "failed to load model from file %s", model_path); + errx(1, "Failed to load model"); - if (pledge("stdio", NULL) == -1) - err(1, "secondary pledge failed"); + if (daemon_mode) + loop(model, sock); /* doesn't return */ - process_request(model, prompt); + if (pledge("stdio", NULL) == -1) + err(1, "Failed to narrow pledge for normal mode"); - llama_model_free(model); - llama_backend_free(); + llm_process_request(model, prompt, stdout); + llm_free(model); return 0; } |
