Change project name, update prompt, tune parameters.HEADmaster

1 files changed, 165 insertions, 1 deletions

diff --git a/main.c b/main.c
index c5ecd8a..a15bfa5 100644
--- a/main.c
+++ b/main.c
@@ -4,9 +4,168 @@
 #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 precise lexicographer. You are defining the EXACT target word provided.\n" \
+	"Do NOT confuse the word with phonetically or visually similar words.\n\n" \
+	"Output format strictly as follows:\n" \
+	"<WORD IN UPPERCASE> (<part of speech>) — <Formality: Conversational|Formal|Literary|Archaic>\n\n" \
+	"DEFINITION:\n" \
+	"<terse, highly accurate definition>\n\n" \
+	"WHEN TO USE:\n" \
+	"<1 sentence on when to use>\n\n" \
+	"EXAMPLES:\n" \
+	"1. <example 1>\n" \
+	"2. <example 2>\n\n" \
+	"Do not include intro, markdown, or extra commentary."
+
+#define CHATML_FMT \
+	"<|im_start|>system\n%s<|im_end|>\n" \
+	"<|im_start|>user\nTarget word: %s\nDefine this specific word:<|im_end|>\n" \
+	"<|im_start|>assistant\n"
+
+#define MAX_TOKENS 300
+
+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_penalties(
+		64,      /* last_n: number of tokens to penalize (64 is standard) */
+		1.1f,    /* repeat_penalty */
+		0.0f,    /* frequency_penalty */
+		0.0f     /* presence_penalty */
+	));
+
+	/* Add Top-P sampling (p = 0.9f, min_keep = 1) */
+	llama_sampler_chain_add(smpl, llama_sampler_init_top_p(0.9f, 1));
+
+	llama_sampler_chain_add(smpl, llama_sampler_init_temp(0.1f)); /* 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;
+	}
+
+	printf("\n");
+	fflush(stdout);
+
+	llama_sampler_free(smpl);
+	llama_free(ctx);
+}
+
 int main(int argc , char *argv[])
 {
 	struct llama_model *model;
@@ -21,6 +180,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 +193,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();