feat: add embedding endpoint

Author: tikikun

config.json

@@ -8,6 +8,7 @@
 	"custom_config": {
 		"llama_model_path": "/Users/alandao/Documents/codes/nitro.cpp_temp/models/llama2_7b_chat_uncensored.Q4_0.gguf",
 		"ctx_len": 2048,
-		"ngl": 100
+		"ngl": 100,
+		"embedding":true
 	}
 }

controllers/llamaCPP.cc

@@ -3,8 +3,9 @@
 #include "nitro_utils.h"
 #include <chrono>
 #include <cstring>
-#include <thread>
+#include <drogon/HttpResponse.h>
 #include <regex>
+#include <thread>
 
 using namespace inferences;
 
@@ -198,3 +199,29 @@ void llamaCPP::chatCompletion(
                                                       "chat_completions.txt");
   callback(resp);
 }
+
+void llamaCPP::embedding(
+    const HttpRequestPtr &req,
+    std::function<void(const HttpResponsePtr &)> &&callback) {
+  auto lock = llama.lock();
+
+  const auto &jsonBody = req->getJsonObject();
+
+  llama.rewind();
+  llama_reset_timings(llama.ctx);
+  if (jsonBody->isMember("content") != 0) {
+    llama.prompt = (*jsonBody)["content"].asString();
+  } else {
+    llama.prompt = "";
+  }
+  llama.params.n_predict = 0;
+  llama.loadPrompt();
+  llama.beginCompletion();
+  llama.doCompletion();
+
+  const json data = format_embedding_response(llama);
+  auto resp = drogon::HttpResponse::newHttpResponse();
+  resp->setBody(data.dump());
+  resp->setContentTypeString("application/json");
+  callback(resp);
+}

controllers/llamaCPP.h

@@ -1308,14 +1308,15 @@ static void append_to_generated_text_from_generated_token_probs(
 using namespace drogon;
 
 namespace inferences {
- class llamaCPP : public drogon::HttpController<llamaCPP> {
+class llamaCPP : public drogon::HttpController<llamaCPP> {
 public:
   llamaCPP() {
     gpt_params params;
     auto conf = drogon::app().getCustomConfig();
     params.model = conf["llama_model_path"].asString();
     params.n_gpu_layers = conf["ngl"].asInt();
     params.n_ctx = conf["ctx_len"].asInt();
+    params.embedding = conf["embedding"].asBool();
 #ifdef GGML_USE_CUBLAS
     LOG_INFO << "Setting up GGML CUBLAS PARAMS";
     params.mul_mat_q = false;
@@ -1345,15 +1346,17 @@ namespace inferences {
   METHOD_LIST_BEGIN
   // list path definitions here;
   METHOD_ADD(llamaCPP::chatCompletion, "chat_completion");
+  METHOD_ADD(llamaCPP::embedding,"embedding");
   // PATH_ADD("/llama/chat_completion", Post);
   METHOD_LIST_END
   void chatCompletion(const HttpRequestPtr &req,
                       std::function<void(const HttpResponsePtr &)> &&callback);
+  void embedding(const HttpRequestPtr &req,
+                 std::function<void(const HttpResponsePtr &)> &&callback);
 
 private:
   llama_server_context llama;
   size_t sent_count = 0;
   size_t sent_token_probs_index = 0;
-};   
-}
-;
+};
+}; // namespace inferences