Skip to content
This repository was archived by the owner on Jul 4, 2025. It is now read-only.

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

Merged
merged 2 commits into from
Jan 24, 2024
Merged

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

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9218cd9
feat: Add input as vector
hiro-v Jan 23, 2024
5f9a0a4
fix: Final update for embedding to support both single and vector of …
hiro-v Jan 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 36 additions & 31 deletions controllers/llamaCPP.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand All @@ -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,
Expand Down Expand Up @@ -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;
Expand Down