pgvector-java/examples/djl/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
packagecom.example;

importjava.io.IOException;
importjava.sql.Connection;
importjava.sql.DriverManager;
importjava.sql.PreparedStatement;
importjava.sql.ResultSet;
importjava.sql.SQLException;
importjava.sql.Statement;
importjava.util.ArrayList;
importjava.util.List;
importai.djl.ModelException;
importai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
importai.djl.inference.Predictor;
importai.djl.repository.zoo.Criteria;
importai.djl.repository.zoo.ZooModel;
importai.djl.translate.TranslateException;
importcom.pgvector.PGvector;

publicclassExample {
publicstaticvoidmain(String[] args) throwsIOException, ModelException, SQLException, TranslateException {
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 vector(384))");

ZooModel<String, float[]> model = loadModel("sentence-transformers/all-MiniLM-L6-v2");

String[] input = {
"The dog is barking",
"The cat is purring",
"The bear is growling"
 };
List<float[]> embeddings = generateEmbeddings(model, 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, newPGvector(embeddings.get(i)));
insertStmt.executeUpdate();
 }

longdocumentId = 2;
PreparedStatementneighborStmt = conn.prepareStatement("SELECT * FROM documents WHERE id != ? ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = ?) LIMIT 5");
neighborStmt.setLong(1, documentId);
neighborStmt.setLong(2, documentId);
ResultSetrs = neighborStmt.executeQuery();
while (rs.next()) {
System.out.println(rs.getString("content"));
 }

conn.close();
 }

privatestaticZooModel<String, float[]> loadModel(Stringid) throwsIOException, ModelException {
returnCriteria.builder()
 .setTypes(String.class, float[].class)
 .optModelUrls("djl://ai.djl.huggingface.pytorch/" + id)
 .optEngine("PyTorch")
 .optTranslatorFactory(newTextEmbeddingTranslatorFactory())
 .build()
 .loadModel();
 }

privatestaticList<float[]> generateEmbeddings(ZooModel<String, float[]> model, String[] input) throwsTranslateException {
Predictor<String, float[]> predictor = model.newPredictor();
List<float[]> embeddings = newArrayList<>(input.length);
for (Stringtext : input) {
embeddings.add(predictor.predict(text));
 }
returnembeddings;
 }
}