Added a property to get the KV cache size from a context

This commit is contained in:
Martin Evans 2023-09-11 00:10:08 +01:00
parent 400c7d4711
commit bba801f4b7
4 changed files with 22 additions and 1 deletions

View File

@ -1,8 +1,9 @@
using System.Text;
using LLama.Common;
namespace LLama.Unittest
{
public class BasicTest
public sealed class BasicTest
: IDisposable
{
private readonly ModelParams _params;
@ -28,6 +29,7 @@ namespace LLama.Unittest
Assert.Equal(32000, _model.VocabCount);
Assert.Equal(2048, _model.ContextSize);
Assert.Equal(4096, _model.EmbeddingSize);
Assert.Equal(Encoding.UTF8, _model.Encoding);
}
[Fact]

View File

@ -30,6 +30,7 @@ namespace LLama.Unittest
Assert.Equal(768, _context.ContextSize);
Assert.Equal(4096, _context.EmbeddingSize);
Assert.Equal(32000, _context.VocabCount);
Assert.Equal(0, _context.KVCacheTokenCount);
}
}
}

View File

@ -41,6 +41,11 @@ namespace LLama
/// </summary>
public int EmbeddingSize => _ctx.EmbeddingSize;
/// <summary>
/// Get the number of tokens in the KV Cache for this context
/// </summary>
public int KVCacheTokenCount => _ctx.KVCacheTokenCount;
/// <summary>
/// The model params set for this model.
/// </summary>

View File

@ -28,6 +28,19 @@ namespace LLama.Native
/// </summary>
public int EmbeddingSize => ThrowIfDisposed().EmbeddingSize;
/// <summary>
/// Get the number of tokens in the KV Cache for this context
/// </summary>
public int KVCacheTokenCount
{
get
{
if (IsClosed)
throw new ObjectDisposedException("Cannot use this `SafeLLamaContextHandle` - it has been disposed");
return NativeApi.llama_get_kv_cache_token_count(this);
}
}
/// <summary>
/// Get the model which this context is using
/// </summary>