topdog/supersonic
1
0
Fork 0
mirror of https://github.com/tencentmusic/supersonic.git synced 2025-12-12 04:27:39 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

(improvement)(chat) Change the storage of InMemoryEmbeddingStore entity to a Set for deduplication. (#564)

Browse Source
This commit is contained in:
lexluo09
2023-12-22 14:49:53 +08:00
committed by GitHub
parent 090abbceed
commit 0cbfe473dd
1 changed files with 7 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
common/src/main/java/com/tencent/supersonic/common/util/embedding/InMemoryS2EmbeddingStore.java
View File

@@ -25,8 +25,9 @@ import java.util.Map;
import java.util.Map.Entry; import java.util.Map.Entry;
import java.util.Objects; import java.util.Objects;
import java.util.PriorityQueue; import java.util.PriorityQueue;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CopyOnWriteArraySet;
import java.util.stream.Collectors; import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.MapUtils; import org.apache.commons.collections4.MapUtils;
@@ -49,7 +50,7 @@ public class InMemoryS2EmbeddingStore implements S2EmbeddingStore {
try { try {
if (Files.exists(filePath)) { if (Files.exists(filePath)) {
embeddingStore = InMemoryEmbeddingStore.fromFile(filePath); embeddingStore = InMemoryEmbeddingStore.fromFile(filePath);
embeddingStore.entries = new CopyOnWriteArrayList<>(embeddingStore.entries); embeddingStore.entries = new CopyOnWriteArraySet<>(embeddingStore.entries);
log.info("embeddingStore reload from file:{}", filePath); log.info("embeddingStore reload from file:{}", filePath);
} }
} catch (Exception e) { } catch (Exception e) {
@@ -215,7 +216,7 @@ public class InMemoryS2EmbeddingStore implements S2EmbeddingStore {
} }
private static final InMemoryEmbeddingStoreJsonCodec CODEC = loadCodec(); private static final InMemoryEmbeddingStoreJsonCodec CODEC = loadCodec();
private List<InMemoryEmbeddingStore.Entry<Embedded>> entries = new CopyOnWriteArrayList<>(); private Set<Entry<Embedded>> entries = new CopyOnWriteArraySet<>();
@Override @Override
public String add(Embedding embedding) { public String add(Embedding embedding) {
@@ -237,23 +238,15 @@ public class InMemoryS2EmbeddingStore implements S2EmbeddingStore {
} }
public void add(String id, Embedding embedding, Embedded embedded) { public void add(String id, Embedding embedding, Embedded embedded) {
if (checkEmbeddingNotExists(embedding)) {
entries.add(new InMemoryEmbeddingStore.Entry<>(id, embedding, embedded)); entries.add(new InMemoryEmbeddingStore.Entry<>(id, embedding, embedded));
} }
}
private boolean checkEmbeddingNotExists(Embedding embedding) {
return entries.stream().noneMatch(entry -> entry.embedding.equals(embedding));
}
@Override @Override
public List<String> addAll(List<Embedding> embeddings) { public List<String> addAll(List<Embedding> embeddings) {
List<String> ids = new ArrayList<>(); List<String> ids = new ArrayList<>();
for (Embedding embedding : embeddings) { for (Embedding embedding : embeddings) {
if (checkEmbeddingNotExists(embedding)) {
ids.add(add(embedding)); ids.add(add(embedding));
} }
}
return ids; return ids;
} }
@@ -265,10 +258,8 @@ public class InMemoryS2EmbeddingStore implements S2EmbeddingStore {
List<String> ids = new ArrayList<>(); List<String> ids = new ArrayList<>();
for (int i = 0; i < embeddings.size(); i++) { for (int i = 0; i < embeddings.size(); i++) {
if (checkEmbeddingNotExists(embeddings.get(i))) {
ids.add(add(embeddings.get(i), embedded.get(i))); ids.add(add(embeddings.get(i), embedded.get(i)));
} }
}
return ids; return ids;
} }