add <|eom_id|> to stopTokens

meta-llama · Dec 6, 2024 · 74526a2 · 74526a2
1 parent 935a923
commit 74526a2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/...ient-local/src/main/kotlin/com/llama/llamastack/client/local/InferenceServiceLocalImpl.kt b/...ient-local/src/main/kotlin/com/llama/llamastack/client/local/InferenceServiceLocalImpl.kt
@@ -29,7 +29,7 @@ constructor(
     private var sequenceLengthKey: String = "seq_len"
 
     override fun onResult(p0: String?) {
-        if (p0.equals(PromptFormatLocal.getStopToken(modelName))) {
+        if (PromptFormatLocal.getStopTokens(modelName).any { it == p0 }) {
             onResultComplete = true
             return
         }
@@ -62,8 +62,9 @@ constructor(
             PromptFormatLocal.getTotalFormattedPrompt(params.messages(), modelName)
 
         // Developer can pass in their sequence length but if not then it will default to a
-        // particular dynamic value. This is to ensure enough value is provided to give a reasonably complete response.
-        // 0.75 is the approximate words per token. And 64 is buffer for tokens for generate response.
+        // particular dynamic value. This is to ensure enough value is provided to give a reasonably
+        // complete response. 0.75 is the approximate words per token. And 64 is buffer for tokens
+        // for generate response.
         val seqLength =
             params._additionalQueryParams().values(sequenceLengthKey).lastOrNull()?.toInt()
                 ?: ((formattedPrompt.length * 0.75) + 64).toInt()

diff --git a/...-client-local/src/main/kotlin/com/llama/llamastack/client/local/util/PromptFormatLocal.kt b/...-client-local/src/main/kotlin/com/llama/llamastack/client/local/util/PromptFormatLocal.kt
@@ -40,13 +40,13 @@ object PromptFormatLocal {
         }
     }
 
-    fun getStopToken(modelName: String?): String {
+    fun getStopTokens(modelName: String?): List<String> {
         return when (modelName) {
             "LLAMA_3",
             "LLAMA_3_1",
             "LLAMA_3_2",
-            "LLAMA_GUARD_3" -> "<|eot_id|>"
-            else -> ""
+            "LLAMA_GUARD_3" -> listOf("<|eot_id|>", "<|eom_id|>")
+            else -> listOf("")
         }
     }