Critical Bug: Context shifting doesn't work!

[aropb]

The tokens list cannot be changed, it is the input list of tokens!

Error line:

LLamaSharp/LLama/LLamaStatelessExecutor.cs

Line 137 in 9fe066d

tokens.Clear();

Llama.cpp source:

https://github.com/ggml-org/llama.cpp/blob/60325fa56f61c228464c9f065db3aa6a61f2156e/examples/main/main.cpp#L334
https://github.com/ggml-org/llama.cpp/blob/60325fa56f61c228464c9f065db3aa6a61f2156e/examples/main/main.cpp#L542

I'm trying to fix and rewrite it correctly, but so far it's not working. I don't understand all the calculations. Let's figure it out together?

[aropb]

Here is a more accurate code, but there is one problem that I cannot solve yet. After the context shift, there is an endless looping and output of one character.

    public async IAsyncEnumerable<string> InferAsync(string text, IInferenceParams inferenceParams, [EnumeratorCancellation] CancellationToken token)
    {
        // Clear Context memory
        _logger.LogDebug("LLamaStatelessExecutor InferAsync start");
        Context.NativeHandle.MemoryClear();

        // Reset the sampling pipeline (if there is one)
        inferenceParams?.SamplingPipeline?.Reset();

        // Sanity check inference params
        inferenceParams ??= new InferenceParams();
        if (inferenceParams.TokensKeep > Context.ContextSize)
        {
            throw new ArgumentOutOfRangeException(nameof(inferenceParams), $"TokensKeep ({inferenceParams.TokensKeep}) cannot be larger than ContextSize ({Context.ContextSize})");
        }

        // Create decoders for the token stream
        StreamingTokenDecoder decoder = new(Context)
        {
            DecodeSpecialTokens = inferenceParams.DecodeSpecialTokens
        };

        AntipromptProcessor antiprocessor = new(inferenceParams.AntiPrompts);

        if (ApplyTemplate)
        {
            LLamaTemplate template = new(_weights.NativeHandle) { AddAssistant = true };
            if (SystemMessage != null) { template.Add("system", SystemMessage); }

            template.Add("user", text);
            text = PromptTemplateTransformer.ToModelPrompt(template);
        }

        // Tokenize the prompt
        var tokens = Context.Tokenize(text: text, addBos: true, special: true).ToList();

        _logger.LogDebug("LLamaStatelessExecutor InferAsync decode prompt");

        // Evaluate the prompt, in chunks smaller than the max batch size
        var (result, _, n_past) = await Context.DecodeAsync(tokens, LLamaSeqId.Zero, _batch, 0);
        if (result != DecodeResult.Ok) { throw new LLamaDecodeError(result); }

        _logger.LogDebug("LLamaStatelessExecutor InferAsync inference loop start");

        // Begin loop, evaluating one token at a time
        int maxTokens = inferenceParams.MaxTokens < 0 ? int.MaxValue : inferenceParams.MaxTokens;
        for (int i = 0; i < maxTokens && !token.IsCancellationRequested; i++)
        {
            // Sample with the pipeline
            LLamaToken id = inferenceParams.SamplingPipeline.Sample(Context.NativeHandle, _batch.TokenCount - 1);

            // Check if this token should end generation
            if (id.IsEndOfGeneration(_weights.Vocab)) { break; }

            // Decode this token into text
            decoder.Add(id);
            string decoded = decoder.Read();
            yield return decoded;

            // Check if any of the antiprompts have been generated
            if (antiprocessor.Add(decoded)) { break; }

            // when run out of context
            // https://github.com/ggml-org/llama.cpp/blob/60325fa56f61c228464c9f065db3aa6a61f2156e/examples/main/main.cpp#L334
            // https://github.com/ggml-org/llama.cpp/blob/60325fa56f61c228464c9f065db3aa6a61f2156e/examples/main/main.cpp#L542

            if (n_past + 1 >= Context.ContextSize)
            {
                // number of tokens to keep when resetting context
                bool canAddBos = Context.Vocab.ShouldAddBOS;
                int tokensKeep = inferenceParams.TokensKeep;

                if (tokensKeep < 0 || tokensKeep > tokens.Count)
                {
                    tokensKeep = tokens.Count;
                }
                else
                {
                    tokensKeep += Convert.ToInt32(canAddBos);
                }

                int n_left = n_past - tokensKeep;
                int n_discard = n_left / 2;

                Context.NativeHandle.MemorySequenceRemove(LLamaSeqId.Zero, tokensKeep, tokensKeep + n_discard);
                Context.NativeHandle.MemorySequenceAdd(LLamaSeqId.Zero, tokensKeep + n_discard, n_past, -n_discard);

                n_past -= n_discard;
            }

            // Evaluate with this new token
            _batch.Clear();
            _batch.Add(id, n_past++, LLamaSeqId.Zero, true);

            var returnCode = await Context.DecodeAsync(_batch, token);
            if (returnCode != 0) { throw new LLamaDecodeError(returnCode); }
        }

        _logger.LogDebug("LLamaStatelessExecutor InferAsync success");
    }

[aropb]

@zsogitbe @martindevans

As a result, after many hours of research, this option works well. But apparently it is not the most optimal, since new tokens are not shifted in the context, but are being recreated. I would appreciate it if someone could find the right solution.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Runtime.CompilerServices;
using Microsoft.Extensions.Logging;
using LLama;
using LLama.Abstractions;
using LLama.Exceptions;
using LLama.Native;
using LLama.Common;
using LLama.Transformers;

