Semantic Text - Use offsets from chunked inference response (#118893)

Author: Mikep86

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextField.java

@@ -425,15 +425,16 @@ public static Chunk toSemanticTextFieldChunk(
         ChunkedInference.Chunk chunk,
         boolean useInferenceMetadataFieldsFormat
     ) {
-        // TODO: Use offsets from ChunkedInferenceServiceResults
-        // TODO: When using legacy semantic text format, build chunk text from offsets
-        assert chunk.matchedText() != null; // TODO: Remove once offsets are available from chunk
-        int startOffset = useInferenceMetadataFieldsFormat ? input.indexOf(chunk.matchedText()) + offsetAdjustment : -1;
-        return new Chunk(
-            useInferenceMetadataFieldsFormat ? null : chunk.matchedText(),
-            useInferenceMetadataFieldsFormat ? startOffset : -1,
-            useInferenceMetadataFieldsFormat ? startOffset + chunk.matchedText().length() : -1,
-            chunk.bytesReference()
-        );
+        String text = null;
+        int startOffset = -1;
+        int endOffset = -1;
+        if (useInferenceMetadataFieldsFormat) {
+            startOffset = chunk.textOffset().start() + offsetAdjustment;
+            endOffset = chunk.textOffset().end() + offsetAdjustment;
+        } else {
+            text = input.substring(chunk.textOffset().start(), chunk.textOffset().end());
+        }
+
+        return new Chunk(text, startOffset, endOffset, chunk.bytesReference());
     }
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/filter/ShardBulkInferenceActionFilterTests.java

@@ -416,7 +416,7 @@ private static BulkItemRequest[] randomBulkItemRequest(
             } else {
                 Map<String, List<String>> inputTextMap = Map.of(field, List.of(inputText));
                 semanticTextField = randomSemanticText(indexVersion, field, model, List.of(inputText), requestContentType);
-                model.putResult(inputText, toChunkedResult(inputTextMap, semanticTextField));
+                model.putResult(inputText, toChunkedResult(indexVersion, inputTextMap, semanticTextField));
             }
 
             if (useInferenceMetadataFieldsFormat) {

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldTests.java

@@ -286,7 +286,11 @@ public static Object randomSemanticTextInput() {
         }
     }
 
-    public static ChunkedInference toChunkedResult(Map<String, List<String>> matchedTextMap, SemanticTextField field) throws IOException {
+    public static ChunkedInference toChunkedResult(
+        IndexVersion indexVersion,
+        Map<String, List<String>> matchedTextMap,
+        SemanticTextField field
+    ) {
         switch (field.inference().modelSettings().taskType()) {
             case SPARSE_EMBEDDING -> {
                 List<ChunkedInferenceEmbeddingSparse.SparseEmbeddingChunk> chunks = new ArrayList<>();
@@ -297,14 +301,10 @@ public static ChunkedInference toChunkedResult(Map<String, List<String>> matched
 
                     ListIterator<String> matchedTextIt = entryFieldMatchedText.listIterator();
                     for (var chunk : entryChunks) {
+                        String matchedText = matchedTextIt.next();
+                        ChunkedInference.TextOffset offset = createOffset(indexVersion, chunk, matchedText);
                         var tokens = parseWeightedTokens(chunk.rawEmbeddings(), field.contentType());
-                        chunks.add(
-                            new ChunkedInferenceEmbeddingSparse.SparseEmbeddingChunk(
-                                tokens,
-                                matchedTextIt.next(),
-                                new ChunkedInference.TextOffset(chunk.startOffset(), chunk.endOffset())
-                            )
-                        );
+                        chunks.add(new ChunkedInferenceEmbeddingSparse.SparseEmbeddingChunk(tokens, matchedText, offset));
                     }
                 }
                 return new ChunkedInferenceEmbeddingSparse(chunks);
@@ -318,6 +318,8 @@ public static ChunkedInference toChunkedResult(Map<String, List<String>> matched
 
                     ListIterator<String> matchedTextIt = entryFieldMatchedText.listIterator();
                     for (var chunk : entryChunks) {
+                        String matchedText = matchedTextIt.next();
+                        ChunkedInference.TextOffset offset = createOffset(indexVersion, chunk, matchedText);
                         double[] values = parseDenseVector(
                             chunk.rawEmbeddings(),
                             field.inference().modelSettings().dimensions(),
@@ -326,8 +328,8 @@ public static ChunkedInference toChunkedResult(Map<String, List<String>> matched
                         chunks.add(
                             new ChunkedInferenceEmbeddingFloat.FloatEmbeddingChunk(
                                 FloatConversionUtils.floatArrayOf(values),
-                                matchedTextIt.next(),
-                                new ChunkedInference.TextOffset(chunk.startOffset(), chunk.endOffset())
+                                matchedText,
+                                offset
                             )
                         );
                     }
@@ -353,6 +355,24 @@ private static List<String> validateAndGetMatchedTextForField(
         return fieldMatchedText;
     }
 
+    /**
+     * Create a {@link ChunkedInference.TextOffset} instance with valid offset values. When using the legacy semantic text format, the
+     * offset values are not written to {@link SemanticTextField.Chunk}, so we cannot read them from there. Instead, use the knowledge that
+     * the matched text corresponds to one complete input value (i.e. one input value -> one chunk) to calculate the offset values.
+     *
+     * @param indexVersion The index version
+     * @param chunk The chunk to get/calculate offset values for
+     * @param matchedText The matched text to calculate offset values for
+     * @return A {@link ChunkedInference.TextOffset} instance with valid offset values
+     */
+    private static ChunkedInference.TextOffset createOffset(IndexVersion indexVersion, SemanticTextField.Chunk chunk, String matchedText) {
+        final boolean useInferenceMetadataFields = InferenceMetadataFieldsMapper.isEnabled(indexVersion);
+        final int startOffset = useInferenceMetadataFields ? chunk.startOffset() : 0;
+        final int endOffset = useInferenceMetadataFields ? chunk.endOffset() : matchedText.length();
+
+        return new ChunkedInference.TextOffset(startOffset, endOffset);
+    }
+
     private static double[] parseDenseVector(BytesReference value, int numDims, XContentType contentType) {
         try (XContentParser parser = XContentHelper.createParserNotCompressed(XContentParserConfiguration.EMPTY, value, contentType)) {
             parser.nextToken();