fix: more cases of unknown characters in generation streaming (#295)

withcatai · Aug 9, 2024 · ecaef63 · ecaef63
1 parent 097b3ec
commit ecaef63
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/src/evaluator/LlamaChat/LlamaChat.ts b/src/evaluator/LlamaChat/LlamaChat.ts
@@ -2217,6 +2217,9 @@ class GenerateResponseState<const Functions extends ChatModelFunctions | undefin
     public waitOnPartialCharactersOrWhiteSpaceTokens() {
         if (this.currentText.endsWith(UNKNOWN_UNICODE_CHAR) || (
             (this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix) && this.currentText?.trim() === ""
+        ) || (
+            this.currentText === "" && this.locksToReleaseOnValidGeneration.length > 0 &&
+            !this.llamaChat.model.isSpecialToken(this.currentToken)
         )) {
             if (this.currentQueuedTokenRelease != null)
                 this.locksToReleaseOnValidGeneration.push(this.currentQueuedTokenRelease.createTextIndexLock(0));

diff --git a/src/evaluator/LlamaCompletion.ts b/src/evaluator/LlamaCompletion.ts
@@ -664,6 +664,8 @@ export class LlamaCompletion {
 
                 if (text.endsWith(UNKNOWN_UNICODE_CHAR) || (
                     (grammar?.trimWhitespaceSuffix || trimWhitespaceSuffix) && text.trim() === ""
+                ) || (
+                    text === "" && locksToReleaseOnValidGeneration.length > 0 && !model.isSpecialToken(token)
                 )) {
                     locksToReleaseOnValidGeneration.push(queuedTokenRelease.createTextIndexLock(0));
                 } else {