Add getTopLogits function to Huffman

Author: tobiasvonderheidt

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

@@ -36,4 +36,23 @@ object Huffman {
 
         return cipherBits
     }
+
+    /**
+     * Function to get the top 2^bitsPerToken logits for the last token of the prompt. Keeps track of the corresponding token IDs in a map.
+     *
+     * Parameter `bits_per_word` from Stegasuras was renamed to `bitsPerToken`.
+     *
+     * @param logits Logits for the last token of the prompt (= last row of logits matrix).
+     * @param bitsPerToken Number of bits to encode/decode per cover text token (= height of Huffman tree). Determined by Settings object.
+     * @return Map of top 2^bitsPerToken logits and the corresponding token IDs.
+     */
+    private fun getTopLogits(logits: FloatArray, bitsPerToken: Int = Settings.bitsPerToken): Map<Int, Float> {
+        val topLogits = logits
+            .mapIndexed{ token, logit -> token to logit }   // Convert to List<Pair<Int, Float>> so token IDs won't get lost
+            .sortedByDescending { it.second }               // Sort pairs descending based on logits
+            .take(1 shl bitsPerToken)                    // Take top 2^bitsPerToken pairs
+            .toMap()                                        // Convert to Map<Int, Float> for Huffman tree (ensures there can't be any duplicate token IDs)
+
+        return topLogits
+    }
 }