Merge pull request #248 from AsakusaRinne/kernel_memory_dev

feat(kernel-memory): avoid loading model twice.

LLama.KernelMemory/BuilderExtensions.cs

@@ -4,6 +4,9 @@ using System.Collections.Generic;
 using System.Linq;
 using System.Text;
 using System.Threading.Tasks;
+using LLama;
+using LLama.Common;
+using Microsoft.KernelMemory.AI;
 
 namespace LLamaSharp.KernelMemory
 {
@@ -24,6 +27,18 @@ namespace LLamaSharp.KernelMemory
             return builder;
         }
 
+        /// <summary>
+        /// Adds LLamaSharpTextEmbeddingGeneration to the KernelMemoryBuilder.
+        /// </summary>
+        /// <param name="builder">The KernelMemoryBuilder instance.</param>
+        /// <param name="textEmbeddingGeneration">The LLamaSharpTextEmbeddingGeneration instance.</param>
+        /// <returns>The KernelMemoryBuilder instance with LLamaSharpTextEmbeddingGeneration added.</returns>
+        public static KernelMemoryBuilder WithLLamaSharpTextEmbeddingGeneration(this KernelMemoryBuilder builder, LLamaSharpTextEmbeddingGeneration textEmbeddingGeneration)
+        {
+            builder.WithCustomEmbeddingGeneration(textEmbeddingGeneration);
+            return builder;
+        }
+
         /// <summary>
         /// Adds LLamaSharpTextGeneration to the KernelMemoryBuilder.
         /// </summary>
@@ -36,6 +51,18 @@ namespace LLamaSharp.KernelMemory
             return builder;
         }
 
+        /// <summary>
+        /// Adds LLamaSharpTextGeneration to the KernelMemoryBuilder.
+        /// </summary>
+        /// <param name="builder">The KernelMemoryBuilder instance.</param>
+        /// <param name="textGeneration">The LlamaSharpTextGeneration instance.</param>
+        /// <returns>The KernelMemoryBuilder instance with LLamaSharpTextGeneration added.</returns>
+        public static KernelMemoryBuilder WithLLamaSharpTextGeneration(this KernelMemoryBuilder builder, LlamaSharpTextGeneration textGeneration)
+        {
+            builder.WithCustomTextGeneration(textGeneration);
+            return builder;
+        }
+
         /// <summary>
         /// Adds LLamaSharpTextEmbeddingGeneration and LLamaSharpTextGeneration to the KernelMemoryBuilder.
         /// </summary>
@@ -44,8 +71,18 @@ namespace LLamaSharp.KernelMemory
         /// <returns>The KernelMemoryBuilder instance with LLamaSharpTextEmbeddingGeneration and LLamaSharpTextGeneration added.</returns>
         public static KernelMemoryBuilder WithLLamaSharpDefaults(this KernelMemoryBuilder builder, LLamaSharpConfig config)
         {
-            builder.WithLLamaSharpTextEmbeddingGeneration(config);
-            builder.WithLLamaSharpTextGeneration(config);
+            var parameters = new ModelParams(config.ModelPath)
+            {
+                ContextSize = config?.ContextSize ?? 2048,
+                Seed = config?.Seed ?? 0,
+                GpuLayerCount = config?.GpuLayerCount ?? 20
+            };
+            var weights = LLamaWeights.LoadFromFile(parameters);
+            var context = weights.CreateContext(parameters);
+            var executor = new StatelessExecutor(weights, parameters);
+            var embedder = new LLamaEmbedder(weights, parameters);
+            builder.WithLLamaSharpTextEmbeddingGeneration(new LLamaSharpTextEmbeddingGeneration(embedder));
+            builder.WithLLamaSharpTextGeneration(new LlamaSharpTextGeneration(weights, context, executor));
             return builder;
         }
     }

LLama.KernelMemory/LLamaSharpTextEmbeddingGeneration.cs

@@ -1,4 +1,5 @@
 using LLama;
+using LLama.Abstractions;
 using LLama.Common;
 using Microsoft.SemanticKernel.AI.Embeddings;
 using System;
