From d47b6afe4d67fe67ca3ee51d9f1ee1cf7fd01925 Mon Sep 17 00:00:00 2001
From: Martin Evans <martindevans@gmail.com>
Date: Tue, 13 Feb 2024 02:09:35 +0000
Subject: [PATCH] Normalizing embeddings in `LLamaEmbedder`. As is done in
 llama.cpp:
 https://github.com/ggerganov/llama.cpp/blob/2891c8aa9af17f4ff636ff3868bc34ff72b56e25/examples/embedding/embedding.cpp#L92

---
 LLama.Unittest/LLamaEmbedderTests.cs | 29 +++++------------------
 LLama/LLamaEmbedder.cs               | 35 ++++++++++++++++++++--------
 2 files changed, 31 insertions(+), 33 deletions(-)

diff --git a/LLama.Unittest/LLamaEmbedderTests.cs b/LLama.Unittest/LLamaEmbedderTests.cs
index 05269048..ca15dc6f 100644
--- a/LLama.Unittest/LLamaEmbedderTests.cs
+++ b/LLama.Unittest/LLamaEmbedderTests.cs
@@ -25,18 +25,6 @@ public sealed class LLamaEmbedderTests
         _embedder.Dispose();
     }
 
-    private static float Magnitude(float[] a)
-    {
-        return MathF.Sqrt(a.Zip(a, (x, y) => x * y).Sum());
-    }
-
-    private static void Normalize(float[] a)
-    {
-        var mag = Magnitude(a);
-        for (var i = 0; i < a.Length; i++)
-            a[i] /= mag;
-    }
-
     private static float Dot(float[] a, float[] b)
     {
         Assert.Equal(a.Length, b.Length);
@@ -46,21 +34,16 @@ public sealed class LLamaEmbedderTests
     [Fact]
     public async Task EmbedCompare()
     {
-        var cat = await _embedder.GetEmbeddings("cat");
-        var kitten = await _embedder.GetEmbeddings("kitten");
-        var spoon = await _embedder.GetEmbeddings("spoon");
+        var cat = await _embedder.GetEmbeddings("The cat is cute");
+        var kitten = await _embedder.GetEmbeddings("The kitten is kawaii");
+        var spoon = await _embedder.GetEmbeddings("The spoon is not real");
 
-        Normalize(cat);
-        Normalize(kitten);
-        Normalize(spoon);
+        _testOutputHelper.WriteLine($"Cat    = [{string.Join(",", cat.AsMemory().Slice(0, 7).ToArray())}...]");
+        _testOutputHelper.WriteLine($"Kitten = [{string.Join(",", kitten.AsMemory().Slice(0, 7).ToArray())}...]");
+        _testOutputHelper.WriteLine($"Spoon  = [{string.Join(",", spoon.AsMemory().Slice(0, 7).ToArray())}...]");
 
         var close = 1 - Dot(cat, kitten);
         var far = 1 - Dot(cat, spoon);
-
         Assert.True(close < far);
-
-        _testOutputHelper.WriteLine($"Cat    = [{string.Join(",", cat.AsMemory().Slice(0, 7).ToArray())}...]");
-        _testOutputHelper.WriteLine($"Kitten = [{string.Join(",", kitten.AsMemory().Slice(0, 7).ToArray())}...]");
-        _testOutputHelper.WriteLine($"Spoon  = [{string.Join(",", spoon.AsMemory().Slice(0, 7).ToArray())}...]");
     }
 }
\ No newline at end of file
diff --git a/LLama/LLamaEmbedder.cs b/LLama/LLamaEmbedder.cs
index e25398bb..86ecceb5 100644
--- a/LLama/LLamaEmbedder.cs
+++ b/LLama/LLamaEmbedder.cs
@@ -88,20 +88,35 @@ namespace LLama
             // Remove everything we just evaluated from the context cache
             Context.NativeHandle.KvCacheClear();
 
+            // Normalize the embeddings vector
+            // https://github.com/ggerganov/llama.cpp/blob/2891c8aa9af17f4ff636ff3868bc34ff72b56e25/examples/embedding/embedding.cpp#L92
+            Normalize(embeddings);
+
             return embeddings;
+        }
 
-            float[] GetEmbeddingsArray()
-            {
-                var embeddings = NativeApi.llama_get_embeddings(Context.NativeHandle);
-                if (embeddings == null)
-                    return Array.Empty<float>();
-                return embeddings.ToArray();
-            }
+        private float[] GetEmbeddingsArray()
+        {
+            var embeddings = NativeApi.llama_get_embeddings(Context.NativeHandle);
+            if (embeddings == null)
+                return Array.Empty<float>();
+            return embeddings.ToArray();
         }
 
-        /// <summary>
-        /// 
-        /// </summary>
+        private static void Normalize(Span<float> embeddings)
+        {
+            // Calculate length
+            var lengthSqr = 0.0;
+            foreach (var value in embeddings)
+                lengthSqr += value * value;
+            var length = (float)Math.Sqrt(lengthSqr);
+
+            // Normalize
+            for (var i = 0; i < embeddings.Length; i++)
+                embeddings[i] /= length;
+        }
+
+        /// <inheritdoc />
         public void Dispose()
         {
             Context.Dispose();