Skip to content

Commit 5f10d45

Browse files
slarenarthw
authored andcommitted
llama : sanitize tokens in the upper bound (ggml-org#9359)
1 parent 6739008 commit 5f10d45

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

src/llama.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16081,7 +16081,7 @@ static int llama_decode_internal(
1608116081
}
1608216082

1608316083
for (uint32_t i = 0; i < n_tokens_all; ++i) {
16084-
if (batch_all.token[i] < 0) {
16084+
if (batch_all.token[i] < 0 || (uint32_t)batch_all.token[i] >= lctx.model.vocab.n_vocab) {
1608516085
LLAMA_LOG_ERROR("%s: invalid token[%d] = %d", __func__, i, batch_all.token[i]);
1608616086
return -1;
1608716087
}
@@ -16380,7 +16380,7 @@ static int llama_encode_internal(
1638016380
}
1638116381

1638216382
for (uint32_t i = 0; i < n_tokens; ++i) {
16383-
if (batch.token[i] < 0) {
16383+
if (batch.token[i] < 0 || (uint32_t)batch.token[i] >= lctx.model.vocab.n_vocab) {
1638416384
LLAMA_LOG_ERROR("%s: invalid token[%d] = %d", __func__, i, batch.token[i]);
1638516385
return -1;
1638616386
}

0 commit comments

Comments
 (0)