namespace LLamaSharp
{
///

/// This executor infer the input as one-time job. Previous inputs won't impact on the
/// response to current input.
///

public sealed class LLamaStatelessExecutor : ILLamaExecutor, IDisposable
{
private readonly LLamaWeights _weights;
private readonly ILogger _logger;

    public bool IsMultiModal => false;

    public LLavaWeights ClipModel => default;

    public List<byte[]> Images { get; } = [];

    /// <summary>
    /// The context used by the executor when running the inference.
    /// </summary>
    public LLamaContext Context { get; }

    /// <summary>
    /// If true, applies the default template to the prompt as defined in the rules for <a href="https://github.com/ggerganov/llama.cpp/wiki/Templates-supported-by-llama_chat_apply_template">llama_chat_apply_template</a> template.  
    /// </summary>
    public bool ApplyTemplate { get; init; }
    
    /// <summary>
    /// The system message to use with the prompt. Only used when <see cref="ApplyTemplate" /> is true.
    /// </summary>
    public string SystemMessage { get; init; }

    /// <summary>
    /// Create a new stateless executor which will use the given model
    /// </summary>
    /// <param name="weights"></param>
    /// <param name="params"></param>
    /// <param name="logger"></param>
    public LLamaStatelessExecutor(LLamaWeights weights, IContextParams @params, ILogger logger)
    {
        Context = weights.CreateContext(@params, logger);
        _weights = weights;
        _logger = logger;
    }
    
    public void Dispose()
    {
        Context.Dispose();
    }

    public async IAsyncEnumerable<string> InferAsync(string text, IInferenceParams inferenceParams, [EnumeratorCancellation] CancellationToken token)
    {
        _logger.LogDebug("LLamaStatelessExecutor InferAsync start");

        // Clear Context memory
        Context.NativeHandle.MemoryClear();

        // Reset the sampling pipeline (if there is one)
        inferenceParams?.SamplingPipeline?.Reset();

        // Sanity check inference params
        inferenceParams ??= new InferenceParams();
        if (inferenceParams.TokensKeep > Context.ContextSize)
        {
            throw new ArgumentOutOfRangeException(nameof(inferenceParams), $"TokensKeep ({inferenceParams.TokensKeep}) cannot be larger than ContextSize ({Context.ContextSize})");
        }

        // Create decoders for the token stream
        StreamingTokenDecoder decoder = new(Context)
        {
            DecodeSpecialTokens = inferenceParams.DecodeSpecialTokens
        };

        AntipromptProcessor antiprocessor = new(inferenceParams.AntiPrompts);

        if (ApplyTemplate)
        {
            LLamaTemplate template = new(_weights.NativeHandle) { AddAssistant = true };
            if (SystemMessage != null) { template.Add("system", SystemMessage); }

            template.Add("user", text);
            text = PromptTemplateTransformer.ToModelPrompt(template);
        }

        // Tokenize the prompt
        List<LLamaToken> inputTokens = Context.Tokenize(text: text, addBos: true, special: true).ToList();

        // Generate tokens
        List<LLamaToken> tokens = new();

        // Evaluate the prompt, in chunks smaller than the max batch size
        LLamaBatch batch = new();
        var (result, _, n_past) = await Context.DecodeAsync(inputTokens, LLamaSeqId.Zero, batch, 0);
        if (result != DecodeResult.Ok) { throw new LLamaDecodeError(result); }

        // Begin loop, evaluating one token at a time
        int maxTokens = inferenceParams.MaxTokens < 0 ? int.MaxValue : inferenceParams.MaxTokens;
        for (int i = 0; i < maxTokens && !token.IsCancellationRequested; i++)
        {
            // Sample with the pipeline
            LLamaToken llamaToken = inferenceParams.SamplingPipeline.Sample(Context.NativeHandle, batch.TokenCount - 1);
            tokens.Add(llamaToken);

            // Check if this token should end generation
            if (llamaToken.IsEndOfGeneration(_weights.Vocab)) { break; }

            // Decode this token into text
            decoder.Add(llamaToken);
            string decoded = decoder.Read();

            // Check if any of the antiprompts have been generated
            if (antiprocessor.Add(decoded)) { break; }

            // Return decoded new token
            yield return decoded;

            // When run out of context
            if (inputTokens.Count + tokens.Count >= Context.ContextSize)
            {
                // Number of tokens to keep when resetting context
                bool canAddBos = Context.Vocab.ShouldAddBOS;
                int tokensKeep = inferenceParams.TokensKeep;

                if (tokensKeep < 0 || tokensKeep > inputTokens.Count)
                {
                    tokensKeep = inputTokens.Count;
                }
                else
                {
                    tokensKeep += Convert.ToInt32(canAddBos);
                }

                int n_left = n_past - tokensKeep;
                int n_discard = n_left / 2;

                tokens.RemoveRange(0, n_discard);
                Context.NativeHandle.MemorySequenceRemove(LLamaSeqId.Zero, tokensKeep, n_past);

                batch.Clear();
                (result, _, n_past) = await Context.DecodeAsync(tokens, LLamaSeqId.Zero, batch, tokensKeep);
                if (result != DecodeResult.Ok) { throw new LLamaDecodeError(result); }
            }
            else
            {
                // Evaluate with this new token
                batch.Clear();
                batch.Add(llamaToken, n_past++, LLamaSeqId.Zero, true);

                var returnCode = await Context.DecodeAsync(batch, token);
                if (returnCode != DecodeResult.Ok) { throw new LLamaDecodeError(returnCode); }
            }
        }

        _logger.LogDebug("LLamaStatelessExecutor InferAsync success");
    }
}

}