feat: improve context shift strategy

giladgd · giladgd · commit 1676d4108553 · 2024-10-27T01:49:40.000+02:00
diff --git a/src/evaluator/LlamaChat/LlamaChat.ts b/src/evaluator/LlamaChat/LlamaChat.ts
@@ -2166,21 +2166,24 @@ class GenerateResponseState<const Functions extends ChatModelFunctions | undefin
     }
 
     public async alignCurrentSequenceStateWithCurrentTokens() {
-        let {firstDifferentIndex} = this.llamaChat.sequence.compareContextTokens(this.tokens);
-
-        // we need to decode at least one token to generate a response
-        if (firstDifferentIndex === this.tokens.length && firstDifferentIndex > 0)
-            firstDifferentIndex -= 1;
-
-        this.tokens.splice(0, firstDifferentIndex);
-
-        if (firstDifferentIndex < this.llamaChat.sequence.nextTokenIndex) {
+        if (this.tokens.length === 1 && this.llamaChat.sequence.nextTokenIndex !== 0) {
             await this.llamaChat.sequence.eraseContextTokenRanges([{
-                start: firstDifferentIndex,
+                start: 0,
                 end: this.llamaChat.sequence.nextTokenIndex
             }]);
-            this.ensureNotAborted();
+            return;
         }
+
+        const lastToken = this.tokens[this.tokens.length - 1]!;
+
+        // we need to decode at least one token to generate a response
+        this.tokens.pop();
+        await this.llamaChat.sequence.adaptStateToTokens(this.tokens, false);
+        this.tokens.push(lastToken);
+        this.ensureNotAborted();
+
+        const firstDifferentIndex = this.llamaChat.sequence.nextTokenIndex;
+        this.tokens.splice(0, firstDifferentIndex);
     }
 
     public async evaluateWithoutGeneratingNewTokens() {
diff --git a/src/evaluator/LlamaChat/utils/contextShiftStrategies/eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy.ts b/src/evaluator/LlamaChat/utils/contextShiftStrategies/eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy.ts
@@ -30,6 +30,8 @@ export async function eraseFirstResponseAndKeepFirstSystemChatContextShiftStrate
         initialCharactersRemovalCount,
         tokenizer,
         chatWrapper,
+        failedCompressionErrorMessage: "Failed to compress chat history for context shift due to a too long prompt or system message that cannot be compressed without affecting the generation quality. " +
+            "Consider increasing the context size or shortening the long prompt or system message.",
         compressChatHistory({chatHistory, charactersToRemove, estimatedCharactersPerToken}) {
             const res = chatHistory.map(item => structuredClone(item));
             let charactersLeftToRemove = charactersToRemove;
@@ -66,6 +68,8 @@ export async function eraseFirstResponseAndKeepFirstSystemChatContextShiftStrate
             }
 
             function removeHistoryThatLedToModelResponseAtIndex(index: number) {
+                let removedItems = 0;
+
                 for (let i = index - 1; i >= 0; i--) {
                     const historyItem = res[i];
 
@@ -79,13 +83,19 @@ export async function eraseFirstResponseAndKeepFirstSystemChatContextShiftStrate
                         break; // keep the first system message
 
                     if (historyItem.type === "user" || historyItem.type === "system") {
-                        const newText = truncateLlamaTextAndRoundToWords(LlamaText.fromJSON(historyItem.text), charactersLeftToRemove);
+                        const newText = truncateLlamaTextAndRoundToWords(
+                            LlamaText.fromJSON(historyItem.text),
+                            charactersLeftToRemove,
+                            undefined,
+                            false
+                        );
                         const newTextString = newText.toString();
                         const historyItemString = LlamaText.fromJSON(historyItem.text).toString();
 
                         if (newText.values.length === 0) {
                             res.splice(i, 1);
                             i++;
+                            removedItems++;
                             charactersLeftToRemove -= historyItemString.length;
                         } else if (newTextString.length < historyItemString.length) {
                             charactersLeftToRemove -= historyItemString.length - newTextString.length;
@@ -98,6 +108,66 @@ export async function eraseFirstResponseAndKeepFirstSystemChatContextShiftStrate
                         void (historyItem satisfies never);
                     }
                 }
+
+                return removedItems;
+            }
+
+            function compressHistoryThatLedToModelResponseAtIndex(index: number, keepTokensCount: number = 0) {
+                let removedItems = 0;
+                let promptStartIndex: number | undefined = undefined;
+
+                for (let i = index - 1; i >= 0; i--) {
+                    const historyItem = res[i];
+
+                    if (historyItem == null)
+                        continue;
+
+                    if (historyItem.type === "model") {
+                        promptStartIndex = i + 1;
+                        break;
+                    }
+
+                    if (i === 0 && historyItem.type === "system") {
+                        promptStartIndex = i + 1;
+                        break; // keep the first system message
+                    }
+                }
+
+                if (promptStartIndex == null || promptStartIndex >= index)
+                    return 0;
+
+                for (let i = promptStartIndex; i < index && charactersLeftToRemove > 0; i++) {
+                    const historyItem = res[i];
+
+                    if (historyItem == null || historyItem.type !== "user")
+                        continue;
+
+                    let removeChars = Math.min(charactersLeftToRemove, historyItem.text.length);
+                    if (keepTokensCount > 0) {
+                        removeChars -= Math.floor(keepTokensCount * estimatedCharactersPerToken);
+                        if (removeChars < 0)
+                            removeChars = 0;
+
+                        keepTokensCount -= Math.min(
+                            keepTokensCount,
+                            Math.max(0, historyItem.text.length - removeChars) / estimatedCharactersPerToken
+                        );
+                    }
+
+                    const newText = truncateTextAndRoundToWords(historyItem.text, removeChars, undefined, false);
+                    if (newText.length === 0) {
+                        res.splice(i, 1);
+                        i--;
+                        index--;
+                        removedItems++;
+                        charactersLeftToRemove -= historyItem.text.length;
+                    } else {
+                        charactersLeftToRemove -= historyItem.text.length - newText.length;
+                        historyItem.text = newText;
+                    }
+                }
+
+                return removedItems;
             }
 
             function compressFirstModelResponse() {
@@ -116,7 +186,7 @@ export async function eraseFirstResponseAndKeepFirstSystemChatContextShiftStrate
                             continue;
 
                         if (typeof item === "string") {
-                            const newText = truncateTextAndRoundToWords(item, charactersLeftToRemove);
+                            const newText = truncateTextAndRoundToWords(item, charactersLeftToRemove, undefined, true);
 
                             if (newText === "") {
                                 historyItem.response.splice(t, 1);
@@ -139,14 +209,14 @@ export async function eraseFirstResponseAndKeepFirstSystemChatContextShiftStrate
                     if (historyItem.response.length === 0) {
                         // if the model response is removed from the history,
                         // the things that led to it are not important anymore
-                        removeHistoryThatLedToModelResponseAtIndex(i);
+                        i -= removeHistoryThatLedToModelResponseAtIndex(i);
                         res.splice(i, 1);
                         i--;
                     }
                 }
             }
 
-            function compressLastModelResponse(minCharactersToKeep: number = 20) {
+            function compressLastModelResponse(minCharactersToKeep: number = 60) {
                 const lastHistoryItem = res[res.length - 1];
 
                 if (lastHistoryItem == null || lastHistoryItem.type !== "model")
@@ -157,14 +227,27 @@ export async function eraseFirstResponseAndKeepFirstSystemChatContextShiftStrate
                 if (lastResponseItem == null || typeof lastResponseItem !== "string")
                     return;
 
-                const nextTextLength = lastResponseItem.length - charactersLeftToRemove;
-                const charactersToRemoveFromText = charactersLeftToRemove + Math.max(0, nextTextLength - minCharactersToKeep);
-                const newText = truncateTextAndRoundToWords(lastResponseItem, charactersToRemoveFromText);
+                compressHistoryThatLedToModelResponseAtIndex(res.length - 1, maxTokensCount / 4);
+
+                if (charactersLeftToRemove <= 0)
+                    return;
+
+                const nextTextLength = Math.max(
+                    Math.min(lastResponseItem.length, minCharactersToKeep),
+                    lastResponseItem.length - charactersLeftToRemove
+                );
+                const charactersToRemoveFromText = lastResponseItem.length - nextTextLength;
+                const newText = truncateTextAndRoundToWords(lastResponseItem, charactersToRemoveFromText, undefined, true);
 
                 if (newText.length < lastResponseItem.length) {
                     lastHistoryItem.response[lastHistoryItem.response.length - 1] = newText;
                     charactersLeftToRemove -= lastResponseItem.length - newText.length;
                 }
+
+                if (charactersLeftToRemove <= 0)
+                    return;
+
+                compressHistoryThatLedToModelResponseAtIndex(res.length - 1);
             }
 
             compressFunctionCalls();
diff --git a/src/evaluator/LlamaCompletion.ts b/src/evaluator/LlamaCompletion.ts
@@ -660,20 +660,22 @@ export class LlamaCompletion {
 
             let shouldContextShift = false;
 
-            let {firstDifferentIndex} = sequence.compareContextTokens(inputTokens);
-
-            // we need to decode at least one token to generate a response
-            if (firstDifferentIndex === inputTokens.length && firstDifferentIndex > 0)
-                firstDifferentIndex -= 1;
-
-            inputTokens.splice(0, firstDifferentIndex);
-
-            if (firstDifferentIndex < sequence.nextTokenIndex) {
+            if (inputTokens.length === 1 && sequence.nextTokenIndex !== 0)
                 await sequence.eraseContextTokenRanges([{
-                    start: firstDifferentIndex,
+                    start: 0,
                     end: sequence.nextTokenIndex
                 }]);
+            else {
+                const lastToken = inputTokens[inputTokens.length - 1]!;
+
+                // we need to decode at least one token to generate a response
+                inputTokens.pop();
+                await sequence.adaptStateToTokens(inputTokens, false);
+                inputTokens.push(lastToken);
                 ensureNotAborted();
+
+                const firstDifferentIndex = sequence.nextTokenIndex;
+                inputTokens.splice(0, firstDifferentIndex);
             }
 
             const evaluationIterator = sequence.evaluate(inputTokens, removeNullFields({
diff --git a/src/evaluator/LlamaContext/LlamaContext.ts b/src/evaluator/LlamaContext/LlamaContext.ts
@@ -909,6 +909,61 @@ export class LlamaContextSequence {
         };
     }
 
+    /**
+     * Erase parts of the context state to align it with the given tokens.
+     *
+     * If the given tokens do not align with the current context state, the context state will be erased to align with the given tokens.
+     *
+     * To find the first different token index between the context state and the given tokens, access the `nextTokenIndex` property.
+     *
+     * If `allowShift` is `true` (the default), shifting tokens may happen to align the context state with the given tokens,
+     * which incurs token evaluation of the shifted tokens.
+     */
+    public async adaptStateToTokens(tokens: Token[], allowShift: boolean = true) {
+        if (this.model.fileInsights.isRecurrent || !allowShift) {
+            const {firstDifferentIndex} = this.compareContextTokens(tokens);
+            if (firstDifferentIndex < this._nextTokenIndex)
+                await this.eraseContextTokenRanges([{
+                    start: firstDifferentIndex,
+                    end: this._nextTokenIndex
+                }]);
+
+            return;
+        }
+
+        const eraseRanges: ContextTokensDeleteRange[] = [];
+
+        let tokensIndex = 0;
+        let differentTokenIndex: number | undefined = undefined;
+        for (let i = 0; i < this._contextTokens.length && tokensIndex < tokens.length; i++) {
+            if (compareTokens(this._contextTokens[i], tokens[tokensIndex])) {
+                if (differentTokenIndex != null) {
+                    eraseRanges.push({
+                        start: differentTokenIndex,
+                        end: i
+                    });
+
+                    differentTokenIndex = undefined;
+                }
+
+                tokensIndex++;
+                continue;
+            }
+
+            if (differentTokenIndex == null)
+                differentTokenIndex = i;
+        }
+
+        if (differentTokenIndex != null)
+            eraseRanges.push({
+                start: differentTokenIndex,
+                end: this._nextTokenIndex
+            });
+
+        if (eraseRanges.length > 0)
+            await this.eraseContextTokenRanges(eraseRanges);
+    }
+
     /**
      * Clear the history of the sequence.
      * If `prependBos` was enabled, the BOS token will be prepended to the sequence again.
@@ -975,15 +1030,23 @@ export class LlamaContextSequence {
                 if (deletionSuccessful)
                     deletionSuccessful &&= this._context._ctx.removeTokenCellsFromSequence(this._sequenceId, range.start, range.end);
 
-                if (deletionSuccessful && lastDeleteRangeEndPos != null && removedTokens > 0 && lastDeleteRangeEndPos !== range.start)
+                if (deletionSuccessful && lastDeleteRangeEndPos != null && removedTokens > 0 && lastDeleteRangeEndPos !== range.start) {
                     this._context._ctx.shiftSequenceTokenCells(this._sequenceId, lastDeleteRangeEndPos, range.start, -removedTokens);
+                    const shiftedTokens = range.start - lastDeleteRangeEndPos;
+                    this._tokenMeter.useTokens(shiftedTokens, "input");
+                }
 
                 removedTokens += range.end - range.start;
                 lastDeleteRangeEndPos = range.end;
             }
 
-            if (deletionSuccessful && lastDeleteRangeEndPos != null && removedTokens > 0 && lastDeleteRangeEndPos !== this._nextTokenIndex)
+            if (deletionSuccessful && lastDeleteRangeEndPos != null && removedTokens > 0 &&
+                lastDeleteRangeEndPos !== this._nextTokenIndex
+            ) {
                 this._context._ctx.shiftSequenceTokenCells(this._sequenceId, lastDeleteRangeEndPos, this._nextTokenIndex, -removedTokens);
+                const shiftedTokens = this._nextTokenIndex - lastDeleteRangeEndPos;
+                this._tokenMeter.useTokens(shiftedTokens, "input");
+            }
 
             this._nextTokenIndex -= removedTokens;
 
diff --git a/src/gguf/insights/GgufInsights.ts b/src/gguf/insights/GgufInsights.ts
@@ -104,6 +104,16 @@ export class GgufInsights {
         return true;
     }
 
+    public get isRecurrent() {
+        switch (this._ggufFileInfo.metadata?.general?.architecture) {
+            case GgufArchitectureType.mamba:
+            case GgufArchitectureType.rwkv6:
+                return true;
+        }
+
+        return false;
+    }
+
     public estimateModelResourceRequirements({gpuLayers}: {gpuLayers: number}): GgufInsightsResourceRequirements {
         const {cpu, gpu} = this._getTensorResourceSplit(gpuLayers);
 
diff --git a/src/utils/truncateTextAndRoundToWords.ts b/src/utils/truncateTextAndRoundToWords.ts
diff --git a/test/modelDependent/llama3/chatSession.test.ts b/test/modelDependent/llama3/chatSession.test.ts
