wip: prompt just about works.HEADmaster

2 files changed, 184 insertions, 1 deletions

diff --git a/main.c b/main.c
index c5ecd8a..2f9c75f 100644
--- a/main.c
+++ b/main.c
@@ -4,9 +4,149 @@
 #include <unistd.h>
 
 #include "llama.h"
+#include "mem.h"
 
 #define MODEL_PATH "Qwen2.5-7B-Instruct-Q4_K_M.gguf"
 
+#define SYSTEM_PROMPT \
+	"You are a dictionary assistant. " \
+	"Given a target word, provide its definition and formality on a single line, " \
+	"followed by a blank line, and then a single example sentence.\n\n" \
+	"Output format:\n" \
+	"<word>: <terse definition>. <Formality>.\n\n" \
+	"Example: <example sentence>\n\n" \
+	"Formality must be one of: Formal, Literary, Archaic, or Conversational. " \
+	"Keep the definition terse. Do not include any intro, markdown, or extra chatter."
+
+#define CHATML_FMT \
+	"<|im_start|>system\n%s<|im_end|>\n" \
+	"<|im_start|>user\nDefine the word: %s<|im_end|>\n" \
+	"<|im_start|>assistant\n"
+
+#define MAX_TOKENS 200
+
+static int build_prompt(char **buf, const char *word)
+{
+	char *p;
+	int len, written;
+
+	if (!buf || !word)
+		return -1;
+
+	len = snprintf(NULL, 0, CHATML_FMT, SYSTEM_PROMPT, word);
+	if (len <= 0)
+		return -1;
+
+	p = MALLOC((size_t)len + 1);
+
+	written = snprintf(p, (size_t)len + 1, CHATML_FMT, SYSTEM_PROMPT, word);
+	if (written <= 0) {
+		free(p);
+		return -1;
+	}
+	
+	*buf = p;
+	return written;
+}
+
+static void process_request(struct llama_model *model, const char *word)
+{
+	char *prompt;
+	int prompt_len;
+	int n_prompt_tokens, i;
+
+	struct llama_context_params cparams;
+	struct llama_context *ctx;
+	llama_token *prompt_tokens;
+	struct llama_batch batch;
+	struct llama_sampler_chain_params sparams;
+	struct llama_sampler *smpl;
+
+	llama_token new_token_id;
+
+	prompt = NULL;
+	prompt_len = build_prompt(&prompt, word);
+
+	if (prompt_len <= 0) {
+		fprintf(stderr, "Error: failed to construct prompt for '%s'\n", word);
+		return;
+	}
+	
+	cparams = llama_context_default_params();
+	cparams.n_ctx = 512; /* keep context small to keep RAM usage low */
+	cparams.n_threads = 4; /* todo: tune */
+	
+	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);
+
+	n_prompt_tokens = -llama_tokenize(vocab, prompt, prompt_len, NULL, 
+		0, true, true);
+
+	if (n_prompt_tokens <= 0) {
+		fprintf(stderr, "Error: tokenization sizing failed\n");
+		llama_free(ctx);
+		free(prompt);
+		return;
+	}
+
+	prompt_tokens = MALLOC((size_t)n_prompt_tokens * sizeof(llama_token));
+
+	if (llama_tokenize(vocab, prompt, prompt_len, prompt_tokens, 
+		n_prompt_tokens, true, true) < 0) {
+		fprintf(stderr, "Error: Tokenization execution failed\n");
+		free(prompt_tokens);
+		llama_free(ctx);
+		free(prompt);
+		return;
+	}
+
+	free(prompt); /* no longer required */
+
+	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_temp(0.2f)); /* todo: tune */
+	llama_sampler_chain_add(smpl, llama_sampler_init_dist(LLAMA_DEFAULT_SEED));
+
+	for (i = 0; i < MAX_TOKENS; i++) {	
+		new_token_id = llama_sampler_sample(smpl, ctx, -1);
+		if (llama_vocab_is_eog(vocab, new_token_id))
+			break;
+
+		char buf[128];
+		int n = llama_token_to_piece(vocab, new_token_id, buf, 
+			sizeof(buf), 0, true);
+
+		if (n > 0) {
+			fwrite(buf, 1, (size_t)n, stdout);
+			fflush(stdout);
+		}
+
+		batch = llama_batch_get_one(&new_token_id, 1);
+		if (llama_decode(ctx, batch) != 0)
+			break;
+	}
+
+	llama_sampler_free(smpl);
+	llama_free(ctx);
+}
+
 int main(int argc , char *argv[])
 {
 	struct llama_model *model;
@@ -21,6 +161,9 @@ int main(int argc , char *argv[])
 	if (pledge("stdio rpath", NULL) == -1)
 		err(1, "initial pledge failed");
 
+	if (argc < 2)
+		errx(1, "usage: %s [prompt]", argv[0]);
+
 	llama_backend_init();
 
 	mparams = llama_model_default_params();
@@ -31,10 +174,12 @@ int main(int argc , char *argv[])
 	if (!model)
 		errx(1, "failed to load model from file %s", MODEL_PATH);
 
-	/* drop rpath after loading the model */
 	if (pledge("stdio", NULL) == -1)
 		err(1, "secondary pledge failed");
 
+	const char *word = argv[1];
+	process_request(model, word);
+
 	llama_model_free(model);
 	llama_backend_free();
 
diff --git a/mem.h b/mem.h
new file mode 100644
index 0000000..2cdb36e
--- /dev/null
+++ b/mem.h
@@ -0,0 +1,38 @@
+#ifndef MEM_H
+#define MEM_H
+
+#include <err.h>
+#include <stdlib.h>
+
+#define MALLOC(s)      xmalloc((s), __FILE__, __LINE__)
+#define CALLOC(n, s)   xcalloc((n), (s), __FILE__, __LINE__)
+#define REALLOC(p, s)  xrealloc((p), (s), __FILE__, __LINE__)
+
+static inline void *xmalloc(size_t s, const char *file, int line)
+{
+	void *p;
+
+	if (!(p = malloc(s)))
+		err(1, "%s:%d: malloc", file, line);
+	return p;
+}
+
+static inline void *xcalloc(size_t n, size_t s, const char *file, int line)
+{
+	void *p;
+
+	if (!(p = calloc(n, s)))
+		err(1, "%s:%d: calloc", file, line);
+	return p;
+}
+
+static inline void *xrealloc(void *ptr, size_t s, const char *file, int line)
+{
+	void *p;
+
+	if (!(p = realloc(ptr, s)))
+		err(1, "%s:%d: realloc", file, line);
+	return p;
+}
+
+#endif /* MEM_H */