Merge pull request #384 from martindevans/fixed_metadata_large_values

Fixed loading of very large metadata values
This commit is contained in:
Martin Evans 2023-12-24 20:09:42 +00:00 committed by GitHub
commit 82baee7133
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 43 additions and 32 deletions

View File

@ -1,10 +1,9 @@
using System; using System;
using System.Buffers;
using System.Collections.Generic; using System.Collections.Generic;
using System.Diagnostics;
using System.Text; using System.Text;
using LLama.Exceptions; using LLama.Exceptions;
using LLama.Extensions; using LLama.Extensions;
using EncodingExtensions = LLama.Extensions.EncodingExtensions;
namespace LLama.Native namespace LLama.Native
{ {
@ -196,7 +195,7 @@ namespace LLama.Native
} }
} }
} }
#endregion #endregion
#region context #region context
/// <summary> /// <summary>
@ -217,15 +216,26 @@ namespace LLama.Native
/// <param name="index">The index to get</param> /// <param name="index">The index to get</param>
/// <param name="buffer">A temporary buffer to store key characters in. Must be large enough to contain the key.</param> /// <param name="buffer">A temporary buffer to store key characters in. Must be large enough to contain the key.</param>
/// <returns>The key, null if there is no such key or if the buffer was too small</returns> /// <returns>The key, null if there is no such key or if the buffer was too small</returns>
public Memory<byte>? MetadataKeyByIndex(int index, Memory<byte> buffer) public Memory<byte>? MetadataKeyByIndex(int index)
{ {
int keyLength;
unsafe unsafe
{ {
using var pin = buffer.Pin(); // Check if the key exists, without getting any bytes of data
var keyLength = NativeApi.llama_model_meta_key_by_index(this, index, (byte*)pin.Pointer, buffer.Length); keyLength = NativeApi.llama_model_meta_key_by_index(this, index, (byte*)0, 0);
if (keyLength < 0) if (keyLength < 0)
return null; return null;
return buffer.Slice(0, keyLength); }
// get a buffer large enough to hold it
var buffer = new byte[keyLength + 1];
unsafe
{
using var pin = buffer.AsMemory().Pin();
keyLength = NativeApi.llama_model_meta_key_by_index(this, index, (byte*)pin.Pointer, buffer.Length);
Debug.Assert(keyLength >= 0);
return buffer.AsMemory().Slice(0, keyLength);
} }
} }
@ -235,15 +245,26 @@ namespace LLama.Native
/// <param name="index">The index to get</param> /// <param name="index">The index to get</param>
/// <param name="buffer">A temporary buffer to store value characters in. Must be large enough to contain the value.</param> /// <param name="buffer">A temporary buffer to store value characters in. Must be large enough to contain the value.</param>
/// <returns>The value, null if there is no such value or if the buffer was too small</returns> /// <returns>The value, null if there is no such value or if the buffer was too small</returns>
public Memory<byte>? MetadataValueByIndex(int index, Memory<byte> buffer) public Memory<byte>? MetadataValueByIndex(int index)
{ {
int valueLength;
unsafe unsafe
{ {
using var pin = buffer.Pin(); // Check if the key exists, without getting any bytes of data
var keyLength = NativeApi.llama_model_meta_val_str_by_index(this, index, (byte*)pin.Pointer, buffer.Length); valueLength = NativeApi.llama_model_meta_val_str_by_index(this, index, (byte*)0, 0);
if (keyLength < 0) if (valueLength < 0)
return null; return null;
return buffer.Slice(0, keyLength); }
// get a buffer large enough to hold it
var buffer = new byte[valueLength + 1];
unsafe
{
using var pin = buffer.AsMemory().Pin();
valueLength = NativeApi.llama_model_meta_val_str_by_index(this, index, (byte*)pin.Pointer, buffer.Length);
Debug.Assert(valueLength >= 0);
return buffer.AsMemory().Slice(0, valueLength);
} }
} }
@ -251,29 +272,19 @@ namespace LLama.Native
{ {
var result = new Dictionary<string, string>(); var result = new Dictionary<string, string>();
var dest = ArrayPool<byte>.Shared.Rent(1024); for (var i = 0; i < MetadataCount; i++)
try
{ {
for (var i = 0; i < MetadataCount; i++) var keyBytes = MetadataKeyByIndex(i);
{ if (keyBytes == null)
Array.Clear(dest, 0, dest.Length); continue;
var key = Encoding.UTF8.GetStringFromSpan(keyBytes.Value.Span);
var keyBytes = MetadataKeyByIndex(i, dest.AsMemory()); var valBytes = MetadataValueByIndex(i);
if (keyBytes == null) if (valBytes == null)
continue; continue;
var key = Encoding.UTF8.GetStringFromSpan(keyBytes.Value.Span); var val = Encoding.UTF8.GetStringFromSpan(valBytes.Value.Span);
var valBytes = MetadataValueByIndex(i, dest.AsMemory()); result[key] = val;
if (valBytes == null)
continue;
var val = Encoding.UTF8.GetStringFromSpan(valBytes.Value.Span);
result[key] = val;
}
}
finally
{
ArrayPool<byte>.Shared.Return(dest);
} }
return result; return result;