Skip to content

Commit 112a1c5

Browse files
tazarovtazarov
authored
feat: EF move to embeddings sub-package (#49)
Closes #48
1 parent df49126 commit 112a1c5

24 files changed

Lines changed: 174 additions & 168 deletions

‎README.md‎

Lines changed: 87 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -73,72 +73,72 @@ package tech.amikos;
7373

7474
import tech.amikos.chromadb.*;
7575
import tech.amikos.chromadb.Collection;
76+
import tech.amikos.chromadb.embeddings.DefaultEmbeddingFunction;
7677

7778
import java.util.*;
7879

7980
public class Main {
80-
public static void main(String[] args) {
81-
try {
82-
Client client = new Client(System.getenv("CHROMA_URL"));
83-
client.reset();
84-
EmbeddingFunction ef = new DefaultEmbeddingFunction();
85-
Collection collection = client.createCollection("test-collection", null, true, ef);
86-
List<Map<String, String>> metadata = new ArrayList<>();
87-
metadata.add(new HashMap<String, String>() {{
88-
put("type", "scientist");
89-
}});
90-
metadata.add(new HashMap<String, String>() {{
91-
put("type", "spy");
92-
}});
93-
collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
94-
Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
95-
System.out.println(qr);
96-
} catch (Exception e) {
97-
System.out.println(e);
98-
}
81+
public static void main(String[] args) {
82+
try {
83+
Client client = new Client(System.getenv("CHROMA_URL"));
84+
client.reset();
85+
EmbeddingFunction ef = new DefaultEmbeddingFunction();
86+
Collection collection = client.createCollection("test-collection", null, true, ef);
87+
List<Map<String, String>> metadata = new ArrayList<>();
88+
metadata.add(new HashMap<String, String>() {{
89+
put("type", "scientist");
90+
}});
91+
metadata.add(new HashMap<String, String>() {{
92+
put("type", "spy");
93+
}});
94+
collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
95+
Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
96+
System.out.println(qr);
97+
} catch (Exception e) {
98+
System.out.println(e);
9999
}
100+
}
100101
}
101102
```
102103

103104
### Example OpenAI Embedding Function
104105

105-
In this example we rely on `tech.amikos.chromadb.OpenAIEmbeddingFunction` to generate embeddings for our documents.
106+
In this example we rely on `tech.amikos.chromadb.embeddings.openai.OpenAIEmbeddingFunction` to generate embeddings for our documents.
106107

107108
| **Important**: Ensure you have `OPENAI_API_KEY` environment variable set
108109

109110
```java
110111
package tech.amikos;
111112

112-
import com.google.gson.internal.LinkedTreeMap;
113113
import tech.amikos.chromadb.Client;
114114
import tech.amikos.chromadb.Collection;
115115
import tech.amikos.chromadb.EmbeddingFunction;
116-
import tech.amikos.chromadb.OpenAIEmbeddingFunction;
116+
import tech.amikos.chromadb.embeddings.openai.OpenAIEmbeddingFunction;
117117

118118
import java.util.*;
119119

120120
public class Main {
121-
public static void main(String[] args) {
122-
try {
123-
Client client = new Client(System.getenv("CHROMA_URL"));
124-
String apiKey = System.getenv("OPENAI_API_KEY");
125-
EmbeddingFunction ef = new OpenAIEmbeddingFunction(apiKey, "text-embedding-3-small");
126-
Collection collection = client.createCollection("test-collection", null, true, ef);
127-
List<Map<String, String>> metadata = new ArrayList<>();
128-
metadata.add(new HashMap<String, String>() {{
129-
put("type", "scientist");
130-
}});
131-
metadata.add(new HashMap<String, String>() {{
132-
put("type", "spy");
133-
}});
134-
collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
135-
Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
136-
System.out.println(qr);
137-
} catch (Exception e) {
138-
e.printStackTrace();
139-
System.out.println(e);
140-
}
121+
public static void main(String[] args) {
122+
try {
123+
Client client = new Client(System.getenv("CHROMA_URL"));
124+
String apiKey = System.getenv("OPENAI_API_KEY");
125+
EmbeddingFunction ef = new OpenAIEmbeddingFunction(apiKey, "text-embedding-3-small");
126+
Collection collection = client.createCollection("test-collection", null, true, ef);
127+
List<Map<String, String>> metadata = new ArrayList<>();
128+
metadata.add(new HashMap<String, String>() {{
129+
put("type", "scientist");
130+
}});
131+
metadata.add(new HashMap<String, String>() {{
132+
put("type", "spy");
133+
}});
134+
collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
135+
Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
136+
System.out.println(qr);
137+
} catch (Exception e) {
138+
e.printStackTrace();
139+
System.out.println(e);
141140
}
141+
}
142142
}
143143
```
144144

@@ -174,7 +174,7 @@ curl http://localhost:11434/api/embeddings -d '{\n "model": "llama2",\n "promp
174174

175175
### Example Cohere Embedding Function
176176

177-
In this example we rely on `tech.amikos.chromadb.CohereEmbeddingFunction` to generate embeddings for our documents.
177+
In this example we rely on `tech.amikos.chromadb.embeddings.cohere.CohereEmbeddingFunction` to generate embeddings for our documents.
178178

179179
| **Important**: Ensure you have `COHERE_API_KEY` environment variable set
180180

@@ -183,32 +183,33 @@ package tech.amikos;
183183

184184
import tech.amikos.chromadb.*;
185185
import tech.amikos.chromadb.Collection;
186+
import tech.amikos.chromadb.embeddings.cohere.CohereEmbeddingFunction;
186187

187188
import java.util.*;
188189

189190
public class Main {
190-
public static void main(String[] args) {
191-
try {
192-
Client client = new Client(System.getenv("CHROMA_URL"));
193-
client.reset();
194-
String apiKey = System.getenv("COHERE_API_KEY");
195-
EmbeddingFunction ef = new CohereEmbeddingFunction(apiKey);
196-
Collection collection = client.createCollection("test-collection", null, true, ef);
197-
List<Map<String, String>> metadata = new ArrayList<>();
198-
metadata.add(new HashMap<String, String>() {{
199-
put("type", "scientist");
200-
}});
201-
metadata.add(new HashMap<String, String>() {{
202-
put("type", "spy");
203-
}});
204-
collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
205-
Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
206-
System.out.println(qr);
207-
} catch (Exception e) {
208-
e.printStackTrace();
209-
System.out.println(e);
210-
}
191+
public static void main(String[] args) {
192+
try {
193+
Client client = new Client(System.getenv("CHROMA_URL"));
194+
client.reset();
195+
String apiKey = System.getenv("COHERE_API_KEY");
196+
EmbeddingFunction ef = new CohereEmbeddingFunction(apiKey);
197+
Collection collection = client.createCollection("test-collection", null, true, ef);
198+
List<Map<String, String>> metadata = new ArrayList<>();
199+
metadata.add(new HashMap<String, String>() {{
200+
put("type", "scientist");
201+
}});
202+
metadata.add(new HashMap<String, String>() {{
203+
put("type", "spy");
204+
}});
205+
collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
206+
Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
207+
System.out.println(qr);
208+
} catch (Exception e) {
209+
e.printStackTrace();
210+
System.out.println(e);
211211
}
212+
}
212213
}
213214
```
214215

@@ -220,7 +221,7 @@ The above should output:
220221

221222
### Example Hugging Face Sentence Transformers Embedding Function
222223

223-
In this example we rely on `tech.amikos.chromadb.HuggingFaceEmbeddingFunction` to generate embeddings for our documents.
224+
In this example we rely on `tech.amikos.chromadb.embeddings.hf.HuggingFaceEmbeddingFunction` to generate embeddings for our documents.
224225

225226
| **Important**: Ensure you have `HF_API_KEY` environment variable set
226227

@@ -229,31 +230,32 @@ package tech.amikos;
229230

230231
import tech.amikos.chromadb.*;
231232
import tech.amikos.chromadb.Collection;
233+
import tech.amikos.chromadb.embeddings.hf.HuggingFaceEmbeddingFunction;
232234

233235
import java.util.*;
234236

235237
public class Main {
236-
public static void main(String[] args) {
237-
try {
238-
Client client = new Client(System.getenv("CHROMA_URL"));
239-
client.reset();
240-
String apiKey = System.getenv("HF_API_KEY");
241-
EmbeddingFunction ef = new HuggingFaceEmbeddingFunction(apiKey);
242-
Collection collection = client.createCollection("test-collection", null, true, ef);
243-
List<Map<String, String>> metadata = new ArrayList<>();
244-
metadata.add(new HashMap<String, String>() {{
245-
put("type", "scientist");
246-
}});
247-
metadata.add(new HashMap<String, String>() {{
248-
put("type", "spy");
249-
}});
250-
collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
251-
Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
252-
System.out.println(qr);
253-
} catch (Exception e) {
254-
System.out.println(e);
255-
}
238+
public static void main(String[] args) {
239+
try {
240+
Client client = new Client(System.getenv("CHROMA_URL"));
241+
client.reset();
242+
String apiKey = System.getenv("HF_API_KEY");
243+
EmbeddingFunction ef = new HuggingFaceEmbeddingFunction(apiKey);
244+
Collection collection = client.createCollection("test-collection", null, true, ef);
245+
List<Map<String, String>> metadata = new ArrayList<>();
246+
metadata.add(new HashMap<String, String>() {{
247+
put("type", "scientist");
248+
}});
249+
metadata.add(new HashMap<String, String>() {{
250+
put("type", "spy");
251+
}});
252+
collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
253+
Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
254+
System.out.println(qr);
255+
} catch (Exception e) {
256+
System.out.println(e);
256257
}
258+
}
257259
}
258260
```
259261

‎pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -228,7 +228,7 @@
228228
<!-- <additionalProperty>ignoreImportMappings=false</additionalProperty>-->
229229
<!-- </additionalProperties>-->
230230
<importMappings>
231-
<importMapping>CreateEmbeddingRequest=tech.amikos.openai.CreateEmbeddingRequest
231+
<importMapping>CreateEmbeddingRequest=tech.amikos.chromadb.embeddings.openai.CreateEmbeddingRequest
232232
</importMapping>
233233
</importMappings>
234234
<generateApiTests>false</generateApiTests>

src/main/java/tech/amikos/chromadb/DefaultEmbeddingFunction.java renamed to src/main/java/tech/amikos/chromadb/embeddings/DefaultEmbeddingFunction.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package tech.amikos.chromadb;
1+
package tech.amikos.chromadb.embeddings;
22

33
import ai.djl.huggingface.tokenizers.Encoding;
44
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
@@ -11,6 +11,8 @@
1111
import org.nd4j.linalg.api.ops.impl.transforms.clip.ClipByValue;
1212
import org.nd4j.linalg.factory.Nd4j;
1313
import org.nd4j.shade.guava.primitives.Floats;
14+
import tech.amikos.chromadb.EFException;
15+
import tech.amikos.chromadb.EmbeddingFunction;
1416

1517
import java.io.*;
1618
import java.net.URL;

src/main/java/tech/amikos/cohere/CohereClient.java renamed to src/main/java/tech/amikos/chromadb/embeddings/cohere/CohereClient.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package tech.amikos.cohere;
1+
package tech.amikos.chromadb.embeddings.cohere;
22

33
import com.google.gson.Gson;
44
import okhttp3.*;

src/main/java/tech/amikos/chromadb/CohereEmbeddingFunction.java renamed to src/main/java/tech/amikos/chromadb/embeddings/cohere/CohereEmbeddingFunction.java

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,8 @@
1-
package tech.amikos.chromadb;
1+
package tech.amikos.chromadb.embeddings.cohere;
22

3-
import tech.amikos.cohere.CohereClient;
4-
import tech.amikos.cohere.CreateEmbeddingRequest;
5-
import tech.amikos.cohere.CreateEmbeddingResponse;
3+
import tech.amikos.chromadb.EmbeddingFunction;
64

7-
import java.util.Arrays;
85
import java.util.List;
9-
import java.util.stream.Collectors;
106

117
public class CohereEmbeddingFunction implements EmbeddingFunction {
128

src/main/java/tech/amikos/cohere/CreateEmbeddingRequest.java renamed to src/main/java/tech/amikos/chromadb/embeddings/cohere/CreateEmbeddingRequest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package tech.amikos.cohere;
1+
package tech.amikos.chromadb.embeddings.cohere;
22

33
import com.google.gson.*;
44
import com.google.gson.annotations.SerializedName;

src/main/java/tech/amikos/cohere/CreateEmbeddingResponse.java renamed to src/main/java/tech/amikos/chromadb/embeddings/cohere/CreateEmbeddingResponse.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package tech.amikos.cohere;
1+
package tech.amikos.chromadb.embeddings.cohere;
22

33
import com.google.gson.Gson;
44
import com.google.gson.annotations.SerializedName;

src/main/java/tech/amikos/hf/CreateEmbeddingRequest.java renamed to src/main/java/tech/amikos/chromadb/embeddings/hf/CreateEmbeddingRequest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package tech.amikos.hf;
1+
package tech.amikos.chromadb.embeddings.hf;
22

33
import com.google.gson.Gson;
44
import com.google.gson.annotations.SerializedName;

src/main/java/tech/amikos/hf/CreateEmbeddingResponse.java renamed to src/main/java/tech/amikos/chromadb/embeddings/hf/CreateEmbeddingResponse.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package tech.amikos.hf;
1+
package tech.amikos.chromadb.embeddings.hf;
22

33
import com.google.gson.Gson;
44

src/main/java/tech/amikos/hf/HuggingFaceClient.java renamed to src/main/java/tech/amikos/chromadb/embeddings/hf/HuggingFaceClient.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package tech.amikos.hf;
1+
package tech.amikos.chromadb.embeddings.hf;
22

33
import com.google.gson.Gson;
44
import okhttp3.*;

0 commit comments

Comments
 (0)