From 36a36c32a3814c036b6290735ba9d4c9111653c7 Mon Sep 17 00:00:00 2001
From: klosax <131523366+klosax@users.noreply.github.com>
Date: Tue, 1 Aug 2023 14:44:28 +0200
Subject: [PATCH] Update gptneox-main.cpp

---
 gptneox-main.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/gptneox-main.cpp b/gptneox-main.cpp
index fa95cb772..235fca7ce 100644
--- a/gptneox-main.cpp
+++ b/gptneox-main.cpp
@@ -716,6 +716,9 @@ int main(int argc, char ** argv) {
 
     }
 
+    uint32_t eos_token_id = 0;
+    int keyidx = gguf_find_key(ggufctx, "tokenizer.ggml.eos_token_id"); if( keyidx != -1 ) {  eos_token_id = gguf_get_val_u32(ggufctx, keyidx); }    
+
     int n_past = 0;
 
     int64_t t_sample_us  = 0;
@@ -794,7 +797,7 @@ int main(int argc, char ** argv) {
         fflush(stdout);
 
         // end of text token
-        if (embd.back() == 0) {
+        if (embd.back() == eos_token_id) {
             break;
         }
     }