more accurate byte average

Author: jan-elastic

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingByteResults.java

@@ -120,15 +120,16 @@ public int hashCode() {
 
     // Note: the field "numberOfMergedEmbeddings" is not serialized, so merging
     // embeddings should happen inbetween serializations.
-    public record Embedding(byte[] values, int numberOfMergedEmbeddings)
+    public record Embedding(byte[] values, int[] sumMergedValues, int numberOfMergedEmbeddings)
         implements
             Writeable,
             ToXContentObject,
             EmbeddingResults.Embedding<Chunk, Embedding> {
+
         public static final String EMBEDDING = "embedding";
 
         public Embedding(byte[] values) {
-            this(values, 1);
+            this(values, null, 1);
         }
 
         public Embedding(StreamInput in) throws IOException {
@@ -201,18 +202,19 @@ public Chunk toChunk(String text, ChunkedInference.TextOffset offset) {
             return new Chunk(values, text, offset);
         }
 
-        // This merge function suffers from round-off errors. TODO: maybe do something smarter?
         @Override
         public Embedding merge(Embedding embedding) {
-            byte[] mergedValues = new byte[values.length];
+            byte[] newValues = new byte[values.length];
+            int[] newSumMergedValues = new int[values.length];
             int newNumberOfMergedEmbeddings = numberOfMergedEmbeddings + embedding.numberOfMergedEmbeddings;
             for (int i = 0; i < values.length; i++) {
+                newSumMergedValues[i] = (numberOfMergedEmbeddings == 1 ? values[i] : sumMergedValues[i])
+                    + (embedding.numberOfMergedEmbeddings == 1 ? embedding.values[i] : embedding.sumMergedValues[i]);
                 // Add (newNumberOfMergedEmbeddings / 2) in the numerator to round towards the
                 // closest byte instead of truncating.
-                mergedValues[i] = (byte) ((numberOfMergedEmbeddings * values[i] + embedding.numberOfMergedEmbeddings * embedding.values[i]
-                    + newNumberOfMergedEmbeddings / 2) / newNumberOfMergedEmbeddings);
+                newValues[i] = (byte) ((newSumMergedValues[i] + newNumberOfMergedEmbeddings / 2) / newNumberOfMergedEmbeddings);
             }
-            return new Embedding(mergedValues, newNumberOfMergedEmbeddings);
+            return new Embedding(newValues, newSumMergedValues, newNumberOfMergedEmbeddings);
         }
     }
 

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingByteResultsTests.java

@@ -17,6 +17,7 @@
 import java.util.List;
 import java.util.Map;
 
+import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
 
 public class TextEmbeddingByteResultsTests extends AbstractWireSerializingTestCase<TextEmbeddingByteResults> {
@@ -104,16 +105,24 @@ public void testTransformToCoordinationFormat() {
     }
 
     public void testGetFirstEmbeddingSize() {
-        var firstEmbeddingSize = new TextEmbeddingByteResults(
-            List.of(
-                new TextEmbeddingByteResults.Embedding(new byte[] { (byte) 23, (byte) 24 }),
-                new TextEmbeddingByteResults.Embedding(new byte[] { (byte) 25, (byte) 26 })
-            )
-        ).getFirstEmbeddingSize();
+        var firstEmbeddingSize = new TextEmbeddingByteResults(List.of(
+            new TextEmbeddingByteResults.Embedding(new byte[] { (byte) 23, (byte) 24 }),
+            new TextEmbeddingByteResults.Embedding(new byte[] { (byte) 25, (byte) 26 })
+        )).getFirstEmbeddingSize();
 
         assertThat(firstEmbeddingSize, is(2));
     }
 
+    public void testEmbeddingMerge() {
+        TextEmbeddingByteResults.Embedding embedding1 = new TextEmbeddingByteResults.Embedding(new byte[] { 1, 1, -128 });
+        TextEmbeddingByteResults.Embedding embedding2 = new TextEmbeddingByteResults.Embedding(new byte[] { 1, 0, 127 });
+        TextEmbeddingByteResults.Embedding embedding3 = new TextEmbeddingByteResults.Embedding(new byte[] { 0, 0, 100 });
+        TextEmbeddingByteResults.Embedding mergedEmbedding = embedding1.merge(embedding2);
+        assertThat(mergedEmbedding, equalTo(new TextEmbeddingByteResults.Embedding(new byte[] { 1, 1, 0 })));
+        mergedEmbedding = mergedEmbedding.merge(embedding3);
+        assertThat(mergedEmbedding, equalTo(new TextEmbeddingByteResults.Embedding(new byte[] { 1, 0, 33 })));
+    }
+
     @Override
     protected Writeable.Reader<TextEmbeddingByteResults> instanceReader() {
         return TextEmbeddingByteResults::new;

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunkerTests.java

@@ -461,7 +461,7 @@ public void testVeryLongInput_Byte() {
         assertThat(embedding.chunks(), hasSize(512));
 
         // The first merged chunk consists of 20 small chunks (so 400 words) and the weight
-        // is the average of the weights 2 ... 21, with some round-off errors.
+        // is the average of the weights 2 ... 21, so 11.5, which is rounded to 12.
         assertThat(embedding.chunks().get(0).matchedText(), startsWith("word0 word1 "));
         assertThat(embedding.chunks().get(0).matchedText(), endsWith(" word398 word399"));
         assertThat(embedding.chunks().get(0), instanceOf(TextEmbeddingByteResults.Chunk.class));
@@ -470,7 +470,7 @@ public void testVeryLongInput_Byte() {
 
         // The last merged chunk consists of 19 small chunks (so 380 words) and the weight
         // is the average of the weights 9983 ... 10001 modulo 256 (bytes overflowing), so
-        // the average of -1, 0, 1, ... , 17, with some round-off errors.
+        // the average of -1, 0, 1, ... , 17, so 8.
         assertThat(embedding.chunks().get(511).matchedText(), startsWith(" word199620 word199621 "));
         assertThat(embedding.chunks().get(511).matchedText(), endsWith(" word199998 word199999"));
         assertThat(embedding.chunks().get(511), instanceOf(TextEmbeddingByteResults.Chunk.class));