Merge pull request #372 from janhq/feat/batch_embedding

feat(embedding): Add input as vector for /embeddings

Author: hiro-v

controllers/llamaCPP.cc

@@ -21,13 +21,8 @@ std::shared_ptr<inferenceState> create_inference_state(llamaCPP *instance) {
 
 // --------------------------------------------
 
-std::string create_embedding_payload(const std::vector<float> &embedding,
+Json::Value create_embedding_payload(const std::vector<float> &embedding,
                                      int prompt_tokens) {
-  Json::Value root;
-
-  root["object"] = "list";
-
-  Json::Value dataArray(Json::arrayValue);
   Json::Value dataItem;
 
   dataItem["object"] = "embedding";
@@ -39,20 +34,7 @@ std::string create_embedding_payload(const std::vector<float> &embedding,
   dataItem["embedding"] = embeddingArray;
   dataItem["index"] = 0;
 
-  dataArray.append(dataItem);
-  root["data"] = dataArray;
-
-  root["model"] = "_";
-
-  Json::Value usage;
-  usage["prompt_tokens"] = prompt_tokens;
-  usage["total_tokens"] = prompt_tokens; // Assuming total tokens equals prompt
-                                         // tokens in this context
-  root["usage"] = usage;
-
-  Json::StreamWriterBuilder writer;
-  writer["indentation"] = ""; // Compact output
-  return Json::writeString(writer, root);
+  return dataItem;
 }
 
 std::string create_full_return_json(const std::string &id,
@@ -406,19 +388,42 @@ void llamaCPP::embedding(
     std::function<void(const HttpResponsePtr &)> &&callback) {
   const auto &jsonBody = req->getJsonObject();
 
-  json prompt;
-  if (jsonBody->isMember("input") != 0) {
-    prompt = (*jsonBody)["input"].asString();
-  } else {
-    prompt = "";
+  Json::Value responseData(Json::arrayValue);
+
+  if (jsonBody->isMember("input")) {
+    const Json::Value &input = (*jsonBody)["input"];
+    if (input.isString()) {
+      // Process the single string input
+      const int task_id = llama.request_completion(
+          {{"prompt", input.asString()}, {"n_predict", 0}}, false, true, -1);
+      task_result result = llama.next_result(task_id);
+      std::vector<float> embedding_result = result.result_json["embedding"];
+      responseData.append(create_embedding_payload(embedding_result, 0));
+    } else if (input.isArray()) {
+      // Process each element in the array input
+      for (const auto &elem : input) {
+        if (elem.isString()) {
+          const int task_id = llama.request_completion(
+              {{"prompt", elem.asString()}, {"n_predict", 0}}, false, true, -1);
+          task_result result = llama.next_result(task_id);
+          std::vector<float> embedding_result = result.result_json["embedding"];
+          responseData.append(create_embedding_payload(embedding_result, 0));
+        }
+      }
+    }
   }
-  const int task_id = llama.request_completion(
-      {{"prompt", prompt}, {"n_predict", 0}}, false, true, -1);
-  task_result result = llama.next_result(task_id);
-  std::vector<float> embedding_result = result.result_json["embedding"];
+
   auto resp = nitro_utils::nitroHttpResponse();
-  std::string embedding_resp = create_embedding_payload(embedding_result, 0);
-  resp->setBody(embedding_resp);
+  Json::Value root;
+  root["data"] = responseData;
+  root["model"] = "_";
+  root["object"] = "list";
+  Json::Value usage;
+  usage["prompt_tokens"] = 0;
+  usage["total_tokens"] = 0;
+  root["usage"] = usage;
+
+  resp->setBody(Json::writeString(Json::StreamWriterBuilder(), root));
   resp->setContentTypeString("application/json");
   callback(resp);
   return;