Add chunking to test inference service

Author: kderusso

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

@@ -13,6 +13,8 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.inference.ChunkingSettings;
+import org.elasticsearch.inference.ChunkingStrategy;
 import org.elasticsearch.inference.InferenceService;
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.ModelConfigurations;
@@ -22,9 +24,13 @@
 import org.elasticsearch.inference.TaskSettings;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xpack.inference.chunking.WordBoundaryChunker;
+import org.elasticsearch.xpack.inference.chunking.WordBoundaryChunkingSettings;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Random;
 
@@ -105,6 +111,33 @@ 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) {
+        if (chunkingSettings == null) {
+            return 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()));
+                }
+            }
+        } else {
+            // Won't implement till we need it
+            throw new UnsupportedOperationException("Test inference service only supports word chunking strategies");
+        }
+
+        return chunkedInputs;
+    }
+
     public static class TestServiceModel extends Model {
 
         public TestServiceModel(

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

@@ -155,7 +155,7 @@ public void chunkedInfer(
             switch (model.getConfigurations().getTaskType()) {
                 case ANY, TEXT_EMBEDDING -> {
                     ServiceSettings modelServiceSettings = model.getServiceSettings();
-                    listener.onResponse(makeChunkedResults(input, modelServiceSettings.dimensions()));
+                    listener.onResponse(makeChunkedResults(input, modelServiceSettings.dimensions(), chunkingSettings));
                 }
                 default -> listener.onFailure(
                     new ElasticsearchStatusException(
@@ -175,6 +175,11 @@ 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);
 

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

@@ -45,6 +45,7 @@
 import java.util.Map;
 
 public class TestSparseInferenceServiceExtension implements InferenceServiceExtension {
+
     @Override
     public List<Factory> getInferenceServiceFactories() {
         return List.of(TestInferenceService::new);
@@ -143,7 +144,7 @@ public void chunkedInfer(
             ActionListener<List<ChunkedInference>> listener
         ) {
             switch (model.getConfigurations().getTaskType()) {
-                case ANY, SPARSE_EMBEDDING -> listener.onResponse(makeChunkedResults(input));
+                case ANY, SPARSE_EMBEDDING -> listener.onResponse(makeChunkedResults(input, chunkingSettings));
                 default -> listener.onFailure(
                     new ElasticsearchStatusException(
                         TaskType.unsupportedTaskTypeErrorMsg(model.getConfigurations().getTaskType(), name()),
@@ -165,6 +166,11 @@ 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) {
             List<ChunkedInference> results = new ArrayList<>();
             for (int i = 0; i < input.size(); i++) {

x-pack/plugin/inference/src/main/java/module-info.java

@@ -42,6 +42,7 @@
     exports org.elasticsearch.xpack.inference.services;
     exports org.elasticsearch.xpack.inference;
     exports org.elasticsearch.xpack.inference.action.task;
+    exports org.elasticsearch.xpack.inference.chunking;
     exports org.elasticsearch.xpack.inference.telemetry;
 
     provides org.elasticsearch.features.FeatureSpecification with org.elasticsearch.xpack.inference.InferenceFeatures;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/WordBoundaryChunkingSettings.java

@@ -58,6 +58,14 @@ public Map<String, Object> asMap() {
         return map;
     }
 
+    public int maxChunkSize() {
+        return maxChunkSize;
+    }
+
+    public int overlap() {
+        return overlap;
+    }
+
     public static WordBoundaryChunkingSettings fromMap(Map<String, Object> map) {
         ValidationException validationException = new ValidationException();
 

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

@@ -37,7 +37,7 @@ setup:
 
   - do:
       indices.create:
-        index: default-chunking
+        index: default-chunking-sparse
         body:
           mappings:
             properties:
@@ -49,7 +49,7 @@ setup:
 
   - do:
       indices.create:
-        index: custom-chunking
+        index: default-chunking-dense
         body:
           mappings:
             properties:
@@ -58,14 +58,26 @@ setup:
               inference_field:
                 type: semantic_text
                 inference_id: dense-inference-id
+
+  - do:
+      indices.create:
+        index: custom-chunking-sparse
+        body:
+          mappings:
+            properties:
+              keyword_field:
+                type: keyword
+              inference_field:
+                type: semantic_text
+                inference_id: sparse-inference-id
                 chunking_settings:
                   strategy: word
                   max_chunk_size: 10
-                  overlap: 5
+                  overlap: 1
 
   - do:
       index:
-        index: default-chunking
+        index: default-chunking-sparse
         id: doc_1
         body:
           keyword_field: "default sentence chunking"
@@ -74,7 +86,7 @@ setup:
 
   - do:
       index:
-        index: custom-chunking
+        index: custom-chunking-sparse
         id: doc_2
         body:
           keyword_field: "custom word chunking"
@@ -86,13 +98,13 @@ setup:
 
   - do:
       indices.get_mapping:
-        index: default-chunking
+        index: default-chunking-sparse
 
   - is_false: default-chunking.mappings.properties.inference_field.chunking_settings
 
   - do:
       indices.get_mapping:
-        index: custom-chunking
+        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 }
@@ -103,7 +115,7 @@ setup:
 
   - do:
       search:
-        index: default-chunking
+        index: default-chunking-sparse
         body:
           query:
             semantic:
@@ -122,7 +134,7 @@ setup:
 
   - do:
       search:
-        index: custom-chunking
+        index: custom-chunking-sparse
         body:
           query:
             semantic:
@@ -138,5 +150,5 @@ setup:
   - 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: " Lucene internally and enjoys all the features it provides." }
+  - match: { hits.hits.0.highlight.inference_field.1: " the features it provides." }