1 files changed, 19 insertions, 2 deletions

diff --git a/main.c b/main.c
index d79a9e3..602a8fd 100644
--- a/main.c
+++ b/main.c
@@ -81,6 +81,7 @@ static void process_request(struct llama_model *model, const char *word)
 	cparams = llama_context_default_params();
 	cparams.n_ctx = 512;    /* context size in tokens */
 	cparams.n_threads = 4;
+	cparams.n_threads_batch = 4;
 	
 	ctx = llama_init_from_model(model, cparams);
 	if (!ctx) {
@@ -90,14 +91,27 @@ static void process_request(struct llama_model *model, const char *word)
 	}
 
 	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, true, 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;
 	}
 
@@ -106,9 +120,9 @@ static void process_request(struct llama_model *model, const char *word)
 	if (llama_tokenize(vocab, prompt, prompt_len, prompt_tokens, 
 		n_prompt_tokens, true, true) < 0) {
 		fprintf(stderr, "Error: Tokenization failed\n");
+		free(prompt);
 		free(prompt_tokens);
 		llama_free(ctx);
-		free(prompt);
 		return;
 	}
 
@@ -146,6 +160,9 @@ static void process_request(struct llama_model *model, const char *word)
 		 * 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
@@ -161,7 +178,7 @@ static void process_request(struct llama_model *model, const char *word)
 
 		/* Convert numeric token id to printable text */
 		int n = llama_token_to_piece(vocab, new_token_id, buf, 
-			sizeof(buf), 0, true);
+			sizeof(buf), 0, false);
 
 		if (n > 0) {
 			fwrite(buf, 1, (size_t)n, stdout);