Remove matched text from chunks (#123607)

Author: Mikep86

server/src/main/java/org/elasticsearch/inference/ChunkedInference.java

@@ -21,18 +21,17 @@ public interface ChunkedInference {
      * Implementations of this function serialize their embeddings to {@link BytesReference} for storage in semantic text fields.
      *
      * @param xcontent provided by the SemanticTextField
-     * @return an iterator of the serialized {@link Chunk} which includes the matched text (input) and bytes reference (output/embedding).
+     * @return an iterator of the serialized {@link Chunk} which includes the offset into the input text and bytes reference
+     * (output/embedding).
      */
-    Iterator<Chunk> chunksAsMatchedTextAndByteReference(XContent xcontent) throws IOException;
+    Iterator<Chunk> chunksAsByteReference(XContent xcontent) throws IOException;
 
     /**
-     * A chunk of inference results containing matched text, the substring location
-     * in the original text and the bytes reference.
-     * @param matchedText
+     * A chunk of inference results containing the substring location in the original text and the bytes reference.
      * @param textOffset
      * @param bytesReference
      */
-    record Chunk(String matchedText, TextOffset textOffset, BytesReference bytesReference) {}
+    record Chunk(TextOffset textOffset, BytesReference bytesReference) {}
 
     record TextOffset(int start, int end) {}
 }

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())
                         )
                     )
@@ -41,7 +40,7 @@ public static List<ChunkedInference> listOf(List<String> inputs, SparseEmbedding
     }
 
     @Override
-    public Iterator<Chunk> chunksAsMatchedTextAndByteReference(XContent xcontent) throws IOException {
+    public Iterator<Chunk> chunksAsByteReference(XContent xcontent) throws IOException {
         var asChunk = new ArrayList<Chunk>();
         for (var chunk : chunks()) {
             asChunk.add(chunk.toChunk(xcontent));

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

@@ -7,17 +7,16 @@
 
 package org.elasticsearch.xpack.core.inference.results;
 
-import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.inference.ChunkedInference;
 import org.elasticsearch.xcontent.XContent;
 
+import java.util.Collections;
 import java.util.Iterator;
-import java.util.stream.Stream;
 
 public record ChunkedInferenceError(Exception exception) implements ChunkedInference {
 
     @Override
-    public Iterator<Chunk> chunksAsMatchedTextAndByteReference(XContent xcontent) {
-        return Stream.of(exception).map(e -> new Chunk(e.getMessage(), new TextOffset(0, 0), BytesArray.EMPTY)).iterator();
+    public Iterator<Chunk> chunksAsByteReference(XContent xcontent) {
+        return Collections.emptyIterator();
     }
 }

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);
     }
 
     /**

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

@@ -175,17 +175,15 @@ 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(matchedText, offset, toBytesReference(xcontent, weightedTokens));
+            return new ChunkedInference.Chunk(offset, toBytesReference(xcontent, weightedTokens));
         }
 
         private static BytesReference toBytesReference(XContent xContent, List<WeightedToken> tokens) throws IOException {

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

@@ -187,18 +187,18 @@ 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(matchedText, offset, toBytesReference(xcontent, embedding));
+            return new ChunkedInference.Chunk(offset, toBytesReference(xcontent, embedding));
         }
 
         private static BytesReference toBytesReference(XContent xContent, byte[] value) throws IOException {

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

@@ -221,15 +221,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);
         }
     }
 
-    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(matchedText, offset, toBytesReference(xcontent, embedding));
+            return new ChunkedInference.Chunk(offset, toBytesReference(xcontent, embedding));
         }
 
         /**

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())
                             )
                         )

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())))
                     )
                 );
             }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/filter/ShardBulkInferenceActionFilter.java

@@ -606,7 +606,7 @@ static IndexRequest getIndexRequestOrNull(DocWriteRequest<?> docWriteRequest) {
 
     private static class EmptyChunkedInference implements ChunkedInference {
         @Override
-        public Iterator<Chunk> chunksAsMatchedTextAndByteReference(XContent xcontent) {
+        public Iterator<Chunk> chunksAsByteReference(XContent xcontent) {
             return Collections.emptyIterator();
         }
     }