llama : add method for getting padding token id

fairydreaming · Jun 14, 2024 · 442e276 · 442e276
1 parent b885765
commit 442e276
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/examples/main/main.cpp b/examples/main/main.cpp
@@ -510,7 +510,7 @@ int main(int argc, char ** argv) {
     }
 
     embd_inp.clear();
-    embd_inp.push_back(0);
+    embd_inp.push_back(llama_token_pad(model));
 
     while ((n_remain != 0 && !is_antiprompt) || params.interactive) {
         // predict

diff --git a/llama.cpp b/llama.cpp
@@ -19344,6 +19344,10 @@ llama_token llama_token_eot(const struct llama_model * model) {
     return model->vocab.special_eot_id;
 }
 
+llama_token llama_token_pad(const struct llama_model * model) {
+    return model->vocab.special_pad_id;
+}
+
 int32_t llama_tokenize(
     const struct llama_model * model,
                   const char * text,

diff --git a/llama.h b/llama.h
@@ -861,6 +861,7 @@ extern "C" {
     LLAMA_API llama_token llama_token_cls(const struct llama_model * model); // classification
     LLAMA_API llama_token llama_token_sep(const struct llama_model * model); // sentence separator
     LLAMA_API llama_token llama_token_nl (const struct llama_model * model); // next-line
+    LLAMA_API llama_token llama_token_pad(const struct llama_model * model); // padding
 
     // Returns -1 if unknown, 1 for true or 0 for false.
     LLAMA_API int32_t         llama_add_bos_token(const struct llama_model * model);