Adding support for JinaAI late chunking

Author: dan-rubinstein

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

@@ -19,4 +19,8 @@ public interface TaskSettings extends ToXContentObject, VersionedNamedWriteable
     boolean isEmpty();
 
     TaskSettings updatedTaskSettings(Map<String, Object> newSettings);
+
+    default Boolean isLateChunkingEnabled() {
+        return null;
+    }
 }

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

@@ -87,6 +87,15 @@ public EmbeddingRequestChunker(
         List<ChunkInferenceInput> inputs,
         int maxNumberOfInputsPerBatch,
         ChunkingSettings defaultChunkingSettings
+    ) {
+        this(inputs, maxNumberOfInputsPerBatch, false, defaultChunkingSettings);
+    }
+
+    public EmbeddingRequestChunker(
+        List<ChunkInferenceInput> inputs,
+        int maxNumberOfInputsPerBatch,
+        Boolean isLateChunkingEnabled,
+        ChunkingSettings defaultChunkingSettings
     ) {
         this.resultEmbeddings = new ArrayList<>(inputs.size());
         this.resultOffsetStarts = new ArrayList<>(inputs.size());
@@ -133,13 +142,24 @@ public EmbeddingRequestChunker(
             }
         }
 
-        AtomicInteger counter = new AtomicInteger();
-        this.batchRequests = allRequests.stream()
-            .collect(Collectors.groupingBy(it -> counter.getAndIncrement() / maxNumberOfInputsPerBatch))
-            .values()
-            .stream()
-            .map(BatchRequest::new)
-            .toList();
+        if (isLateChunkingEnabled != null && isLateChunkingEnabled) {
+            // This must be true for late chunking cases otherwise we can't pass all chunks in a single request
+            assert (maxNumberOfInputsPerBatch >= MAX_CHUNKS);
+            this.batchRequests = allRequests.stream()
+                .collect(Collectors.groupingBy(Request::inputIndex))
+                .values()
+                .stream()
+                .map(BatchRequest::new)
+                .toList();
+        } else {
+            AtomicInteger counter = new AtomicInteger();
+            this.batchRequests = allRequests.stream()
+                .collect(Collectors.groupingBy(it -> counter.getAndIncrement() / maxNumberOfInputsPerBatch))
+                .values()
+                .stream()
+                .map(BatchRequest::new)
+                .toList();
+        }
     }
 
     /**

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/jinaai/JinaAIService.java

@@ -282,6 +282,7 @@ protected void doChunkedInfer(
         List<EmbeddingRequestChunker.BatchRequestAndListener> batchedRequests = new EmbeddingRequestChunker<>(
             inputs.getInputs(),
             EMBEDDING_MAX_BATCH_SIZE,
+            jinaaiModel.getTaskSettings().isLateChunkingEnabled(),
             jinaaiModel.getConfigurations().getChunkingSettings()
         ).batchRequestsWithListeners(listener);
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/jinaai/embeddings/JinaAIEmbeddingsTaskSettings.java

@@ -24,6 +24,7 @@
 import java.util.Objects;
 
 import static org.elasticsearch.inference.InputType.invalidInputTypeMessage;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalBoolean;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalEnum;
 import static org.elasticsearch.xpack.inference.services.jinaai.JinaAIService.VALID_INPUT_TYPE_VALUES;
 
@@ -36,6 +37,7 @@ public class JinaAIEmbeddingsTaskSettings implements TaskSettings {
     public static final String NAME = "jinaai_embeddings_task_settings";
     public static final JinaAIEmbeddingsTaskSettings EMPTY_SETTINGS = new JinaAIEmbeddingsTaskSettings((InputType) null);
     static final String INPUT_TYPE = "input_type";
+    static final String LATE_CHUNKING = "late_chunking";
 
     public static JinaAIEmbeddingsTaskSettings fromMap(Map<String, Object> map) {
         if (map == null || map.isEmpty()) {
@@ -53,11 +55,13 @@ public static JinaAIEmbeddingsTaskSettings fromMap(Map<String, Object> map) {
             validationException
         );
 
+        Boolean lateChunking = extractOptionalBoolean(map, LATE_CHUNKING, validationException);
+
         if (validationException.validationErrors().isEmpty() == false) {
             throw validationException;
         }
 
-        return new JinaAIEmbeddingsTaskSettings(inputType);
+        return new JinaAIEmbeddingsTaskSettings(inputType, lateChunking);
     }
 
     /**
@@ -77,7 +81,8 @@ public static JinaAIEmbeddingsTaskSettings of(
     ) {
         var inputTypeToUse = getValidInputType(originalSettings, requestTaskSettings);
 
-        return new JinaAIEmbeddingsTaskSettings(inputTypeToUse);
+        return new JinaAIEmbeddingsTaskSettings(inputTypeToUse, requestTaskSettings.lateChunking);
+        // TODO: Check the above
     }
 
     private static InputType getValidInputType(
@@ -94,14 +99,22 @@ private static InputType getValidInputType(
     }
 
     private final InputType inputType;
+    private final Boolean lateChunking;
 
     public JinaAIEmbeddingsTaskSettings(StreamInput in) throws IOException {
-        this(in.readOptionalEnum(InputType.class));
+        this(in.readOptionalEnum(InputType.class), in.readOptionalBoolean());
+    }
+
+    public JinaAIEmbeddingsTaskSettings(@Nullable InputType inputType, @Nullable Boolean lateChunking) {
+        validateInputType(inputType);
+        this.inputType = inputType;
+        this.lateChunking = lateChunking;
     }
 
     public JinaAIEmbeddingsTaskSettings(@Nullable InputType inputType) {
         validateInputType(inputType);
         this.inputType = inputType;
+        this.lateChunking = null;
     }
 
     private static void validateInputType(InputType inputType) {
@@ -124,6 +137,11 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field(INPUT_TYPE, inputType);
         }
 
+        // TODO: Add a transport version
+        if (lateChunking != null) {
+            builder.field(LATE_CHUNKING, lateChunking);
+        }
+
         builder.endObject();
         return builder;
     }
@@ -132,6 +150,11 @@ public InputType getInputType() {
         return inputType;
     }
 
+    @Override
+    public Boolean isLateChunkingEnabled() {
+        return lateChunking;
+    }
+
     @Override
     public String getWriteableName() {
         return NAME;
@@ -145,19 +168,20 @@ public TransportVersion getMinimalSupportedVersion() {
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalEnum(inputType);
+        out.writeOptionalBoolean(lateChunking);
     }
 
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         JinaAIEmbeddingsTaskSettings that = (JinaAIEmbeddingsTaskSettings) o;
-        return Objects.equals(inputType, that.inputType);
+        return Objects.equals(inputType, that.inputType) && lateChunking == that.lateChunking;
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(inputType);
+        return Objects.hash(inputType, lateChunking);
     }
 
     @Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/jinaai/request/JinaAIEmbeddingsRequestEntity.java

@@ -34,6 +34,7 @@ public record JinaAIEmbeddingsRequestEntity(
     private static final String CLASSIFICATION = "classification";
     private static final String INPUT_FIELD = "input";
     private static final String MODEL_FIELD = "model";
+    private static final String LATE_CHUNKING = "late_chunking";
     public static final String TASK_TYPE_FIELD = "task";
     static final String EMBEDDING_TYPE_FIELD = "embedding_type";
 
@@ -49,6 +50,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         builder.field(INPUT_FIELD, input);
         builder.field(MODEL_FIELD, model);
 
+        if (taskSettings.isLateChunkingEnabled() != null) {
+            builder.field(LATE_CHUNKING, taskSettings.isLateChunkingEnabled());
+        }
+
         if (embeddingType != null) {
             builder.field(EMBEDDING_TYPE_FIELD, embeddingType.toRequestString());
         }