From e12f3fb88414199226c12b7a3bb972a4a8055c9c Mon Sep 17 00:00:00 2001
From: Tom 7 <499244+tom7@users.noreply.github.com>
Date: Fri, 19 May 2023 12:50:07 -0400
Subject: [PATCH] Fix bug in main.cpp where penalize_nl=false has no effect. It
 modifies the underlying logits array, but at this point we are already
 working on the candidates copy.

---
 examples/main/main.cpp | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/examples/main/main.cpp b/examples/main/main.cpp
index 18673ed2e..3d18a5239 100644
--- a/examples/main/main.cpp
+++ b/examples/main/main.cpp
@@ -428,7 +428,12 @@ int main(int argc, char ** argv) {
                     last_n_tokens.data() + last_n_tokens.size() - last_n_repeat,
                     last_n_repeat, alpha_frequency, alpha_presence);
                 if (!penalize_nl) {
-                    logits[llama_token_nl()] = nl_logit;
+                  for (size_t idx = 0; idx < candidates_p.size; idx++) {
+                    if (candidates_p.data[idx].id == llama_token_nl()) {
+                      candidates_p.data[idx].logit = nl_logit;
+                      break;
+                    }
+                  }
                 }
 
                 if (temp <= 0) {