From dfeda32abdcf35b2b8f4e98ea82b35409d6465e9 Mon Sep 17 00:00:00 2001
From: vvhg1 <victor.hobe.gelting@gmail.com>
Date: Fri, 6 Oct 2023 18:26:18 +0200
Subject: [PATCH] infill tokens correction

---
 examples/infill/infill.cpp | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/examples/infill/infill.cpp b/examples/infill/infill.cpp
index 9ec75ce42..9542d7b6c 100644
--- a/examples/infill/infill.cpp
+++ b/examples/infill/infill.cpp
@@ -234,9 +234,12 @@ int main(int argc, char ** argv) {
     LOG("add_bos: %d\n", add_bos);
 
     std::vector<llama_token> embd_inp;
-    std::vector<llama_token> inp_pfx = ::llama_tokenize(ctx, params.input_prefix, add_bos);
-    std::vector<llama_token> inp_sfx = ::llama_tokenize(ctx, params.input_suffix, add_bos);
+    std::vector<llama_token> inp_pfx = ::llama_tokenize(ctx, params.input_prefix, false);
+    std::vector<llama_token> inp_sfx = ::llama_tokenize(ctx, params.input_suffix, false);
     inp_pfx.insert(inp_pfx.begin(), llama_token_prefix(ctx));
+    if (add_bos) {
+        inp_pfx.insert(inp_pfx.begin(), llama_token_bos(ctx));
+    }
     inp_sfx.insert(inp_sfx.begin(), llama_token_suffix(ctx));
     embd_inp = inp_pfx;
     embd_inp.insert(embd_inp.end(), inp_sfx.begin(), inp_sfx.end());