From 6726e3f29a2442e58211bc1171694e34a46feba1 Mon Sep 17 00:00:00 2001
From: Georgi Gerganov <ggerganov@gmail.com>
Date: Sat, 7 Sep 2024 21:52:25 +0300
Subject: [PATCH] llama : check that the input tokens are valid

ggml-ci
---
 src/llama.cpp | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/src/llama.cpp b/src/llama.cpp
index 6bbaf9fc9..190564fa4 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -16066,6 +16066,13 @@ static int llama_decode_internal(
         return -1;
     }
 
+    for (uint32_t i = 0; i < n_tokens_all; ++i) {
+        if (batch_all.token[i] < 0) {
+            LLAMA_LOG_ERROR("%s: invalid token[%d] = %d", __func__, i, batch_all.token[i]);
+            return -1;
+        }
+    }
+
     const auto & model   = lctx.model;
     const auto & hparams = model.hparams;
     const auto & cparams = lctx.cparams;
@@ -16358,6 +16365,13 @@ static int llama_encode_internal(
         return -1;
     }
 
+    for (uint32_t i = 0; i < n_tokens; ++i) {
+        if (batch.token[i] < 0) {
+            LLAMA_LOG_ERROR("%s: invalid token[%d] = %d", __func__, i, batch.token[i]);
+            return -1;
+        }
+    }
+
     const auto & model   = lctx.model;
     const auto & hparams = model.hparams;
     const auto & cparams = lctx.cparams;