Add BinaryEmbedding (#6398)

* Add BinaryEmbedding Also: - Renames the polymorphic discriminators to conform with typical lingo for these types. - Adds an Embedding.Dimensions virtual property.
dotnet · joperezr · May 14, 2025 · May 9, 2025 · May 8, 2025 · May 8, 2025
commit 49bb2c5ea0326f7f055ac42c538e59834c159b13
@@ -0,0 +1,111 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Buffers;
+using System.Collections;
+using System.ComponentModel;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Extensions.AI;
+
+/// <summary>Represents an embedding composed of a bit vector.</summary>
+public sealed class BinaryEmbedding : Embedding
+{
+    /// <summary>The embedding vector this embedding represents.</summary>
+    private BitArray _vector;
+
+    /// <summary>Initializes a new instance of the <see cref="BinaryEmbedding"/> class with the embedding vector.</summary>
+    /// <param name="vector">The embedding vector this embedding represents.</param>
+    /// <exception cref="ArgumentNullException"><paramref name="vector"/> is <see langword="null"/>.</exception>
+    public BinaryEmbedding(BitArray vector)
+    {
+        _vector = Throw.IfNull(vector);
+    }
+
+    /// <summary>Gets or sets the embedding vector this embedding represents.</summary>
+    [JsonConverter(typeof(VectorConverter))]
+    public BitArray Vector
+    {
+        get => _vector;
+        set => _vector = Throw.IfNull(value);
+    }
+
+    /// <inheritdoc />
+    [JsonIgnore]
+    public override int Dimensions => _vector.Length;
+
+    /// <summary>Provides a <see cref="JsonConverter{BitArray}"/> for serializing <see cref="BitArray"/> instances.</summary>
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    public sealed class VectorConverter : JsonConverter<BitArray>
+    {
+        /// <inheritdoc/>
+        public override BitArray Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
+        {
+            _ = Throw.IfNull(typeToConvert);
+            _ = Throw.IfNull(options);
+
+            if (reader.TokenType != JsonTokenType.String)
+            {
+                throw new JsonException("Expected string property.");
+            }
+
+            ReadOnlySpan<byte> utf8;
+            byte[]? tmpArray = null;
+            if (!reader.HasValueSequence && !reader.ValueIsEscaped)
+            {
+                utf8 = reader.ValueSpan;
+            }
+            else
+            {
+                // This path should be rare.
+                int length = reader.HasValueSequence ? checked((int)reader.ValueSequence.Length) : reader.ValueSpan.Length;
+                tmpArray = ArrayPool<byte>.Shared.Rent(length);
+                utf8 = tmpArray.AsSpan(0, reader.CopyString(tmpArray));
+            }
+
+            BitArray result = new(utf8.Length);
+
+            for (int i = 0; i < utf8.Length; i++)
+            {
+                result[i] = utf8[i] switch
+                {
+                    (byte)'0' => false,
+                    (byte)'1' => true,
+                    _ => throw new JsonException("Expected binary character sequence.")
+                };
+            }
+
+            if (tmpArray is not null)
+            {
+                ArrayPool<byte>.Shared.Return(tmpArray);
+            }
+
+            return result;
+        }
+
+        /// <inheritdoc/>
+        public override void Write(Utf8JsonWriter writer, BitArray value, JsonSerializerOptions options)
+        {
+            _ = Throw.IfNull(writer);
+            _ = Throw.IfNull(value);
+            _ = Throw.IfNull(options);
+
+            int length = value.Length;
+
+            byte[] tmpArray = ArrayPool<byte>.Shared.Rent(length);
+
+            Span<byte> utf8 = tmpArray.AsSpan(0, length);
+            for (int i = 0; i < utf8.Length; i++)
+            {
+                utf8[i] = value[i] ? (byte)'1' : (byte)'0';
+            }
+
+            writer.WriteStringValue(utf8);
+
+            ArrayPool<byte>.Shared.Return(tmpArray);
+        }
+    }
+}
@@ -2,20 +2,23 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Diagnostics;
 using System.Text.Json.Serialization;
 
 namespace Microsoft.Extensions.AI;
 
 /// <summary>Represents an embedding generated by a <see cref="IEmbeddingGenerator{TInput, TEmbedding}"/>.</summary>
 /// <remarks>This base class provides metadata about the embedding. Derived types provide the concrete data contained in the embedding.</remarks>
 [JsonPolymorphic(TypeDiscriminatorPropertyName = "$type")]
+[JsonDerivedType(typeof(BinaryEmbedding), typeDiscriminator: "binary")]
+[JsonDerivedType(typeof(Embedding<byte>), typeDiscriminator: "uint8")]
+[JsonDerivedType(typeof(Embedding<sbyte>), typeDiscriminator: "int8")]
 #if NET
-[JsonDerivedType(typeof(Embedding<Half>), typeDiscriminator: "halves")]
+[JsonDerivedType(typeof(Embedding<Half>), typeDiscriminator: "float16")]
 #endif
-[JsonDerivedType(typeof(Embedding<float>), typeDiscriminator: "floats")]
-[JsonDerivedType(typeof(Embedding<double>), typeDiscriminator: "doubles")]
-[JsonDerivedType(typeof(Embedding<byte>), typeDiscriminator: "bytes")]
-[JsonDerivedType(typeof(Embedding<sbyte>), typeDiscriminator: "sbytes")]
+[JsonDerivedType(typeof(Embedding<float>), typeDiscriminator: "float32")]
+[JsonDerivedType(typeof(Embedding<double>), typeDiscriminator: "float64")]
+[DebuggerDisplay("Dimensions = {Dimensions}")]
 public class Embedding
 {
     /// <summary>Initializes a new instance of the <see cref="Embedding"/> class.</summary>
@@ -26,6 +29,13 @@ protected Embedding()
     /// <summary>Gets or sets a timestamp at which the embedding was created.</summary>
     public DateTimeOffset? CreatedAt { get; set; }
 
+    /// <summary>Gets the dimensionality of the embedding vector.</summary>
+    /// <remarks>
+    /// This value corresponds to the number of elements in the embedding vector.
+    /// </remarks>
+    [JsonIgnore]
+    public virtual int Dimensions { get; }
+
     /// <summary>Gets or sets the model ID using in the creation of the embedding.</summary>
     public string? ModelId { get; set; }
 

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Text.Json.Serialization;
 
 namespace Microsoft.Extensions.AI;
 
@@ -19,4 +20,8 @@ public Embedding(ReadOnlyMemory<T> vector)
 
     /// <summary>Gets or sets the embedding vector this embedding represents.</summary>
     public ReadOnlyMemory<T> Vector { get; set; }
+
+    /// <inheritdoc />
+    [JsonIgnore]
+    public override int Dimensions => Vector.Length;
 }
@@ -0,0 +1,95 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections;
+using System.Linq;
+using System.Text.Json;
+using Xunit;
+
+namespace Microsoft.Extensions.AI;
+
+public class BinaryEmbeddingTests
+{
+    [Fact]
+    public void Ctor_Roundtrips()
+    {
+        BitArray vector = new BitArray(new bool[] { false, true, false, true });
+
+        BinaryEmbedding e = new(vector);
+        Assert.Same(vector, e.Vector);
+        Assert.Null(e.ModelId);
+        Assert.Null(e.CreatedAt);
+        Assert.Null(e.AdditionalProperties);
+    }
+
+    [Fact]
+    public void Properties_Roundtrips()
+    {
+        BitArray vector = new BitArray(new bool[] { false, true, false, true });
+
+        BinaryEmbedding e = new(vector);
+
+        Assert.Same(vector, e.Vector);
+        BitArray newVector = new BitArray(new bool[] { true, false, true, false });
+        e.Vector = newVector;
+        Assert.Same(newVector, e.Vector);
+
+        Assert.Null(e.ModelId);
+        e.ModelId = "text-embedding-3-small";
+        Assert.Equal("text-embedding-3-small", e.ModelId);
+
+        Assert.Null(e.CreatedAt);
+        DateTimeOffset createdAt = DateTimeOffset.Parse("2022-01-01T00:00:00Z");
+        e.CreatedAt = createdAt;
+        Assert.Equal(createdAt, e.CreatedAt);
+
+        Assert.Null(e.AdditionalProperties);
+        AdditionalPropertiesDictionary props = new();
+        e.AdditionalProperties = props;
+        Assert.Same(props, e.AdditionalProperties);
+    }
+
+    [Fact]
+    public void Serialization_Roundtrips()
+    {
+        foreach (int length in Enumerable.Range(0, 64).Concat(new[] { 10_000 }))
+        {
+            bool[] bools = new bool[length];
+            Random r = new(42);
+            for (int i = 0; i < length; i++)
+            {
+                bools[i] = r.Next(2) != 0;
+            }
+
+            BitArray vector = new BitArray(bools);
+            BinaryEmbedding e = new(vector);
+
+            string json = JsonSerializer.Serialize(e, TestJsonSerializerContext.Default.Embedding);
+            Assert.Equal($$"""{"$type":"binary","vector":"{{string.Concat(vector.Cast<bool>().Select(b => b ? '1' : '0'))}}"}""", json);
+
+            BinaryEmbedding result = Assert.IsType<BinaryEmbedding>(JsonSerializer.Deserialize(json, TestJsonSerializerContext.Default.Embedding));
+            Assert.Equal(e.Vector, result.Vector);
+        }
+    }
+
+    [Fact]
+    public void Derialization_SupportsEncodedBits()
+    {
+        BinaryEmbedding result = Assert.IsType<BinaryEmbedding>(JsonSerializer.Deserialize(
+            """{"$type":"binary","vector":"\u0030\u0031\u0030\u0031\u0030\u0031"}""",
+            TestJsonSerializerContext.Default.Embedding));
+
+        Assert.Equal(new BitArray(new[] { false, true, false, true, false, true }), result.Vector);
+    }
+
+    [Theory]
+    [InlineData("""{"$type":"binary","vector":"\u0030\u0032"}""")]
+    [InlineData("""{"$type":"binary","vector":"02"}""")]
+    [InlineData("""{"$type":"binary","vector":" "}""")]
+    [InlineData("""{"$type":"binary","vector":10101}""")]
+    public void Derialization_InvalidBinaryEmbedding_Throws(string json)
+    {
+        Assert.Throws<JsonException>(() => JsonSerializer.Deserialize(json, TestJsonSerializerContext.Default.Embedding));
+    }
+}
@@ -14,7 +14,7 @@ public class EmbeddingTests
     public void Embedding_Ctor_Roundtrips()
     {
         float[] floats = [1f, 2f, 3f];
-        UsageDetails usage = new();
+
         AdditionalPropertiesDictionary props = [];
         var createdAt = DateTimeOffset.Parse("2022-01-01T00:00:00Z");
         const string Model = "text-embedding-3-small";
@@ -35,6 +35,32 @@ public void Embedding_Ctor_Roundtrips()
         Assert.Same(floats, array.Array);
     }
 
+    [Fact]
+    public void Embedding_Byte_SerializationRoundtrips()
+    {
+        byte[] bytes = [1, 2, 3];
+        Embedding<byte> e = new(bytes);
+
+        string json = JsonSerializer.Serialize(e, TestJsonSerializerContext.Default.Embedding);
+        Assert.Equal("""{"$type":"uint8","vector":"AQID"}""", json);
+
+        Embedding<byte> result = Assert.IsType<Embedding<byte>>(JsonSerializer.Deserialize(json, TestJsonSerializerContext.Default.Embedding));
+        Assert.Equal(e.Vector.ToArray(), result.Vector.ToArray());
+    }
+
+    [Fact]
+    public void Embedding_SByte_SerializationRoundtrips()
+    {
+        sbyte[] bytes = [1, 2, 3];
+        Embedding<sbyte> e = new(bytes);
+
+        string json = JsonSerializer.Serialize(e, TestJsonSerializerContext.Default.Embedding);
+        Assert.Equal("""{"$type":"int8","vector":[1,2,3]}""", json);
+
+        Embedding<sbyte> result = Assert.IsType<Embedding<sbyte>>(JsonSerializer.Deserialize(json, TestJsonSerializerContext.Default.Embedding));
+        Assert.Equal(e.Vector.ToArray(), result.Vector.ToArray());
+    }
+
 #if NET
     [Fact]
     public void Embedding_Half_SerializationRoundtrips()
@@ -43,7 +69,7 @@ public void Embedding_Half_SerializationRoundtrips()
         Embedding<Half> e = new(halfs);
 
         string json = JsonSerializer.Serialize(e, TestJsonSerializerContext.Default.Embedding);
-        Assert.Equal("""{"$type":"halves","vector":[1,2,3]}""", json);
+        Assert.Equal("""{"$type":"float16","vector":[1,2,3]}""", json);
 
         Embedding<Half> result = Assert.IsType<Embedding<Half>>(JsonSerializer.Deserialize(json, TestJsonSerializerContext.Default.Embedding));
         Assert.Equal(e.Vector.ToArray(), result.Vector.ToArray());
@@ -57,7 +83,7 @@ public void Embedding_Single_SerializationRoundtrips()
         Embedding<float> e = new(floats);
 
         string json = JsonSerializer.Serialize(e, TestJsonSerializerContext.Default.Embedding);
-        Assert.Equal("""{"$type":"floats","vector":[1,2,3]}""", json);
+        Assert.Equal("""{"$type":"float32","vector":[1,2,3]}""", json);
 
         Embedding<float> result = Assert.IsType<Embedding<float>>(JsonSerializer.Deserialize(json, TestJsonSerializerContext.Default.Embedding));
         Assert.Equal(e.Vector.ToArray(), result.Vector.ToArray());
@@ -70,7 +96,7 @@ public void Embedding_Double_SerializationRoundtrips()
         Embedding<double> e = new(floats);
 
         string json = JsonSerializer.Serialize(e, TestJsonSerializerContext.Default.Embedding);
-        Assert.Equal("""{"$type":"doubles","vector":[1,2,3]}""", json);
+        Assert.Equal("""{"$type":"float64","vector":[1,2,3]}""", json);
 
         Embedding<double> result = Assert.IsType<Embedding<double>>(JsonSerializer.Deserialize(json, TestJsonSerializerContext.Default.Embedding));
         Assert.Equal(e.Vector.ToArray(), result.Vector.ToArray());

@@ -2,6 +2,9 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+#if NET
+using System.Collections;
+#endif
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
@@ -148,7 +151,14 @@ public async Task Quantization_Binary_EmbeddingsCompareSuccessfully()
         {
             for (int j = 0; j < embeddings.Count; j++)
             {
-                distances[i, j] = TensorPrimitives.HammingBitDistance(embeddings[i].Bits.Span, embeddings[j].Bits.Span);
+                distances[i, j] = TensorPrimitives.HammingBitDistance<byte>(ToArray(embeddings[i].Vector), ToArray(embeddings[j].Vector));
+
+                static byte[] ToArray(BitArray array)
+                {
+                    byte[] result = new byte[(array.Length + 7) / 8];
+                    array.CopyTo(result, 0);
+                    return result;
+                }
             }
         }