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

This commit is contained in:
Yaohui Liu 2023-11-05 17:16:50 +08:00
parent 3a4a91bc08
commit 46f01bbc94
No known key found for this signature in database
GPG Key ID: E86D01E1809BD23E
4 changed files with 108 additions and 10 deletions

View File

@ -4,6 +4,9 @@ using System.Collections.Generic;
using System.Linq; using System.Linq;
using System.Text; using System.Text;
using System.Threading.Tasks; using System.Threading.Tasks;
using LLama;
using LLama.Common;
using Microsoft.KernelMemory.AI;
namespace LLamaSharp.KernelMemory namespace LLamaSharp.KernelMemory
{ {
@ -24,6 +27,18 @@ namespace LLamaSharp.KernelMemory
return builder; 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> /// <summary>
/// Adds LLamaSharpTextGeneration to the KernelMemoryBuilder. /// Adds LLamaSharpTextGeneration to the KernelMemoryBuilder.
/// </summary> /// </summary>
@ -36,6 +51,18 @@ namespace LLamaSharp.KernelMemory
return builder; 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> /// <summary>
/// Adds LLamaSharpTextEmbeddingGeneration and LLamaSharpTextGeneration to the KernelMemoryBuilder. /// Adds LLamaSharpTextEmbeddingGeneration and LLamaSharpTextGeneration to the KernelMemoryBuilder.
/// </summary> /// </summary>
@ -44,8 +71,18 @@ namespace LLamaSharp.KernelMemory
/// <returns>The KernelMemoryBuilder instance with LLamaSharpTextEmbeddingGeneration and LLamaSharpTextGeneration added.</returns> /// <returns>The KernelMemoryBuilder instance with LLamaSharpTextEmbeddingGeneration and LLamaSharpTextGeneration added.</returns>
public static KernelMemoryBuilder WithLLamaSharpDefaults(this KernelMemoryBuilder builder, LLamaSharpConfig config) public static KernelMemoryBuilder WithLLamaSharpDefaults(this KernelMemoryBuilder builder, LLamaSharpConfig config)
{ {
builder.WithLLamaSharpTextEmbeddingGeneration(config); var parameters = new ModelParams(config.ModelPath)
builder.WithLLamaSharpTextGeneration(config); {
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; return builder;
} }
} }

View File

@ -1,4 +1,5 @@
using LLama; using LLama;
using LLama.Abstractions;
using LLama.Common; using LLama.Common;
using Microsoft.SemanticKernel.AI.Embeddings; using Microsoft.SemanticKernel.AI.Embeddings;
using System; using System;
@ -14,9 +15,11 @@ namespace LLamaSharp.KernelMemory
/// </summary> /// </summary>
public class LLamaSharpTextEmbeddingGeneration : ITextEmbeddingGeneration, IDisposable public class LLamaSharpTextEmbeddingGeneration : ITextEmbeddingGeneration, IDisposable
{ {
private readonly LLamaSharpConfig _config; private readonly LLamaSharpConfig? _config;
private readonly LLamaWeights? _weights;
private readonly LLamaEmbedder _embedder; private readonly LLamaEmbedder _embedder;
private readonly LLamaWeights _weights; private bool _ownsEmbedder = false;
private bool _ownsWeights = false;
/// <summary> /// <summary>
/// Initializes a new instance of the <see cref="LLamaSharpTextEmbeddingGeneration"/> class. /// Initializes a new instance of the <see cref="LLamaSharpTextEmbeddingGeneration"/> class.
@ -28,13 +31,46 @@ namespace LLamaSharp.KernelMemory
var @params = new ModelParams(_config.ModelPath); var @params = new ModelParams(_config.ModelPath);
_weights = LLamaWeights.LoadFromFile(@params); _weights = LLamaWeights.LoadFromFile(@params);
_embedder = new LLamaEmbedder(_weights, @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/> /// <inheritdoc/>
public void Dispose() public void Dispose()
{ {
_embedder.Dispose(); if (_ownsWeights)
_weights.Dispose(); {
_weights?.Dispose();
}
if(_ownsEmbedder)
{
_embedder.Dispose();
}
} }
/// <inheritdoc/> /// <inheritdoc/>

View File

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

View File

@ -1,4 +1,5 @@
using LLama; using LLama;
using LLama.Abstractions;
using LLama.Common; using LLama.Common;
using Microsoft.KernelMemory.AI; using Microsoft.KernelMemory.AI;
using System; using System;
@ -14,10 +15,12 @@ namespace LLamaSharp.KernelMemory
/// </summary> /// </summary>
public class LlamaSharpTextGeneration : ITextGeneration, IDisposable public class LlamaSharpTextGeneration : ITextGeneration, IDisposable
{ {
private readonly LLamaSharpConfig _config; private readonly LLamaSharpConfig? _config;
private readonly LLamaWeights _weights; private readonly LLamaWeights _weights;
private readonly StatelessExecutor _executor; private readonly StatelessExecutor _executor;
private readonly LLamaContext _context; private readonly LLamaContext _context;
private bool _ownsContext = false;
private bool _ownsWeights = false;
/// <summary> /// <summary>
/// Initializes a new instance of the <see cref="LlamaSharpTextGeneration"/> class. /// Initializes a new instance of the <see cref="LlamaSharpTextGeneration"/> class.
@ -35,13 +38,35 @@ namespace LLamaSharp.KernelMemory
_weights = LLamaWeights.LoadFromFile(parameters); _weights = LLamaWeights.LoadFromFile(parameters);
_context = _weights.CreateContext(parameters); _context = _weights.CreateContext(parameters);
_executor = new StatelessExecutor(_weights, 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/> /// <inheritdoc/>
public void Dispose() public void Dispose()
{ {
_context.Dispose(); if (_ownsWeights)
_weights.Dispose(); {
_weights?.Dispose();
}
if (_ownsContext)
{
_context.Dispose();
}
} }
/// <inheritdoc/> /// <inheritdoc/>