Add num_bos num_eos to GenerationConfig

Author: larryliu0820

examples/models/llama/main.cpp

@@ -42,6 +42,18 @@ DEFINE_int32(
     -1,
     "Number of CPU threads for inference. Defaults to -1, which implies we'll use a heuristic to derive the # of performant cores for a specific device.");
 
+DEFINE_int32(
+    num_bos,
+    0,
+    "Number of BOS tokens to prepend to the prompt. Defaults to 0. If > 0, the prompt will be prepended with BOS tokens. This is useful for models that expect one or more BOS token at the start."
+)
+
+DEFINE_int32(
+    num_eos,
+    0,
+    "Number of EOS tokens to append to the prompt. Defaults to 0. If > 0, the prompt will be appended with EOS tokens. This is useful for models that expect one or more EOS token at the end."
+)
+
 DEFINE_bool(warmup, false, "Whether to run a warmup run.");
 
 int32_t main(int32_t argc, char** argv) {

extension/llm/runner/irunner.h

@@ -49,6 +49,10 @@ struct GenerationConfig {
   // Temperature for sampling (higher = more random)
   float temperature = 0.8f;
 
+  // Number of eos and bos to add to the prompt
+  int32_t num_bos = 0;
+  int32_t num_eos = 0;
+
   /**
    * Resolve the maximum number of new tokens to generate based on constraints.
    *

extension/llm/runner/text_llm_runner.cpp

@@ -117,8 +117,8 @@ Error TextLLMRunner::generate_from_pos(
 
   ::tokenizers::Result<std::vector<uint64_t>> encode_res = tokenizer_->encode(
       prompt,
-      /* bos */ 0,
-      /* eos */ 0);
+      /*bos=*/config.num_bos,
+      /*eos=*/config.num_eos);
 
   ET_CHECK_TK_OK_OR_RETURN_ERROR(
       encode_res.error(), "Failed to encode prompt %s", prompt.c_str());