Improved examples [skip ci]

Author: ankane

examples/cohere/src/main/java/com/example/Example.java

@@ -42,7 +42,8 @@ public static void main(String[] args) throws IOException, InterruptedException,
             "The cat is purring",
             "The bear is growling"
         };
-        List<byte[]> embeddings = fetchEmbeddings(input, "search_document", apiKey);
+        List<byte[]> embeddings = embed(input, "search_document", apiKey);
+
         for (int i = 0; i < input.length; i++) {
             PreparedStatement insertStmt = conn.prepareStatement("INSERT INTO documents (content, embedding) VALUES (?, ?)");
             insertStmt.setString(1, input[i]);
@@ -51,7 +52,7 @@ public static void main(String[] args) throws IOException, InterruptedException,
         }
 
         String query = "forest";
-        byte[] queryEmbedding = fetchEmbeddings(new String[] {query}, "search_query", apiKey).get(0);
+        byte[] queryEmbedding = embed(new String[] {query}, "search_query", apiKey).get(0);
         PreparedStatement neighborStmt = conn.prepareStatement("SELECT * FROM documents ORDER BY embedding <~> ? LIMIT 5");
         neighborStmt.setObject(1, new PGbit(queryEmbedding));
         ResultSet rs = neighborStmt.executeQuery();
@@ -63,7 +64,7 @@ public static void main(String[] args) throws IOException, InterruptedException,
     }
 
     // https://docs.cohere.com/reference/embed
-    private static List<byte[]> fetchEmbeddings(String[] texts, String inputType, String apiKey) throws IOException, InterruptedException {
+    private static List<byte[]> embed(String[] texts, String inputType, String apiKey) throws IOException, InterruptedException {
         ObjectMapper mapper = new ObjectMapper();
         ObjectNode root = mapper.createObjectNode();
         for (String v : texts) {

examples/hybrid/src/main/java/com/example/Example.java

@@ -37,7 +37,7 @@ public static void main(String[] args) throws IOException, ModelException, SQLEx
             "The cat is purring",
             "The bear is growling"
         };
-        List<float[]> embeddings = generateEmbeddings(model, input);
+        List<float[]> embeddings = embed(model, input);
 
         for (int i = 0; i < input.length; i++) {
             PreparedStatement insertStmt = conn.prepareStatement("INSERT INTO documents (content, embedding) VALUES (?, ?)");
@@ -47,7 +47,7 @@ public static void main(String[] args) throws IOException, ModelException, SQLEx
         }
 
         String query = "growling bear";
-        float[] queryEmbedding = generateEmbeddings(model, new String[] {query}).get(0);
+        float[] queryEmbedding = embed(model, new String[] {query}).get(0);
         double k = 60;
 
         PreparedStatement queryStmt = conn.prepareStatement(HYBRID_SQL);
@@ -74,7 +74,7 @@ private static ZooModel<String, float[]> loadModel(String id) throws IOException
             .loadModel();
     }
 
-    private static List<float[]> generateEmbeddings(ZooModel<String, float[]> model, String[] input) throws TranslateException {
+    private static List<float[]> embed(ZooModel<String, float[]> model, String[] input) throws TranslateException {
         Predictor<String, float[]> predictor = model.newPredictor();
         List<float[]> embeddings = new ArrayList<>(input.length);
         for (String text : input) {

examples/openai/src/main/java/com/example/Example.java

@@ -44,7 +44,7 @@ public static void main(String[] args) throws IOException, InterruptedException,
             "The cat is purring",
             "The bear is growling"
         };
-        List<float[]> embeddings = fetchEmbeddings(input, apiKey);
+        List<float[]> embeddings = embed(input, apiKey);
 
         for (int i = 0; i < input.length; i++) {
             PreparedStatement insertStmt = conn.prepareStatement("INSERT INTO documents (content, embedding) VALUES (?, ?)");
@@ -53,10 +53,10 @@ public static void main(String[] args) throws IOException, InterruptedException,
             insertStmt.executeUpdate();
         }
 
-        long documentId = 2;
-        PreparedStatement neighborStmt = conn.prepareStatement("SELECT * FROM documents WHERE id != ? ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = ?) LIMIT 5");
-        neighborStmt.setObject(1, documentId);
-        neighborStmt.setObject(2, documentId);
+        String query = "forest";
+        float[] queryEmbedding = embed(new String[] {query}, apiKey).get(0);
+        PreparedStatement neighborStmt = conn.prepareStatement("SELECT content FROM documents ORDER BY embedding <=> ? LIMIT 5");
+        neighborStmt.setObject(1, new PGvector(queryEmbedding));
         ResultSet rs = neighborStmt.executeQuery();
         while (rs.next()) {
             System.out.println(rs.getString("content"));
@@ -65,7 +65,7 @@ public static void main(String[] args) throws IOException, InterruptedException,
         conn.close();
     }
 
-    private static List<float[]> fetchEmbeddings(String[] input, String apiKey) throws IOException, InterruptedException {
+    private static List<float[]> embed(String[] input, String apiKey) throws IOException, InterruptedException {
         ObjectMapper mapper = new ObjectMapper();
         ObjectNode root = mapper.createObjectNode();
         for (String v : input) {

examples/sparse/src/main/java/com/example/Example.java

@@ -48,7 +48,7 @@ public static void main(String[] args) throws IOException, InterruptedException,
             "The cat is purring",
             "The bear is growling"
         };
-        List<Map<Integer, Float>> embeddings = fetchEmbeddings(input);
+        List<Map<Integer, Float>> embeddings = embed(input);
 
         for (int i = 0; i < input.length; i++) {
             PreparedStatement insertStmt = conn.prepareStatement("INSERT INTO documents (content, embedding) VALUES (?, ?)");
@@ -58,7 +58,7 @@ public static void main(String[] args) throws IOException, InterruptedException,
         }
 
         String query = "forest";
-        Map<Integer, Float> queryEmbedding = fetchEmbeddings(new String[] { query }).get(0);
+        Map<Integer, Float> queryEmbedding = embed(new String[] { query }).get(0);
         PreparedStatement neighborStmt = conn.prepareStatement("SELECT content FROM documents ORDER BY embedding <#> ? LIMIT 5");
         neighborStmt.setObject(1, new PGsparsevec(queryEmbedding, 30522));
         ResultSet rs = neighborStmt.executeQuery();
@@ -69,7 +69,7 @@ public static void main(String[] args) throws IOException, InterruptedException,
         conn.close();
     }
 
-    private static List<Map<Integer, Float>> fetchEmbeddings(String[] inputs) throws IOException, InterruptedException {
+    private static List<Map<Integer, Float>> embed(String[] inputs) throws IOException, InterruptedException {
         ObjectMapper mapper = new ObjectMapper();
         ObjectNode root = mapper.createObjectNode();
         for (String v : inputs) {