TRY-198: fix collect

m10-payments · Dec 27, 2023 · 2fcaa8e · 2fcaa8e
1 parent 116f38d
commit 2fcaa8e
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/java/api/src/main/java/com/exadel/frs/core/trainservice/cache/EmbeddingCollection.java b/java/api/src/main/java/com/exadel/frs/core/trainservice/cache/EmbeddingCollection.java
@@ -10,7 +10,6 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -28,11 +27,10 @@ public static EmbeddingCollection from(final Stream<EnhancedEmbeddingProjection>
         // we copy vector here just in case
         var newMap = stream.map(e -> Map.entry(EmbeddingProjection.from(e), MatrixUtils.createRealVector(e.getEmbeddingData()))
         ).collect(
-                Collectors.toMap(
+                Collectors.toConcurrentMap(
                         Map.Entry::getKey,
                         Map.Entry::getValue,
-                        (v1, v2) -> v1,
-                        ConcurrentHashMap::new
+                        (v1, v2) -> v1
                 )
         );
         return new EmbeddingCollection(newMap);