From 7933154462d99f570707edca02b60d6c7831af17 Mon Sep 17 00:00:00 2001
From: Amith Koujalgi <amith.koujalgi@razorthink.com>
Date: Sun, 12 Nov 2023 21:14:54 +0530
Subject: [PATCH] Added `generateEmbeddings` API

---
 README.md                                          | 14 ++++++++++++++
 .../amithkoujalgi/ollama4j/core/OllamaAPI.java     |  4 ++--
 .../ollama4j/core/models/EmbeddingResponse.java    |  8 +++++---
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index c370adc..fa47e84 100644
--- a/README.md
+++ b/README.md
@@ -175,6 +175,20 @@ public class Main {
 }
 ```
 
+#### Generate embeddings:
+
+```java
+public class Main {
+    public static void main(String[] args) throws Exception {
+        String host = "http://localhost:11434/";
+        OllamaAPI ollamaAPI = new OllamaAPI(host);
+        ollamaAPI.setVerbose(true);
+        List<Double> doubleList = ollamaAPI.generateEmbeddings(OllamaModelType.LLAMA2, "Here is an article about llamas...");
+        doubleList.forEach(System.out::println);
+    }
+}
+```
+
 #### Ask a question to the model with ollama4j
 
 ##### Using sync API:
diff --git a/src/main/java/io/github/amithkoujalgi/ollama4j/core/OllamaAPI.java b/src/main/java/io/github/amithkoujalgi/ollama4j/core/OllamaAPI.java
index 920ed4e..4893ea1 100644
--- a/src/main/java/io/github/amithkoujalgi/ollama4j/core/OllamaAPI.java
+++ b/src/main/java/io/github/amithkoujalgi/ollama4j/core/OllamaAPI.java
@@ -293,12 +293,12 @@ public class OllamaAPI {
      *
      * @param model name of model to generate embeddings from
      * @param prompt text to generate embeddings for
-     * @return embeddings as double[]
+     * @return embeddings
      * @throws IOException
      * @throws ParseException
      * @throws OllamaBaseException
      */
-    public double[] generateEmbeddings(String model, String prompt) throws IOException, ParseException, OllamaBaseException {
+    public List<Double> generateEmbeddings(String model, String prompt) throws IOException, ParseException, OllamaBaseException {
         String url = this.host + "/api/embeddings";
         String jsonData = String.format("{\"model\": \"%s\", \"prompt\": \"%s\"}", model, prompt);
         final HttpPost httpPost = new HttpPost(url);
diff --git a/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/EmbeddingResponse.java b/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/EmbeddingResponse.java
index a4f6b82..807688d 100644
--- a/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/EmbeddingResponse.java
+++ b/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/EmbeddingResponse.java
@@ -2,18 +2,20 @@ package io.github.amithkoujalgi.ollama4j.core.models;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 
+import java.util.List;
+
 public class EmbeddingResponse {
     @JsonProperty("embedding")
-    private double[] embedding;
+    private List<Double> embedding;
 
     public EmbeddingResponse() {
     }
 
-    public double[] getEmbedding() {
+    public List<Double> getEmbedding() {
         return embedding;
     }
 
-    public void setEmbedding(double[] embedding) {
+    public void setEmbedding(List<Double> embedding) {
         this.embedding = embedding;
     }
 }