Jelajahi Sumber

llama : sanitize tokens in the upper bound (#9359)

slaren 1 tahun lalu
induk
melakukan
eae597182c
1 mengubah file dengan 2 tambahan dan 2 penghapusan
  1. 2 2
      src/llama.cpp

+ 2 - 2
src/llama.cpp

@@ -16077,7 +16077,7 @@ static int llama_decode_internal(
     }
     }
 
 
     for (uint32_t i = 0; i < n_tokens_all; ++i) {
     for (uint32_t i = 0; i < n_tokens_all; ++i) {
-        if (batch_all.token[i] < 0) {
+        if (batch_all.token[i] < 0 || (uint32_t)batch_all.token[i] >= lctx.model.vocab.n_vocab) {
             LLAMA_LOG_ERROR("%s: invalid token[%d] = %d", __func__, i, batch_all.token[i]);
             LLAMA_LOG_ERROR("%s: invalid token[%d] = %d", __func__, i, batch_all.token[i]);
             return -1;
             return -1;
         }
         }
@@ -16376,7 +16376,7 @@ static int llama_encode_internal(
     }
     }
 
 
     for (uint32_t i = 0; i < n_tokens; ++i) {
     for (uint32_t i = 0; i < n_tokens; ++i) {
-        if (batch.token[i] < 0) {
+        if (batch.token[i] < 0 || (uint32_t)batch.token[i] >= lctx.model.vocab.n_vocab) {
             LLAMA_LOG_ERROR("%s: invalid token[%d] = %d", __func__, i, batch.token[i]);
             LLAMA_LOG_ERROR("%s: invalid token[%d] = %d", __func__, i, batch.token[i]);
             return -1;
             return -1;
         }
         }