pgvector-java/examples/sparse/src/main/java/com/example/Example.java at master · pgvector/pgvector-java · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
// good resources
// https://opensearch.org/blog/improving-document-retrieval-with-sparse-semantic-encoders/
// https://huggingface.co/opensearch-project/opensearch-neural-sparse-encoding-v1
//
// run with
// text-embeddings-router --model-id opensearch-project/opensearch-neural-sparse-encoding-v1 --pooling splade

packagecom.example;

importjava.io.IOException;
importjava.net.URI;
importjava.net.http.HttpClient;
importjava.net.http.HttpRequest;
importjava.net.http.HttpRequest.BodyPublishers;
importjava.net.http.HttpResponse;
importjava.net.http.HttpResponse.BodyHandlers;
importjava.sql.Connection;
importjava.sql.DriverManager;
importjava.sql.PreparedStatement;
importjava.sql.ResultSet;
importjava.sql.SQLException;
importjava.sql.Statement;
importjava.util.ArrayList;
importjava.util.HashMap;
importjava.util.List;
importjava.util.Map;
importcom.fasterxml.jackson.databind.ObjectMapper;
importcom.fasterxml.jackson.databind.JsonNode;
importcom.fasterxml.jackson.databind.node.ObjectNode;
importcom.pgvector.PGsparsevec;
importcom.pgvector.PGvector;

publicclassExample {
publicstaticvoidmain(String[] args) throwsIOException, InterruptedException, SQLException {
Connectionconn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/pgvector_example");

StatementsetupStmt = conn.createStatement();
setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector");
setupStmt.executeUpdate("DROP TABLE IF EXISTS documents");

PGvector.addVectorType(conn);

StatementcreateStmt = conn.createStatement();
createStmt.executeUpdate("CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding sparsevec(30522))");

String[] input = {
"The dog is barking",
"The cat is purring",
"The bear is growling"
 };
List<Map<Integer, Float>> embeddings = embed(input);

for (inti = 0; i < input.length; i++) {
PreparedStatementinsertStmt = conn.prepareStatement("INSERT INTO documents (content, embedding) VALUES (?, ?)");
insertStmt.setString(1, input[i]);
insertStmt.setObject(2, newPGsparsevec(embeddings.get(i), 30522));
insertStmt.executeUpdate();
 }

Stringquery = "forest";
Map<Integer, Float> queryEmbedding = embed(newString[] { query }).get(0);
PreparedStatementneighborStmt = conn.prepareStatement("SELECT content FROM documents ORDER BY embedding <#> ? LIMIT 5");
neighborStmt.setObject(1, newPGsparsevec(queryEmbedding, 30522));
ResultSetrs = neighborStmt.executeQuery();
while (rs.next()) {
System.out.println(rs.getString("content"));
 }

conn.close();
 }

privatestaticList<Map<Integer, Float>> embed(String[] inputs) throwsIOException, InterruptedException {
ObjectMappermapper = newObjectMapper();
ObjectNoderoot = mapper.createObjectNode();
for (Stringv : inputs) {
root.withArray("inputs").add(v);
 }
Stringjson = mapper.writeValueAsString(root);

HttpClientclient = HttpClient.newHttpClient();
HttpRequestrequest = HttpRequest.newBuilder()
 .uri(URI.create("http://localhost:3000/embed_sparse"))
 .header("Content-Type", "application/json")
 .POST(BodyPublishers.ofString(json))
 .build();
HttpResponse<String> response = client.send(request, BodyHandlers.ofString());

List<Map<Integer, Float>> embeddings = newArrayList<>();
for (JsonNoden : mapper.readTree(response.body())) {
Map<Integer, Float> embedding = newHashMap<Integer, Float>();
for (JsonNodev : n) {
intindex = v.get("index").asInt();
floatvalue = (float) v.get("value").asDouble();
embedding.put(Integer.valueOf(index), Float.valueOf(value));
 }
embeddings.add(embedding);
 }
returnembeddings;
 }
}