Skip to content

Commit eae5971

Browse files
authored
llama : sanitize tokens in the upper bound (ggml-org#9359)
1 parent 00b02bb commit eae5971

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
@@ -16077,7 +16077,7 @@ static int llama_decode_internal(
1607716077
}
1607816078

1607916079
for (uint32_t i = 0; i < n_tokens_all; ++i) {
16080-
if (batch_all.token[i] < 0) {
16080+
if (batch_all.token[i] < 0 || (uint32_t)batch_all.token[i] >= lctx.model.vocab.n_vocab) {
1608116081
LLAMA_LOG_ERROR("%s: invalid token[%d] = %d", __func__, i, batch_all.token[i]);
1608216082
return -1;
1608316083
}
@@ -16376,7 +16376,7 @@ static int llama_encode_internal(
1637616376
}
1637716377

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

0 commit comments

Comments
 (0)