Add functions to suppress special tokens

Author: tobiasvonderheidt

app/src/main/cpp/hips.cpp

@@ -252,6 +252,33 @@ extern "C" JNIEXPORT jstring JNICALL Java_org_vonderheidt_hips_utils_LlamaCpp_de
     return jString;
 }
 
+/**
+ * Function to check if a token is a special token.
+ *
+ * @param env The JNI environment.
+ * @param thiz Java object this function was called with.
+ * @param token Token ID to check.
+ * @param jCtx Memory address of the context.
+ * @return Boolean that is true if the token special, false otherwise.
+ */
+extern "C" JNIEXPORT jboolean JNICALL Java_org_vonderheidt_hips_utils_LlamaCpp_isSpecial(JNIEnv* env, jobject thiz, jint token, jlong jCtx) {
+    // Cast memory address of the context from Java long to C++ pointer
+    auto cppCtx = reinterpret_cast<llama_context*>(jCtx);
+
+    // Get model the context was created with
+    const llama_model* model = llama_get_model(cppCtx);
+
+    // Check if token is special
+    // Token ID doesn't need casting because jint and llama_token are both just int32_t
+    bool cppIsSpecial = llama_token_is_eog(model, token) || llama_token_is_control(model,token);
+
+    // Cast boolean to return it
+    // static_cast because casting booleans is type safe, unlike reinterpret_cast for casting C++ pointers to Java long
+    auto jIsSpecial = static_cast<jboolean>(cppIsSpecial);
+
+    return jIsSpecial;
+}
+
 /**
  * Function to calculate the logit matrix (i.e. predictions for every token in the prompt).
  *

app/src/main/java/org/vonderheidt/hips/utils/Huffman.kt

@@ -41,6 +41,9 @@ object Huffman {
                 // Only last row of logit matrix is needed as it contains logits corresponding to last token of the prompt
                 val logits = LlamaCpp.getLogits(if (isFirstRun) contextTokens else intArrayOf(sampledToken)).last()
 
+                // Suppress special tokens to avoid early termination before all bits of secret message are encoded
+                LlamaCpp.suppressSpecialTokens(logits)
+
                 // Get top 2^bitsPerToken logits for last token of prompt (= height of Huffman tree)
                 val topLogits = getTopLogits(logits)
 
@@ -124,6 +127,9 @@ object Huffman {
             // Calculate the logit matrix again initially from context tokens, then from last cover text token, and get last row
             val logits = LlamaCpp.getLogits(if (isFirstRun) contextTokens else intArrayOf(coverTextToken)).last()
 
+            // Suppress special tokens
+            LlamaCpp.suppressSpecialTokens(logits)
+
             // Get top 2^bitsPerToken logits
             val topLogits = getTopLogits(logits)
 

app/src/main/java/org/vonderheidt/hips/utils/LlamaCpp.kt

@@ -95,6 +95,25 @@ object LlamaCpp {
         }
     }
 
+    /**
+     * Function to suppress special tokens, i.e. eog (end-of-generation) and control tokens.
+     *
+     * Suppressing eog tokens is needed to avoid early termination when generating a cover text.
+     *
+     * Additionally suppressing control tokens is beneficial because the cover text then can't contain any invisible tokens.
+     * This ensures integrity when using a non-digital communication medium.
+     *
+     * @param logits Logits for the last token of the prompt (= last row of logits matrix).
+     */
+    fun suppressSpecialTokens(logits: FloatArray) {
+        // Suppress special tokens by setting their logits to negative values
+        for (token in logits.indices) {
+            if (isSpecial(token)) {
+                logits[token] = -100f
+            }
+        }
+    }
+
     /**
      * Function to check if a token is the end of a sentence. Needed to complete the last sentence of the cover text.
      *
@@ -193,6 +212,15 @@ object LlamaCpp {
      */
     external fun getLogits(tokens: IntArray, ctx: Long = this.ctx): Array<FloatArray>
 
+    /**
+     * Wrapper for the `llama_token_is_eog` and `llama_token_is_control` functions of llama.cpp. Checks if a token is a special token.
+     *
+     * @param token Token ID to check.
+     * @param ctx Memory address of the context.
+     * @return Boolean that is true if the token special, false otherwise.
+     */
+    private external fun isSpecial(token: Int, ctx: Long = this.ctx): Boolean
+
     /**
      * Wrapper for the `llama_sampler_sample` function of llama.cpp. Samples the next token based on the last one.
      *