Change implementation of asMap() to avoid extra nesting in inference …

…results
carlosdelest · Feb 14, 2024 · bd4e19e · bd4e19e
1 parent fc76918
commit bd4e19e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/.../main/java/org/elasticsearch/xpack/core/inference/results/LegacyTextEmbeddingResults.java b/.../main/java/org/elasticsearch/xpack/core/inference/results/LegacyTextEmbeddingResults.java
@@ -80,15 +80,15 @@ public String getResultsField() {
     @Override
     public Map<String, Object> asMap() {
         Map<String, Object> map = new LinkedHashMap<>();
-        map.put(getResultsField(), embeddings.stream().map(Embedding::asMap).collect(Collectors.toList()));
+        map.put(getResultsField(), embeddings.stream().flatMap(v -> v.values.stream()).collect(Collectors.toList()));
 
         return map;
     }
 
     @Override
     public Map<String, Object> asMap(String outputField) {
         Map<String, Object> map = new LinkedHashMap<>();
-        map.put(outputField, embeddings.stream().map(Embedding::asMap).collect(Collectors.toList()));
+        map.put(outputField, embeddings.stream().flatMap(v -> v.values.stream()).collect(Collectors.toList()));
 
         return map;
     }

diff --git a/...re/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingResults.java b/...re/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingResults.java
@@ -21,7 +21,6 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -129,10 +128,7 @@ public List<? extends InferenceResults> transformToLegacyFormat() {
     }
 
     public Map<String, Object> asMap() {
-        Map<String, Object> map = new LinkedHashMap<>();
-        map.put(TEXT_EMBEDDING, embeddings.stream().map(Embedding::asMap).collect(Collectors.toList()));
-
-        return map;
+        return Map.of(TEXT_EMBEDDING, embeddings.stream().flatMap(v -> v.values.stream()).collect(Collectors.toList()));
     }
 
     public record Embedding(List<Float> values) implements Writeable, ToXContentObject, EmbeddingInt {