@@ -14,9 +15,11 @@ namespace LLamaSharp.KernelMemory
     /// </summary>
     public class LLamaSharpTextEmbeddingGeneration : ITextEmbeddingGeneration, IDisposable
     {
-        private readonly LLamaSharpConfig _config;
+        private readonly LLamaSharpConfig? _config;
+        private readonly LLamaWeights? _weights;
         private readonly LLamaEmbedder _embedder;
-        private readonly LLamaWeights _weights;
+        private bool _ownsEmbedder = false;
+        private bool _ownsWeights = false;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="LLamaSharpTextEmbeddingGeneration"/> class.
@@ -28,13 +31,46 @@ namespace LLamaSharp.KernelMemory
             var @params = new ModelParams(_config.ModelPath);
             _weights = LLamaWeights.LoadFromFile(@params);
             _embedder = new LLamaEmbedder(_weights, @params);
+            _ownsWeights = true;
+            _ownsEmbedder = true;
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="LLamaSharpTextEmbeddingGeneration"/> class from reused weights.
+        /// </summary>
+        /// <param name="config">The configuration for LLamaSharp.</param>
+        /// <param name="weights">A LLamaWeights object.</param>
+        public LLamaSharpTextEmbeddingGeneration(LLamaSharpConfig config, LLamaWeights weights)
+        {
+            this._config = config;
+            var @params = new ModelParams(_config.ModelPath);
+            _weights = weights;
+            _embedder = new LLamaEmbedder(_weights, @params);
+            _ownsEmbedder = true;
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="LLamaSharpTextEmbeddingGeneration"/> class from reused embedder.
+        /// </summary>
+        /// <param name="embedder">A LLamaEmbedder object.</param>
+        public LLamaSharpTextEmbeddingGeneration(LLamaEmbedder embedder)
+        {
+            this._config = null;
+            this._weights = null;
+            _embedder = embedder;
         }
 
         /// <inheritdoc/>
         public void Dispose()
         {
-            _embedder.Dispose();
-            _weights.Dispose();
+            if (_ownsWeights)
+            {
+                _weights?.Dispose();
+            }
+            if(_ownsEmbedder)
+            {
+                _embedder.Dispose();
+            }
         }
 
         /// <inheritdoc/>

LLama.KernelMemory/LlamaSharpConfig.cs

@@ -7,7 +7,7 @@ using System.Threading.Tasks;
 namespace LLamaSharp.KernelMemory
 {
     /// <summary>
-    /// Represents the configuration for LLamaSharp.
+    /// Represents the configuration for LLamaSharp. Available properties are `ModelPath`, `ContextSize`, `Seed`, `GpuLayerCount`.
     /// </summary>
     public class LLamaSharpConfig
     {

LLama.KernelMemory/LlamaSharpTextGeneration.cs

@@ -1,4 +1,5 @@
 using LLama;
+using LLama.Abstractions;
 using LLama.Common;
 using Microsoft.KernelMemory.AI;
 using System;
@@ -14,10 +15,12 @@ namespace LLamaSharp.KernelMemory
     /// </summary>
     public class LlamaSharpTextGeneration : ITextGeneration, IDisposable
     {
-        private readonly LLamaSharpConfig _config;
+        private readonly LLamaSharpConfig? _config;
         private readonly LLamaWeights _weights;
         private readonly StatelessExecutor _executor;
         private readonly LLamaContext _context;
+        private bool _ownsContext = false;
+        private bool _ownsWeights = false;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="LlamaSharpTextGeneration"/> class.
@@ -35,13 +38,35 @@ namespace LLamaSharp.KernelMemory
             _weights = LLamaWeights.LoadFromFile(parameters);
             _context = _weights.CreateContext(parameters);
             _executor = new StatelessExecutor(_weights, parameters);
+            _ownsWeights = _ownsContext = true;
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="LlamaSharpTextGeneration"/> class from reused weights, context and executor.
+        /// If executor is not specified, then a StatelessExecutor will be created with `context.Params`. So far only `StatelessExecutor` is expected.
+        /// </summary>
+        /// <param name="weights">A LLamaWeights object.</param>
+        /// <param name="context">A LLamaContext object.</param>
+        /// <param name="executor">An executor. Currently only StatelessExecutor is expected.</param>
+        public LlamaSharpTextGeneration(LLamaWeights weights, LLamaContext context, StatelessExecutor? executor = null)
+        {
+            _config = null;
+            _weights = weights;
+            _context = context;
+            _executor = executor ?? new StatelessExecutor(_weights, _context.Params);
         }
 
         /// <inheritdoc/>
         public void Dispose()
         {
-            _context.Dispose();
-            _weights.Dispose();
+            if (_ownsWeights)
+            {
+                _weights?.Dispose();
+            }
+            if (_ownsContext)
+            {
+                _context.Dispose();
+            }
         }
 
         /// <inheritdoc/>