Sampling: Add stop tokens

This was an oversight in the initial release. Now stop tokens can
be provided as an integer value.

Signed-off-by: kingbri <8082010+kingbri1@users.noreply.github.com>

Author: kingbri1

bindings/binding.cpp

@@ -8,6 +8,7 @@
 #include <vector>
 #include <string>
 #include <sstream>
+#include <unordered_set>
 
 // Static vector to hold previous generation tokens
 // TODO: Remove in continuous batch implementation
@@ -749,7 +750,9 @@ const char* InferToReadbackBuffer(
     const char** rewindStrings,
     const unsigned numRewindStrings,
     const char** stoppingStrings,
-    const unsigned numStoppingStrings)
+    const unsigned numStoppingStrings,
+    const unsigned* stoppingTokens,
+    const unsigned numStoppingTokens)
 {
     if (abortCallback != nullptr) {
         llama_set_abort_callback(context, abortCallback, nullptr);
@@ -832,6 +835,13 @@ const char* InferToReadbackBuffer(
         return nullptr;
     }
 
+    // Create stop token set
+    std::unordered_set<unsigned> stopTokenSet;
+    if (stoppingTokens != nullptr && numStoppingTokens > 0) {
+        stopTokenSet.insert(stoppingTokens, stoppingTokens + numStoppingTokens);
+    }
+
+    // Populate string ban trie
     MatchTrie::MatchTrie matchingTrie;
 
     if (rewindStrings != nullptr && numRewindStrings > 0) {
@@ -885,7 +895,7 @@ const char* InferToReadbackBuffer(
 
     while (true) {
         // Abort if callback is fired
-        if (isEnd) {
+        if (isEnd || (!stopTokenSet.empty() && stopTokenSet.find(newTokenId) != stopTokenSet.end())) {
             finishReason = "StopToken";
             stoppedAt = TokenToPiece(model, newTokenId, decodeSpecial).value_or("");
             break;

bindings/binding.h

@@ -110,7 +110,9 @@ extern "C" {
         const char** rewindStrings,
         unsigned numRewindStrings,
         const char** stoppingStrings,
-        unsigned numStoppingStrings);
+        unsigned numStoppingStrings,
+        const unsigned* stoppingTokens,
+        unsigned numStoppingTokens);
 
 #ifdef __cplusplus
 }

bindings/bindings.ts

@@ -719,9 +719,18 @@ export class Model {
 
         const promptPtr = new TextEncoder().encode(prompt + "\0");
 
+        // These are numbers and strings, TS doesn't understand for some reason
+        const stopTokens = params.stop.filter((e) =>
+            typeof e === "number"
+        ) as number[];
+        const stopStrings = params.stop.filter((e) =>
+            typeof e === "string"
+        ) as string[];
+
         // Use the helper function for both arrays
         const rewindPtrArray = pointerArrayFromStrings(params.banned_strings);
-        const stopPtrArray = pointerArrayFromStrings(params.stop);
+        const stopTokensPtr = new Uint32Array(stopTokens);
+        const stopStringsPtr = pointerArrayFromStrings(stopStrings);
 
         const promptBosToken = params.add_bos_token
             ? this.tokenizer.bosToken?.piece
@@ -754,8 +763,10 @@ export class Model {
             seed,
             rewindPtrArray.inner,
             params.banned_strings.length,
-            stopPtrArray.inner,
-            params.stop.length,
+            stopStringsPtr.inner,
+            stopStrings.length,
+            stopTokensPtr,
+            stopTokens.length,
         );
 
         // Read from the read buffer

bindings/symbols.ts

@@ -196,6 +196,8 @@ export default {
             "u32", // count of rewind strings
             "buffer", // const char** stoppingStrings
             "u32", // count of stop strings
+            "buffer", // const unsigned* stoppingTokens
+            "u32", // count of stop tokens
         ],
         result: "pointer" as const,
         nonblocking: true,

common/sampling.ts

@@ -8,7 +8,7 @@ const GenerationOptionsSchema = z.aliasedObject(
             }),
         stop: z.union([
             z.string().transform((str) => [str]),
-            z.array(z.string()),
+            z.array(z.union([z.string(), z.number()])),
         ])
             .nullish().coalesce([])
             .openapi({