-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #37 from frazik-main/24-fix-localmemory-funcionality
24 fix localmemory funcionality
- Loading branch information
Showing
7 changed files
with
25 additions
and
72 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
48 changes: 9 additions & 39 deletions
48
lib/src/main/java/com/frazik/instructgpt/memory/LocalMemory.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,67 +1,37 @@ | ||
| package com.frazik.instructgpt.memory; | ||
|
|
||
| import com.frazik.instructgpt.embedding.EmbeddingProvider; | ||
| import lombok.extern.slf4j.Slf4j; | ||
| import org.nd4j.linalg.api.ndarray.INDArray; | ||
| import org.nd4j.linalg.factory.Nd4j; | ||
| import lombok.extern.slf4j.Slf4j;; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.Arrays; | ||
| import java.util.List; | ||
|
|
||
|
|
||
| @Slf4j | ||
| public class LocalMemory extends Memory { | ||
|
|
||
| private final List<String> docs; | ||
| private INDArray embs; | ||
| private final EmbeddingProvider embeddingProvider; | ||
|
|
||
| public LocalMemory(EmbeddingProvider embeddingProvider) { | ||
| public LocalMemory() { | ||
| super(); | ||
| this.docs = new ArrayList<>(); | ||
| this.embs = null; | ||
| this.embeddingProvider = embeddingProvider; | ||
| } | ||
|
|
||
| @Override | ||
| public void add(String doc, String key) { | ||
| if (key == null) { | ||
| key = doc; | ||
| } | ||
| double[] embeddings = this.embeddingProvider.get(key); | ||
| INDArray emb = Nd4j.create(embeddings); | ||
|
|
||
| if (this.embs == null) { | ||
| this.embs = Nd4j.expandDims(emb, 0); | ||
| } else { | ||
| this.embs = Nd4j.concat(0, this.embs, Nd4j.expandDims(emb, 0)); | ||
| } | ||
| this.docs.add(doc); | ||
| public void add(String doc) { | ||
| this.docs.add(0, doc); | ||
| } | ||
|
|
||
| @Override | ||
| public List<String> get(String query, int k) { | ||
| if (this.embs == null) { | ||
| return new ArrayList<>(); | ||
| } | ||
| double[] embeddings = embeddingProvider.get(query); | ||
| INDArray scores; | ||
| try (INDArray emb = Nd4j.create(embeddings)) { | ||
|
|
||
| scores = this.embs.mmul(emb); | ||
| } | ||
| int[] idxs = Nd4j.argMax(scores, 0).toIntVector(); | ||
| String[] results = new String[k]; | ||
| for (int i = 0; i < k; i++) { | ||
| results[i] = this.docs.get(idxs[i]); | ||
| public List<String> get(int k) { | ||
| // get last k docs, or all docs if k > docs.size() | ||
| if (k > this.docs.size()) { | ||
| return this.docs; | ||
| } | ||
| return Arrays.asList(results); | ||
| return this.docs.subList(0, k); | ||
| } | ||
| @Override | ||
| public void clear() { | ||
| this.docs.clear(); | ||
| this.embs = null; | ||
| } | ||
|
|
||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,7 +1,7 @@ | ||
| package com.frazik.instructgpt.memory; | ||
| import java.util.List; | ||
| public abstract class Memory { | ||
| public abstract void add(String doc, String key); | ||
| public abstract List<String> get(String query, int k); | ||
| public abstract void add(String doc); | ||
| public abstract List<String> get(int k); | ||
| public abstract void clear(); | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters