summaryrefslogtreecommitdiffstats
path: root/main.c
diff options
context:
space:
mode:
Diffstat (limited to 'main.c')
-rw-r--r--main.c294
1 files changed, 111 insertions, 183 deletions
diff --git a/main.c b/main.c
index 6a2eb36..e89ed7a 100644
--- a/main.c
+++ b/main.c
@@ -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;
}