Replace random sampler with greedy sampler, add inline comments explaining LLAMA code.

1 files changed, 35 insertions, 7 deletions

diff --git a/main.c b/main.c
index d311766..d79a9e3 100644
--- a/main.c
+++ b/main.c
@@ -62,9 +62,11 @@ static void process_request(struct llama_model *model, const char *word)
 	llama_token *prompt_tokens;
 	llama_token new_token_id;
 
+	/* Data structure used to pass tokens into llama_decode() */
+	struct llama_batch batch;
+
 	struct llama_context *ctx;
 	struct llama_context_params cparams;
-	struct llama_batch batch;
 	struct llama_sampler *smpl;
 	struct llama_sampler_chain_params sparams;
 
@@ -77,7 +79,7 @@ static void process_request(struct llama_model *model, const char *word)
 	}
 	
 	cparams = llama_context_default_params();
-	cparams.n_ctx = 512;    /* context size -> RAM usage */
+	cparams.n_ctx = 512;    /* context size in tokens */
 	cparams.n_threads = 4;
 	
 	ctx = llama_init_from_model(model, cparams);
@@ -112,7 +114,10 @@ static void process_request(struct llama_model *model, const char *word)
 
 	free(prompt); /* prompt string no longer required */
 
+	/* 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);
@@ -127,22 +132,34 @@ static void process_request(struct llama_model *model, const char *word)
 	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) */
+		64,      /* last_n: lookback window (64 is standard) */
 		1.1f,    /* repeat_penalty */
 		0.0f,    /* frequency_penalty */
 		0.0f     /* presence_penalty */
 	));
 
-	llama_sampler_chain_add(smpl, llama_sampler_init_top_p(0.9f, 1));
-	llama_sampler_chain_add(smpl, llama_sampler_init_temp(0.1f));
-	llama_sampler_chain_add(smpl, llama_sampler_init_dist(LLAMA_DEFAULT_SEED));
+	/* 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);
+
+		/* Check for end-of-generation (EOG) tokens:
+		 * EOS: end-of-sequence
+		 * EOT: end-of-turn
+		 * Generates garbage until token limit hit or context window 
+		 * exhausted, if omitted */
 		if (llama_vocab_is_eog(vocab, new_token_id))
 			break;
 
+		/* LLMS process words as sub-word tokens. 
+		 * E.g.: unbelievable -> ["un", "believ", "able"] 
+		 * 128-byte buffer is sufficient */
 		char buf[128];
+
+		/* Convert numeric token id to printable text */
 		int n = llama_token_to_piece(vocab, new_token_id, buf, 
 			sizeof(buf), 0, true);
 
@@ -151,9 +168,20 @@ static void process_request(struct llama_model *model, const char *word)
 			fflush(stdout);
 		}
 
+		/* Create batch with 1 token: 
+		 * Prompt has been evaluated. Here, we generate one token at a time. 
+		 * See autoregressive (AR), diffusion (dLLM), non-autoregressive (NAR), 
+		 * speculative/MTP for alternative frameworks. */
 		batch = llama_batch_get_one(&new_token_id, 1);
-		if (llama_decode(ctx, batch) != 0)
+
+		/* llama_decode(): the CPU-heavy forward pass through the transformer model:
+		 *   - allocates memory and KV cache
+		 *   - runs matrix multiplications
+		 *   - generates logits (output prediction vectors) */
+		if (llama_decode(ctx, batch) != 0) {
+			fprintf(stderr, "llama_decode failed!\n");
 			break;
+		}
 	}
 
 	printf("\n");