Removed matched text from chunk (part 2)

Mikep86 · Mikep86 · commit 13e875ed550c · 2025-02-25T15:40:35.000-05:00
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/ChunkedInferenceEmbedding.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/ChunkedInferenceEmbedding.java
@@ -29,7 +29,6 @@ public static List<ChunkedInference> listOf(List<String> inputs, SparseEmbedding
                     List.of(
                         new SparseEmbeddingResults.Chunk(
                             sparseEmbeddingResults.embeddings().get(i).tokens(),
-                            inputs.get(i),
                             new TextOffset(0, inputs.get(i).length())
                         )
                     )
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/EmbeddingResults.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/EmbeddingResults.java
@@ -24,13 +24,11 @@ public interface EmbeddingResults<C extends EmbeddingResults.Chunk, E extends Em
         InferenceServiceResults {
 
     /**
-     * A resulting embedding together with its input text.
+     * A resulting embedding together with the offset into the input text.
      */
     interface Chunk {
         ChunkedInference.Chunk toChunk(XContent xcontent) throws IOException;
 
-        String matchedText();
-
         ChunkedInference.TextOffset offset();
     }
 
@@ -39,9 +37,9 @@ interface Chunk {
      */
     interface Embedding<C extends Chunk> {
         /**
-         * Combines the resulting embedding with the input into a chunk.
+         * Combines the resulting embedding with the offset into the input text into a chunk.
          */
-        C toChunk(String text, ChunkedInference.TextOffset offset);
+        C toChunk(ChunkedInference.TextOffset offset);
     }
 
     /**
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/SparseEmbeddingResults.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/SparseEmbeddingResults.java
@@ -175,14 +175,12 @@ public String toString() {
         }
 
         @Override
-        public Chunk toChunk(String text, ChunkedInference.TextOffset offset) {
-            return new Chunk(tokens, text, offset);
+        public Chunk toChunk(ChunkedInference.TextOffset offset) {
+            return new Chunk(tokens, offset);
         }
     }
 
-    public record Chunk(List<WeightedToken> weightedTokens, String matchedText, ChunkedInference.TextOffset offset)
-        implements
-            EmbeddingResults.Chunk {
+    public record Chunk(List<WeightedToken> weightedTokens, ChunkedInference.TextOffset offset) implements EmbeddingResults.Chunk {
 
         public ChunkedInference.Chunk toChunk(XContent xcontent) throws IOException {
             return new ChunkedInference.Chunk(offset, toBytesReference(xcontent, weightedTokens));
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingByteResults.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingByteResults.java
@@ -187,15 +187,15 @@ public int hashCode() {
         }
 
         @Override
-        public Chunk toChunk(String text, ChunkedInference.TextOffset offset) {
-            return new Chunk(values, text, offset);
+        public Chunk toChunk(ChunkedInference.TextOffset offset) {
+            return new Chunk(values, offset);
         }
     }
 
     /**
      * Serialises the {@code value} array, according to the provided {@link XContent}, into a {@link BytesReference}.
      */
-    public record Chunk(byte[] embedding, String matchedText, ChunkedInference.TextOffset offset) implements EmbeddingResults.Chunk {
+    public record Chunk(byte[] embedding, ChunkedInference.TextOffset offset) implements EmbeddingResults.Chunk {
 
         public ChunkedInference.Chunk toChunk(XContent xcontent) throws IOException {
             return new ChunkedInference.Chunk(offset, toBytesReference(xcontent, embedding));
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingFloatResults.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingFloatResults.java
@@ -221,12 +221,12 @@ public int hashCode() {
         }
 
         @Override
-        public Chunk toChunk(String text, ChunkedInference.TextOffset offset) {
-            return new Chunk(values, text, offset);
+        public Chunk toChunk(ChunkedInference.TextOffset offset) {
+            return new Chunk(values, offset);
         }
     }
 
-    public record Chunk(float[] embedding, String matchedText, ChunkedInference.TextOffset offset) implements EmbeddingResults.Chunk {
+    public record Chunk(float[] embedding, ChunkedInference.TextOffset offset) implements EmbeddingResults.Chunk {
 
         public ChunkedInference.Chunk toChunk(XContent xcontent) throws IOException {
             return new ChunkedInference.Chunk(offset, toBytesReference(xcontent, embedding));
diff --git a/x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestDenseInferenceServiceExtension.java b/x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestDenseInferenceServiceExtension.java
@@ -183,7 +183,6 @@ private List<ChunkedInference> makeChunkedResults(List<String> input, int dimens
                         List.of(
                             new TextEmbeddingFloatResults.Chunk(
                                 nonChunkedResults.embeddings().get(i).values(),
-                                input.get(i),
                                 new ChunkedInference.TextOffset(0, input.get(i).length())
                             )
                         )
diff --git a/x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestSparseInferenceServiceExtension.java b/x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestSparseInferenceServiceExtension.java
@@ -172,13 +172,7 @@ private List<ChunkedInference> makeChunkedResults(List<String> input) {
                 }
                 results.add(
                     new ChunkedInferenceEmbedding(
-                        List.of(
-                            new SparseEmbeddingResults.Chunk(
-                                tokens,
-                                input.get(i),
-                                new ChunkedInference.TextOffset(0, input.get(i).length())
-                            )
-                        )
+                        List.of(new SparseEmbeddingResults.Chunk(tokens, new ChunkedInference.TextOffset(0, input.get(i).length())))
                     )
                 );
             }
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunker.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunker.java
@@ -197,10 +197,7 @@ private ChunkedInference mergeResultsWithInputs(int index) {
         AtomicReferenceArray<EmbeddingResults.Embedding<?>> result = results.get(index);
         for (int i = 0; i < request.size(); i++) {
             EmbeddingResults.Chunk chunk = result.get(i)
-                .toChunk(
-                    request.get(i).chunkText(),
-                    new ChunkedInference.TextOffset(request.get(i).chunk.start(), request.get(i).chunk.end())
-                );
+                .toChunk(new ChunkedInference.TextOffset(request.get(i).chunk.start(), request.get(i).chunk.end()));
             chunks.add(chunk);
         }
         return new ChunkedInferenceEmbedding(chunks);
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/elser/HuggingFaceElserService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/elser/HuggingFaceElserService.java
@@ -121,7 +121,6 @@ private static List<ChunkedInference> translateToChunkedResults(DocumentsOnlyInp
                         List.of(
                             new TextEmbeddingFloatResults.Chunk(
                                 textEmbeddingResults.embeddings().get(i).values(),
-                                inputs.getInputs().get(i),
                                 new ChunkedInference.TextOffset(0, inputs.getInputs().get(i).length())
                             )
                         )
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldTests.java
@@ -177,7 +177,7 @@ public static ChunkedInferenceEmbedding randomChunkedInferenceEmbeddingByte(Mode
             for (int j = 0; j < values.length; j++) {
                 values[j] = randomByte();
             }
-            chunks.add(new TextEmbeddingByteResults.Chunk(values, input, new ChunkedInference.TextOffset(0, input.length())));
+            chunks.add(new TextEmbeddingByteResults.Chunk(values, new ChunkedInference.TextOffset(0, input.length())));
         }
         return new ChunkedInferenceEmbedding(chunks);
     }
@@ -189,7 +189,7 @@ public static ChunkedInferenceEmbedding randomChunkedInferenceEmbeddingFloat(Mod
             for (int j = 0; j < values.length; j++) {
                 values[j] = randomFloat();
             }
-            chunks.add(new TextEmbeddingFloatResults.Chunk(values, input, new ChunkedInference.TextOffset(0, input.length())));
+            chunks.add(new TextEmbeddingFloatResults.Chunk(values, new ChunkedInference.TextOffset(0, input.length())));
         }
         return new ChunkedInferenceEmbedding(chunks);
     }
@@ -205,7 +205,7 @@ public static ChunkedInferenceEmbedding randomChunkedInferenceEmbeddingSparse(Li
             for (var token : input.split("\\s+")) {
                 tokens.add(new WeightedToken(token, withFloats ? randomFloat() : randomIntBetween(1, 255)));
             }
-            chunks.add(new SparseEmbeddingResults.Chunk(tokens, input, new ChunkedInference.TextOffset(0, input.length())));
+            chunks.add(new SparseEmbeddingResults.Chunk(tokens, new ChunkedInference.TextOffset(0, input.length())));
         }
         return new ChunkedInferenceEmbedding(chunks);
     }
@@ -308,7 +308,7 @@ public static ChunkedInference toChunkedResult(
                         String matchedText = matchedTextIt.next();
                         ChunkedInference.TextOffset offset = createOffset(useLegacyFormat, chunk, matchedText);
                         var tokens = parseWeightedTokens(chunk.rawEmbeddings(), field.contentType());
-                        chunks.add(new SparseEmbeddingResults.Chunk(tokens, matchedText, offset));
+                        chunks.add(new SparseEmbeddingResults.Chunk(tokens, offset));
                     }
                 }
                 return new ChunkedInferenceEmbedding(chunks);
@@ -329,7 +329,7 @@ public static ChunkedInference toChunkedResult(
                             field.inference().modelSettings().dimensions(),
                             field.contentType()
                         );
-                        chunks.add(new TextEmbeddingFloatResults.Chunk(FloatConversionUtils.floatArrayOf(values), matchedText, offset));
+                        chunks.add(new TextEmbeddingFloatResults.Chunk(FloatConversionUtils.floatArrayOf(values), offset));
                     }
                 }
                 return new ChunkedInferenceEmbedding(chunks);
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceTests.java
@@ -562,7 +562,6 @@ public void testChunkedInfer_PassesThrough() throws IOException {
                     List.of(
                         new SparseEmbeddingResults.Chunk(
                             List.of(new WeightedToken("hello", 2.1259406f), new WeightedToken("greet", 1.7073475f)),
-                            "input text",
                             new ChunkedInference.TextOffset(0, "input text".length())
                         )
                     )
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceElserServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceElserServiceTests.java
@@ -111,7 +111,6 @@ public void testChunkedInfer_CallsInfer_Elser_ConvertsFloatResponse() throws IOE
                     List.of(
                         new SparseEmbeddingResults.Chunk(
                             List.of(new WeightedToken(".", 0.13315596f)),
-                            "abc",
                             new ChunkedInference.TextOffset(0, "abc".length())
                         )
                     )

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,6 @@ public static List<ChunkedInference> listOf(List<String> inputs, SparseEmbedding`
`29`	`29`	`List.of(`
`30`	`30`	`new SparseEmbeddingResults.Chunk(`
`31`	`31`	`sparseEmbeddingResults.embeddings().get(i).tokens(),`
`32`		`- inputs.get(i),`
`33`	`32`	`new TextOffset(0, inputs.get(i).length())`
`34`	`33`	`)`
`35`	`34`	`)`
Original file line number	Diff line number	Diff line change
`@@ -175,14 +175,12 @@ public String toString() {`
`175`	`175`	`}`
`176`	`176`
`177`	`177`	`@Override`
`178`		`- public Chunk toChunk(String text, ChunkedInference.TextOffset offset) {`
`179`		`- return new Chunk(tokens, text, offset);`
	`178`	`+ public Chunk toChunk(ChunkedInference.TextOffset offset) {`
	`179`	`+ return new Chunk(tokens, offset);`
`180`	`180`	`}`
`181`	`181`	`}`
`182`	`182`
`183`		`- public record Chunk(List<WeightedToken> weightedTokens, String matchedText, ChunkedInference.TextOffset offset)`
`184`		`- implements`
`185`		`- EmbeddingResults.Chunk {`
	`183`	`+ public record Chunk(List<WeightedToken> weightedTokens, ChunkedInference.TextOffset offset) implements EmbeddingResults.Chunk {`
`186`	`184`
`187`	`185`	`public ChunkedInference.Chunk toChunk(XContent xcontent) throws IOException {`
`188`	`186`	`return new ChunkedInference.Chunk(offset, toBytesReference(xcontent, weightedTokens));`
Original file line number	Diff line number	Diff line change
`@@ -187,15 +187,15 @@ public int hashCode() {`
`187`	`187`	`}`
`188`	`188`
`189`	`189`	`@Override`
`190`		`- public Chunk toChunk(String text, ChunkedInference.TextOffset offset) {`
`191`		`- return new Chunk(values, text, offset);`
	`190`	`+ public Chunk toChunk(ChunkedInference.TextOffset offset) {`
	`191`	`+ return new Chunk(values, offset);`
`192`	`192`	`}`
`193`	`193`	`}`
`194`	`194`
`195`	`195`	`/**`
`196`	`196`	`* Serialises the {@code value} array, according to the provided {@link XContent}, into a {@link BytesReference}.`
`197`	`197`	`*/`
`198`		`- public record Chunk(byte[] embedding, String matchedText, ChunkedInference.TextOffset offset) implements EmbeddingResults.Chunk {`
	`198`	`+ public record Chunk(byte[] embedding, ChunkedInference.TextOffset offset) implements EmbeddingResults.Chunk {`
`199`	`199`
`200`	`200`	`public ChunkedInference.Chunk toChunk(XContent xcontent) throws IOException {`
`201`	`201`	`return new ChunkedInference.Chunk(offset, toBytesReference(xcontent, embedding));`
Original file line number	Diff line number	Diff line change
`@@ -221,12 +221,12 @@ public int hashCode() {`
`221`	`221`	`}`
`222`	`222`
`223`	`223`	`@Override`
`224`		`- public Chunk toChunk(String text, ChunkedInference.TextOffset offset) {`
`225`		`- return new Chunk(values, text, offset);`
	`224`	`+ public Chunk toChunk(ChunkedInference.TextOffset offset) {`
	`225`	`+ return new Chunk(values, offset);`
`226`	`226`	`}`
`227`	`227`	`}`
`228`	`228`
`229`		`- public record Chunk(float[] embedding, String matchedText, ChunkedInference.TextOffset offset) implements EmbeddingResults.Chunk {`
	`229`	`+ public record Chunk(float[] embedding, ChunkedInference.TextOffset offset) implements EmbeddingResults.Chunk {`
`230`	`230`
`231`	`231`	`public ChunkedInference.Chunk toChunk(XContent xcontent) throws IOException {`
`232`	`232`	`return new ChunkedInference.Chunk(offset, toBytesReference(xcontent, embedding));`
Original file line number	Diff line number	Diff line change
`@@ -183,7 +183,6 @@ private List<ChunkedInference> makeChunkedResults(List<String> input, int dimens`
`183`	`183`	`List.of(`
`184`	`184`	`new TextEmbeddingFloatResults.Chunk(`
`185`	`185`	`nonChunkedResults.embeddings().get(i).values(),`
`186`		`- input.get(i),`
`187`	`186`	`new ChunkedInference.TextOffset(0, input.get(i).length())`
`188`	`187`	`)`
`189`	`188`	`)`
Original file line number	Diff line number	Diff line change
`@@ -172,13 +172,7 @@ private List<ChunkedInference> makeChunkedResults(List<String> input) {`
`172`	`172`	`}`
`173`	`173`	`results.add(`
`174`	`174`	`new ChunkedInferenceEmbedding(`
`175`		`- List.of(`
`176`		`- new SparseEmbeddingResults.Chunk(`
`177`		`- tokens,`
`178`		`- input.get(i),`
`179`		`- new ChunkedInference.TextOffset(0, input.get(i).length())`
`180`		`- )`
`181`		`- )`
	`175`	`+ List.of(new SparseEmbeddingResults.Chunk(tokens, new ChunkedInference.TextOffset(0, input.get(i).length())))`
`182`	`176`	`)`
`183`	`177`	`);`
`184`	`178`	`}`
Original file line number	Diff line number	Diff line change
`@@ -121,7 +121,6 @@ private static List<ChunkedInference> translateToChunkedResults(DocumentsOnlyInp`
`121`	`121`	`List.of(`
`122`	`122`	`new TextEmbeddingFloatResults.Chunk(`
`123`	`123`	`textEmbeddingResults.embeddings().get(i).values(),`
`124`		`- inputs.getInputs().get(i),`
`125`	`124`	`new ChunkedInference.TextOffset(0, inputs.getInputs().get(i).length())`
`126`	`125`	`)`
`127`	`126`	`)`
Original file line number	Diff line number	Diff line change
`@@ -177,7 +177,7 @@ public static ChunkedInferenceEmbedding randomChunkedInferenceEmbeddingByte(Mode`
`177`	`177`	`for (int j = 0; j < values.length; j++) {`
`178`	`178`	`values[j] = randomByte();`
`179`	`179`	`}`
`180`		`- chunks.add(new TextEmbeddingByteResults.Chunk(values, input, new ChunkedInference.TextOffset(0, input.length())));`
	`180`	`+ chunks.add(new TextEmbeddingByteResults.Chunk(values, new ChunkedInference.TextOffset(0, input.length())));`
`181`	`181`	`}`
`182`	`182`	`return new ChunkedInferenceEmbedding(chunks);`
`183`	`183`	`}`
`@@ -189,7 +189,7 @@ public static ChunkedInferenceEmbedding randomChunkedInferenceEmbeddingFloat(Mod`
`189`	`189`	`for (int j = 0; j < values.length; j++) {`
`190`	`190`	`values[j] = randomFloat();`
`191`	`191`	`}`
`192`		`- chunks.add(new TextEmbeddingFloatResults.Chunk(values, input, new ChunkedInference.TextOffset(0, input.length())));`
	`192`	`+ chunks.add(new TextEmbeddingFloatResults.Chunk(values, new ChunkedInference.TextOffset(0, input.length())));`
`193`	`193`	`}`
`194`	`194`	`return new ChunkedInferenceEmbedding(chunks);`
`195`	`195`	`}`
`@@ -205,7 +205,7 @@ public static ChunkedInferenceEmbedding randomChunkedInferenceEmbeddingSparse(Li`
`205`	`205`	`for (var token : input.split("\\s+")) {`
`206`	`206`	`tokens.add(new WeightedToken(token, withFloats ? randomFloat() : randomIntBetween(1, 255)));`
`207`	`207`	`}`
`208`		`- chunks.add(new SparseEmbeddingResults.Chunk(tokens, input, new ChunkedInference.TextOffset(0, input.length())));`
	`208`	`+ chunks.add(new SparseEmbeddingResults.Chunk(tokens, new ChunkedInference.TextOffset(0, input.length())));`
`209`	`209`	`}`
`210`	`210`	`return new ChunkedInferenceEmbedding(chunks);`
`211`	`211`	`}`
`@@ -308,7 +308,7 @@ public static ChunkedInference toChunkedResult(`
`308`	`308`	`String matchedText = matchedTextIt.next();`
`309`	`309`	`ChunkedInference.TextOffset offset = createOffset(useLegacyFormat, chunk, matchedText);`
`310`	`310`	`var tokens = parseWeightedTokens(chunk.rawEmbeddings(), field.contentType());`
`311`		`- chunks.add(new SparseEmbeddingResults.Chunk(tokens, matchedText, offset));`
	`311`	`+ chunks.add(new SparseEmbeddingResults.Chunk(tokens, offset));`
`312`	`312`	`}`
`313`	`313`	`}`
`314`	`314`	`return new ChunkedInferenceEmbedding(chunks);`
`@@ -329,7 +329,7 @@ public static ChunkedInference toChunkedResult(`
`329`	`329`	`field.inference().modelSettings().dimensions(),`
`330`	`330`	`field.contentType()`
`331`	`331`	`);`
`332`		`- chunks.add(new TextEmbeddingFloatResults.Chunk(FloatConversionUtils.floatArrayOf(values), matchedText, offset));`
	`332`	`+ chunks.add(new TextEmbeddingFloatResults.Chunk(FloatConversionUtils.floatArrayOf(values), offset));`
`333`	`333`	`}`
`334`	`334`	`}`
`335`	`335`	`return new ChunkedInferenceEmbedding(chunks);`