Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve Embedding.ConvertToVectorOfFloats #38

Merged
merged 1 commit into from
Jun 10, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Improive Embedding.ConvertToVectorOfFloats
It's currently incorrect on big endian systems (e.g. IBM s390x), as it blits the little-endian bytes into an array of floats. On big endian the values need to be reversed.

It also assumes the data being sent back from the service is always correct. If it's corrupted in certain ways, such as not being as long as was expected, not being quoted, etc., we can silently get bad data.

It also more allocation than is necessary. It first allocates a string for the base64 string data. Then it allocates a new string with the quotes stripped off. Then it allocates a byte[] with the decoded bytes. And finally it allocates the resulting float[]. We can avoid the initial string by getting the raw memory from the BinaryData. We can avoid the substring by just slicing those bytes. And we can generally avoid the temporary byte[] by renting one from the array pool. That leaves just the required float[].
  • Loading branch information
stephentoub committed Jun 10, 2024
commit 0d023facf63fe41eb121db0ee37ca90b3f4e8ced
44 changes: 39 additions & 5 deletions src/Custom/Embeddings/Embedding.cs
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
using System;
using System.Buffers;
using System.Buffers.Binary;
using System.Buffers.Text;
using System.Collections.Generic;
using System.Runtime.InteropServices;

namespace OpenAI.Embeddings;

Expand Down Expand Up @@ -89,13 +93,43 @@ internal Embedding(int index, BinaryData embeddingProperty, InternalEmbeddingObj
// CUSTOM: Implemented custom logic to transform from BinaryData to ReadOnlyMemory<float>.
private ReadOnlyMemory<float> ConvertToVectorOfFloats(BinaryData binaryData)
{
string base64EncodedVector = binaryData.ToString();
base64EncodedVector = base64EncodedVector.Substring(1, base64EncodedVector.Length - 2);
ReadOnlySpan<byte> base64 = binaryData.ToMemory().Span;

byte[] bytes = Convert.FromBase64String(base64EncodedVector);
float[] vector = new float[bytes.Length / sizeof(float)];
Buffer.BlockCopy(bytes, 0, vector, 0, bytes.Length);
// Remove quotes around base64 string.
if (base64.Length < 2 || base64[0] != (byte)'"' || base64[base64.Length - 1] != (byte)'"')
{
ThrowInvalidData();
}
base64 = base64.Slice(1, base64.Length - 2);

// Decode base64 string to bytes.
byte[] bytes = ArrayPool<byte>.Shared.Rent(Base64.GetMaxDecodedFromUtf8Length(base64.Length));
OperationStatus status = Base64.DecodeFromUtf8(base64, bytes.AsSpan(), out int bytesConsumed, out int bytesWritten);
if (status != OperationStatus.Done || bytesWritten % sizeof(float) != 0)
{
ThrowInvalidData();
}

// Interpret bytes as floats
float[] vector = new float[bytesWritten / sizeof(float)];
bytes.AsSpan(0, bytesWritten).CopyTo(MemoryMarshal.AsBytes(vector.AsSpan()));
if (!BitConverter.IsLittleEndian)
stephentoub marked this conversation as resolved.
Show resolved Hide resolved
{
Span<int> ints = MemoryMarshal.Cast<float, int>(vector.AsSpan());
#if NET8_0_OR_GREATER
BinaryPrimitives.ReverseEndianness(ints, ints);
#else
for (int i = 0; i < ints.Length; i++)
{
ints[i] = BinaryPrimitives.ReverseEndianness(ints[i]);
}
#endif
}

ArrayPool<byte>.Shared.Return(bytes);
return new ReadOnlyMemory<float>(vector);

static void ThrowInvalidData() =>
throw new FormatException("The input is not a valid Base64 string of encoded floats.");
}
}
Loading