Fix DOS index bug

Author: ishaangandhi

common/common.cpp

@@ -1748,6 +1748,12 @@ std::string common_token_to_piece(const struct llama_vocab * vocab, llama_token
     return piece;
 }
 
+bool can_be_detokenized(const struct llama_context * ctx, const std::vector<llama_token> & tokens) {
+    const llama_model * model = llama_get_model(ctx);
+    const llama_vocab * vocab = llama_model_get_vocab(model);
+    return can_detokenize(vocab, tokens.data(), (int32_t) tokens.size());
+}
+
 std::string common_detokenize(const struct llama_context * ctx, const std::vector<llama_token> & tokens, bool special) {
     const llama_model * model = llama_get_model(ctx);
     const llama_vocab * vocab = llama_model_get_vocab(model);

common/common.h

@@ -593,6 +593,10 @@ size_t common_lcs(const llama_tokens & a, const llama_tokens & b);
 // Vocab utils
 //
 
+bool can_be_detokenized(
+    const struct llama_context * ctx,
+    const std::vector<llama_token> & tokens);
+
 // tokenizes a string into a vector of tokens
 // should work similar to Python's `tokenizer.encode`
 std::vector<llama_token> common_tokenize(

examples/server/server.cpp

@@ -2054,6 +2054,11 @@ struct server_context {
             slot.lora = task.params.lora;
         }
 
+        bool can_detokenize = can_be_detokenized(ctx, slot.prompt_tokens);
+        if (!can_detokenize) {
+            send_error(task, "Prompt contains invalid tokens", ERROR_TYPE_INVALID_REQUEST);
+            return false;
+        }
         SLT_DBG(slot, "launching slot : %s\n", safe_json_to_str(slot.to_json()).c_str());
 
         if (slot.n_predict > 0 && slot.params.n_predict > slot.n_predict) {

include/llama.h

@@ -1021,6 +1021,10 @@ extern "C" {
                                int32_t   lstrip,
                                   bool   special);
 
+    /// @details Determine whether the provided tokens are in the vocabulary.
+    /// @return Returns true if every token provided is in the vocabulary.
+    LLAMA_API bool can_detokenize(const struct llama_vocab * vocab, const llama_token * tokens, int32_t n_tokens);
+
     /// @details Convert the provided tokens into text (inverse of llama_tokenize()).
     /// @param text The char pointer must be large enough to hold the resulting text.
     /// @return Returns the number of chars/bytes on success, no more than text_len_max.

src/llama-vocab.cpp

@@ -1279,6 +1279,7 @@ struct llama_vocab::impl {
 
     std::string type_name() const;
 
+    bool is_valid       (llama_token id) const;
     bool is_normal      (llama_token id) const;
     bool is_unknown     (llama_token id) const;
     bool is_control     (llama_token id) const;
@@ -2068,6 +2069,11 @@ std::string llama_vocab::impl::type_name() const{
     }
 }
 
+bool llama_vocab::impl::is_valid(llama_token id) const {
+    GGML_ASSERT(type != LLAMA_VOCAB_TYPE_NONE);
+    return 0 <= id && id < (int32_t) id_to_token.size();
+}
+
 bool llama_vocab::impl::is_normal(llama_token id) const {
     GGML_ASSERT(type != LLAMA_VOCAB_TYPE_NONE);
     return id_to_token[id].attr & LLAMA_TOKEN_ATTR_NORMAL;
@@ -2760,6 +2766,10 @@ std::string llama_vocab::type_name() const{
     return pimpl->type_name();
 }
 
+bool llama_vocab::is_valid(llama_token id) const {
+    return pimpl->is_valid(id);
+}
+
 bool llama_vocab::is_normal(llama_token id) const {
     return pimpl->is_normal(id);
 }
@@ -3262,3 +3272,11 @@ int32_t llama_detokenize(
     return vocab->detokenize(tokens, n_tokens, text, text_len_max, remove_special, unparse_special);
 }
 
+bool can_detokenize(const struct llama_vocab * vocab, const llama_token * tokens, int32_t n_tokens) {
+    for (int32_t i = 0; i < n_tokens; ++i) {
+        if (!vocab->is_valid(tokens[i])) {
+            return false;
+        }
+    }
+    return true;
+}

src/llama-vocab.h

@@ -29,6 +29,7 @@ struct llama_vocab {
 
     std::string type_name() const;
 
+    bool is_valid       (llama_token id) const;
     bool is_normal      (llama_token id) const;
     bool is_unknown     (llama_token id) const;
     bool is_control     (llama_token id) const;