Fix chunking in test inference service and yaml tests

Author: kderusso

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/AbstractTestInferenceService.java

@@ -29,6 +29,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -111,25 +112,24 @@ public void start(Model model, TimeValue timeout, ActionListener<Boolean> listen
     @Override
     public void close() throws IOException {}
 
-    protected List<String> chunkInputs(List<String> input, ChunkingSettings chunkingSettings) {
+    protected List<String> chunkInputs(String input, ChunkingSettings chunkingSettings) {
         if (chunkingSettings == null) {
-            return input;
+            return Collections.singletonList(input);
         }
         List<String> chunkedInputs = new ArrayList<>();
         ChunkingStrategy chunkingStrategy = chunkingSettings.getChunkingStrategy();
         if (chunkingStrategy == ChunkingStrategy.WORD) {
             WordBoundaryChunker chunker = new WordBoundaryChunker();
-            for (String inputString : input) {
-                WordBoundaryChunkingSettings wordBoundaryChunkingSettings = (WordBoundaryChunkingSettings) chunkingSettings;
-                List<WordBoundaryChunker.ChunkOffset> offsets = chunker.chunk(
-                    inputString,
-                    wordBoundaryChunkingSettings.maxChunkSize(),
-                    wordBoundaryChunkingSettings.overlap()
-                );
-                for (WordBoundaryChunker.ChunkOffset offset : offsets) {
-                    chunkedInputs.add(inputString.substring(offset.start(), offset.end()));
-                }
+            WordBoundaryChunkingSettings wordBoundaryChunkingSettings = (WordBoundaryChunkingSettings) chunkingSettings;
+            List<WordBoundaryChunker.ChunkOffset> offsets = chunker.chunk(
+                input,
+                wordBoundaryChunkingSettings.maxChunkSize(),
+                wordBoundaryChunkingSettings.overlap()
+            );
+            for (WordBoundaryChunker.ChunkOffset offset : offsets) {
+                chunkedInputs.add(input.substring(offset.start(), offset.end()));
             }
+
         } else {
             // Won't implement till we need it
             throw new UnsupportedOperationException("Test inference service only supports word chunking strategies");

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestDenseInferenceServiceExtension.java

@@ -36,7 +36,9 @@
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.core.inference.results.ChunkedInferenceEmbedding;
+import org.elasticsearch.xpack.core.inference.results.SparseEmbeddingResults;
 import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
+import org.elasticsearch.xpack.core.ml.search.WeightedToken;
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
@@ -175,27 +177,28 @@ private TextEmbeddingFloatResults makeResults(List<String> input, int dimensions
             return new TextEmbeddingFloatResults(embeddings);
         }
 
-        private List<ChunkedInference> makeChunkedResults(List<String> input, int dimensions, ChunkingSettings chunkingSettings) {
-            List<String> chunkedInputs = chunkInputs(input, chunkingSettings);
-            return makeChunkedResults(chunkedInputs, dimensions);
-        }
-
-        private List<ChunkedInference> makeChunkedResults(List<String> input, int dimensions) {
-            TextEmbeddingFloatResults nonChunkedResults = makeResults(input, dimensions);
-
-            var results = new ArrayList<ChunkedInference>();
-            for (int i = 0; i < input.size(); i++) {
-                results.add(
-                    new ChunkedInferenceEmbedding(
-                        List.of(
-                            new TextEmbeddingFloatResults.Chunk(
-                                nonChunkedResults.embeddings().get(i).values(),
-                                input.get(i),
-                                new ChunkedInference.TextOffset(0, input.get(i).length())
-                            )
+        private List<ChunkedInference> makeChunkedResults(List<String> inputs, int dimensions, ChunkingSettings chunkingSettings) {
+
+            List<ChunkedInference> results = new ArrayList<>();
+            for (int i = 0; i < inputs.size(); i++) {
+                String input = inputs.get(i);
+                TextEmbeddingFloatResults nonChunkedResults = makeResults(inputs, dimensions);
+                List<String> chunkedInput = chunkInputs(input, chunkingSettings);
+                List<TextEmbeddingFloatResults.Chunk> chunks = new ArrayList<>();
+                int offset = 0;
+                for (String c : chunkedInput) {
+                    offset = input.indexOf(c, offset);
+                    int endOffset = offset + c.length();
+                    chunks.add(
+                        new TextEmbeddingFloatResults.Chunk(
+                            nonChunkedResults.embeddings().get(i).values(),
+                            c,
+                            new ChunkedInference.TextOffset(offset, endOffset)
                         )
-                    )
-                );
+                    );
+                }
+                ChunkedInferenceEmbedding chunkedInferenceEmbedding = new ChunkedInferenceEmbedding(chunks);
+                results.add(chunkedInferenceEmbedding);
             }
             return results;
         }

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestSparseInferenceServiceExtension.java

@@ -34,6 +34,7 @@
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.core.inference.results.ChunkedInferenceEmbedding;
+import org.elasticsearch.xpack.core.inference.results.EmbeddingResults;
 import org.elasticsearch.xpack.core.inference.results.SparseEmbeddingResults;
 import org.elasticsearch.xpack.core.ml.search.WeightedToken;
 
@@ -166,29 +167,24 @@ private SparseEmbeddingResults makeResults(List<String> input) {
             return new SparseEmbeddingResults(embeddings);
         }
 
-        private List<ChunkedInference> makeChunkedResults(List<String> input, ChunkingSettings chunkingSettings) {
-            List<String> chunkedInputs = chunkInputs(input, chunkingSettings);
-            return makeChunkedResults(chunkedInputs);
-        }
-
-        private List<ChunkedInference> makeChunkedResults(List<String> input) {
+        private List<ChunkedInference> makeChunkedResults(List<String> inputs, ChunkingSettings chunkingSettings) {
             List<ChunkedInference> results = new ArrayList<>();
-            for (int i = 0; i < input.size(); i++) {
+            for (int i = 0; i < inputs.size(); i++) {
+                String input = inputs.get(i);
                 var tokens = new ArrayList<WeightedToken>();
                 for (int j = 0; j < 5; j++) {
-                    tokens.add(new WeightedToken("feature_" + j, generateEmbedding(input.get(i), j)));
+                    tokens.add(new WeightedToken("feature_" + j, generateEmbedding(input, j)));
                 }
-                results.add(
-                    new ChunkedInferenceEmbedding(
-                        List.of(
-                            new SparseEmbeddingResults.Chunk(
-                                tokens,
-                                input.get(i),
-                                new ChunkedInference.TextOffset(0, input.get(i).length())
-                            )
-                        )
-                    )
-                );
+                List<String> chunkedInput = chunkInputs(input, chunkingSettings);
+                List<SparseEmbeddingResults.Chunk> chunks = new ArrayList<>();
+                int offset = 0;
+                for (String c : chunkedInput) {
+                    offset = input.indexOf(c, offset);
+                    int endOffset = offset + c.length();
+                    chunks.add(new SparseEmbeddingResults.Chunk(tokens, c, new ChunkedInference.TextOffset(offset, endOffset)));
+                }
+                ChunkedInferenceEmbedding chunkedInferenceEmbedding = new ChunkedInferenceEmbedding(chunks);
+                results.add(chunkedInferenceEmbedding);
             }
             return results;
         }

x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/25_semantic_text_field_mapping_chunking.yml

@@ -75,6 +75,22 @@ setup:
                   max_chunk_size: 10
                   overlap: 1
 
+  - do:
+      indices.create:
+        index: custom-chunking-dense
+        body:
+          mappings:
+            properties:
+              keyword_field:
+                type: keyword
+              inference_field:
+                type: semantic_text
+                inference_id: dense-inference-id
+                chunking_settings:
+                  strategy: word
+                  max_chunk_size: 10
+                  overlap: 1
+
   - do:
       index:
         index: default-chunking-sparse
@@ -93,25 +109,57 @@ setup:
           inference_field: "Elasticsearch is an open source, distributed, RESTful, search engine which is built on top of Lucene internally and enjoys all the features it provides."
         refresh: true
 
+  - do:
+      index:
+        index: default-chunking-dense
+        id: doc_3
+        body:
+          keyword_field: "default sentence chunking"
+          inference_field: "Elasticsearch is an open source, distributed, RESTful, search engine which is built on top of Lucene internally and enjoys all the features it provides."
+        refresh: true
+
+  - do:
+      index:
+        index: custom-chunking-dense
+        id: doc_4
+        body:
+          keyword_field: "custom word chunking"
+          inference_field: "Elasticsearch is an open source, distributed, RESTful, search engine which is built on top of Lucene internally and enjoys all the features it provides."
+        refresh: true
+
 ---
 "We return chunking configurations with mappings":
 
   - do:
       indices.get_mapping:
         index: default-chunking-sparse
 
-  - is_false: default-chunking.mappings.properties.inference_field.chunking_settings
+  - is_false: default-chunking-sparse.mappings.properties.inference_field.chunking_settings
 
   - do:
       indices.get_mapping:
         index: custom-chunking-sparse
 
-  - match: { "custom-chunking.mappings.properties.inference_field.chunking_settings.strategy": "word" }
-  - match: { "custom-chunking.mappings.properties.inference_field.chunking_settings.max_chunk_size": 10 }
-  - match: { "custom-chunking.mappings.properties.inference_field.chunking_settings.overlap": 5 }
+  - match: { "custom-chunking-sparse.mappings.properties.inference_field.chunking_settings.strategy": "word" }
+  - match: { "custom-chunking-sparse.mappings.properties.inference_field.chunking_settings.max_chunk_size": 10 }
+  - match: { "custom-chunking-sparse.mappings.properties.inference_field.chunking_settings.overlap": 1 }
+
+  - do:
+      indices.get_mapping:
+        index: default-chunking-dense
+
+  - is_false: default-chunking-dense.mappings.properties.inference_field.chunking_settings
+
+  - do:
+      indices.get_mapping:
+        index: custom-chunking-dense
+
+  - match: { "custom-chunking-dense.mappings.properties.inference_field.chunking_settings.strategy": "word" }
+  - match: { "custom-chunking-dense.mappings.properties.inference_field.chunking_settings.max_chunk_size": 10 }
+  - match: { "custom-chunking-dense.mappings.properties.inference_field.chunking_settings.overlap": 1 }
 
 ---
-"We return different chunks based on configured chunking overrides or model defaults":
+"We return different chunks based on configured chunking overrides or model defaults for sparse embeddings":
 
   - do:
       search:
@@ -149,6 +197,47 @@ setup:
   - match: { hits.total.value: 1 }
   - match: { hits.hits.0._id: "doc_2" }
   - length: { hits.hits.0.highlight.inference_field: 2 }
-  - match: { hits.hits.0.highlight.inference_field.0: "Elasticsearch is an open source, distributed, RESTful, search engine which is built on top of Lucene internally and enjoys all" }
-  - match: { hits.hits.0.highlight.inference_field.1: " the features it provides." }
+  - match: { hits.hits.0.highlight.inference_field.0: "Elasticsearch is an open source, distributed, RESTful, search engine which" }
+  - match: { hits.hits.0.highlight.inference_field.1: " which is built on top of Lucene internally and enjoys" }
 
+---
+"We return different chunks based on configured chunking overrides or model defaults for dense embeddings":
+
+  - do:
+      search:
+        index: default-chunking-dense
+        body:
+          query:
+            semantic:
+              field: "inference_field"
+              query: "What is Elasticsearch?"
+          highlight:
+            fields:
+              inference_field:
+                type: "semantic"
+                number_of_fragments: 2
+
+  - match: { hits.total.value: 1 }
+  - match: { hits.hits.0._id: "doc_3" }
+  - length: { hits.hits.0.highlight.inference_field: 1 }
+  - match: { hits.hits.0.highlight.inference_field.0: "Elasticsearch is an open source, distributed, RESTful, search engine which is built on top of Lucene internally and enjoys all the features it provides." }
+
+  - do:
+      search:
+        index: custom-chunking-dense
+        body:
+          query:
+            semantic:
+              field: "inference_field"
+              query: "What is Elasticsearch?"
+          highlight:
+            fields:
+              inference_field:
+                type: "semantic"
+                number_of_fragments: 2
+
+  - match: { hits.total.value: 1 }
+  - match: { hits.hits.0._id: "doc_4" }
+  - length: { hits.hits.0.highlight.inference_field: 2 }
+  - match: { hits.hits.0.highlight.inference_field.0: "Elasticsearch is an open source, distributed, RESTful, search engine which" }
+  - match: { hits.hits.0.highlight.inference_field.1: " which is built on top of Lucene internally and enjoys" }