QwenPipeline.cs

// Copyright (c) TensorStack. All rights reserved.
// Licensed under the Apache 2.0 License.

using System;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using TensorStack.Common;
using TensorStack.Common.Pipeline;
using TensorStack.Common.Tensor;
using TensorStack.TextGeneration.Cache;
using TensorStack.TextGeneration.Common;
using TensorStack.TextGeneration.Processing;
using TensorStack.TextGeneration.Tokenizers;

namespace TensorStack.TextGeneration.Pipelines.Qwen
{
    public class QwenPipeline : DecoderPipeline<GenerateOptions>,
        IPipeline<GenerateResult, GenerateOptions, GenerateProgress>,
        IPipeline<GenerateResult[], SearchOptions, GenerateProgress>
    {
        /// <summary>
        /// Initializes a new instance of the <see cref="QwenPipeline"/> class.
        /// </summary>
        /// <param name="tokenizerConfig">The tokenizer configuration.</param>
        /// <param name="decoderConfig">The decoder configuration.</param>
        public QwenPipeline(QwenConfig configuration)
            : base(configuration.Tokenizer, configuration.DecoderConfig)
        {
            Configuration = configuration;
        }

        public QwenConfig Configuration { get; }


        /// <summary>
        /// Runs the GreedySearch inference
        /// </summary>
        /// <param name="options">The options.</param>
        /// <param name="cancellationToken">The cancellation token.</param>
        /// <returns></returns>
        public virtual async Task<GenerateResult> RunAsync(GenerateOptions options, IProgress<GenerateProgress> progressCallback = null, CancellationToken cancellationToken = default)
        {
            await TokenizePromptAsync(options);
            var sequence = await GreedySearchAsync(options, progressCallback, cancellationToken);
            using (sequence)
            {
                return new GenerateResult
                {
                    Score = sequence.Score,
                    Result = Tokenizer.Decode(sequence.Tokens),
                    Tokens = sequence.Tokens,
                    LastHiddenState = sequence.LastHiddenState
                };
            }
        }


        /// <summary>
        /// Runs the BeamSearch inference
        /// </summary>
        /// <param name="options">The options.</param>
        /// <param name="progressCallback">The progress callback.</param>
        /// <param name="cancellationToken">The cancellation token that can be used by other objects or threads to receive notice of cancellation.</param>
        public async Task<GenerateResult[]> RunAsync(SearchOptions options, IProgress<GenerateProgress> progressCallback = null, CancellationToken cancellationToken = default)
        {
            await TokenizePromptAsync(options);

            var sequences = await BeamSearchAsync(options, progressCallback, cancellationToken);
            var results = new GenerateResult[sequences.Length];
            for (int beam = 0; beam < sequences.Length; beam++)
            {
                var sequence = sequences[beam];
                using (sequence)
                {
                    results[beam] = new GenerateResult
                    {
                        Beam = beam,
                        Score = sequence.Score,
                        PenaltyScore = sequence.PenaltyScore,
                        Result = Tokenizer.Decode(sequence.Tokens),
                        Tokens = sequence.Tokens,
                        LastHiddenState = sequence.LastHiddenState
                    };
                }
            }
            return results;
        }


        /// <summary>
        /// Tokenize the prompt
        /// </summary>
        /// <param name="options">The options.</param>
        /// <returns>A Task representing the asynchronous operation.</returns>
        protected override async Task TokenizePromptAsync(GenerateOptions options)
        {
            var tokenizerResult = await Tokenizer.EncodeAsync(options.Prompt);
            var inputIds = tokenizerResult.InputIds.Span.Pad(Tokenizer.EOS, options.MinLength);
            var mask = tokenizerResult.Mask.Span.Pad(0, options.MinLength);
            TokenizerOutput = new TokenizerResult(inputIds, mask);
        }


        /// <summary>
        /// Gets the token processors.
        /// </summary>
        /// <param name="options">The options.</param>
        /// <returns>ITokenProcessor[].</returns>
        protected override ITokenProcessor[] GetTokenProcessors(GenerateOptions options)
        {
            return
            [
                new EOSTokenProcessor(options.MinLength, Tokenizer.EOS),
                new MaxLengthTokenProcessor(options.MaxLength)
            ];
        }


        /// <summary>
        /// Initialize the Decoder cache
        /// </summary>
        /// <param name="options">The options.</param>
        /// <returns>A Task&lt;Sequence&gt; representing the asynchronous operation.</returns>
        protected override async Task<Sequence> InitializeAsync(GenerateOptions options)
        {
            var modelMetadata = await Decoder.LoadAsync();
            var kvCache = new KVCacheDecoder(modelMetadata, DecoderConfig.NumHeads, DecoderConfig.NumLayers, DecoderConfig.HiddenSize, DecoderConfig.NumKVHeads, options.MaxLength);
            var sequence = new Sequence(kvCache, Tokenizer.BOS);
            sequence.Initialize(0);

            var position = TokenizerOutput.Length;
            var inputIds = TokenizerOutput.InputIds;
            var positionIds = GetPositionIds(modelMetadata, 0, position);
            var attentionMask = new Tensor<long>([1, position], 1);
            RunDecoderInternal(modelMetadata, sequence, inputIds, positionIds, attentionMask, false);
            return sequence;
        }


        /// <summary>
        /// Run decoder model
        /// </summary>
        /// <param name="sequence">The sequence.</param>
        /// <returns>A Task&lt;Tensor`1&gt; representing the asynchronous operation.</returns>
        protected override async Task<Tensor<float>> RunDecoderAsync(Sequence sequence)
        {
            var modelMetadata = await Decoder.LoadAsync();
            var position = TokenizerOutput.Length + sequence.Tokens.Count;
            var inputIds = new Tensor<long>([1, 1], sequence.Tokens[^1]);
            var positionIds = GetPositionIds(modelMetadata, position);
            var attentionMask = new Tensor<long>([1, position], 1);
            return RunDecoderInternal(modelMetadata, sequence, inputIds, positionIds, attentionMask, true);
        }


        /// <summary>
        /// Runs the decoder
        /// </summary>
        /// <param name="modelMetadata">The model metadata.</param>
        /// <param name="sequence">The sequence.</param>
        /// <param name="inputIds">The input ids.</param>
        /// <param name="positionIds">The position ids.</param>
        /// <param name="attentionMask">The attention mask.</param>
        /// <param name="useBranchCache">if set to <c>true</c> [use branch cache].</param>
        private Tensor<float> RunDecoderInternal(ModelMetadata modelMetadata, Sequence sequence, Tensor<long> inputIds, Tensor<long> positionIds, Tensor<long> attentionMask, bool useBranchCache)
        {
            using (var parameters = new ModelParameters(modelMetadata))
            {
                // Inputs
                parameters.AddInput(inputIds);
                parameters.AddInput(attentionMask);
                if (positionIds != null)
                    parameters.AddInput(positionIds);

                foreach (var pastKeyValue in sequence.Cache)
                    parameters.AddInput(pastKeyValue, false);

                // Outputs
                foreach (var output in modelMetadata.Outputs)
                    parameters.AddOutput();

                // Result
                var modelResult = Decoder.RunInference(parameters);
                using (var logitsResult = modelResult[0])
                {
                    var dimension = logitsResult.GetDimensions();
                    var logits = logitsResult.ToTensor(dimension[1..]);
                    var presentKeyValues = modelResult.ToArray()[1..];
                    sequence.UpdateCache(presentKeyValues, useBranchCache);
                    return logits;
                }
            }
        }


        /// <summary>
        /// Creates the QwenPipeline
        /// </summary>
        /// <param name="provider">The provider.</param>
        /// <param name="modelPath">The model path.</param>
        /// <param name="model">The decoder model.</param>
        /// <returns>QwenPipeline.</returns>
        public static QwenPipeline Create(ExecutionProvider provider, string modelPath, string model = "model.onnx")
        {
            // Qwen-2.5 - https://huggingface.co/onnx-community/Qwen2.5-0.5B
            var numHeads = 14;
            var numLayers = 24;
            var hiddenSize = 896;
            var numKVHeads = 2;
            var vocabSize = 151936;
            var config = new QwenConfig
            {
                Tokenizer = new BPETokenizer(new TokenizerConfig
                {
                    BOS = 151643,
                    EOS = 151643,
                    Path = modelPath
                }),
                DecoderConfig = new DecoderConfig
                {
                    Path = Path.Combine(modelPath, model),
                    VocabSize = vocabSize,
                    NumHeads = numHeads,
                    NumLayers = numLayers,
                    HiddenSize = hiddenSize,
                    NumKVHeads = numKVHeads
                }
            };

            config.DecoderConfig.SetProvider(provider);
            return new QwenPipeline(config);
        }

    }
}