Skip to content

output normalize embedding in '/v1/embeddings' #5956

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Mar 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
15 changes: 15 additions & 0 deletions common/common.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1852,3 +1852,18 @@ void dump_kv_cache_view_seqs(const llama_kv_cache_view & view, int row_size) {

printf("\n=== Done dumping\n");
}

void llama_embd_normalize(const float * inp, float * out, int n) {
double sum = 0.0;
for (int i = 0; i < n; i++) {
sum += inp[i] * inp[i];
}
sum = sqrt(sum);

const float norm = sum > 0.0 ? 1.0f / sum : 0.0f;

for (int i = 0; i < n; i++) {
out[i] = inp[i] * norm;
}
}

7 changes: 7 additions & 0 deletions common/common.h
Original file line number Diff line number Diff line change
Expand Up @@ -260,3 +260,10 @@ void dump_kv_cache_view(const llama_kv_cache_view & view, int row_size = 80);

// Dump the KV cache view showing individual sequences in each cell (long output).
void dump_kv_cache_view_seqs(const llama_kv_cache_view & view, int row_size = 40);

//
// Embedding utils
//

void llama_embd_normalize(const float * inp, float * out, int n);

14 changes: 1 addition & 13 deletions examples/embedding/embedding.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -23,17 +23,6 @@ static void batch_add_seq(llama_batch & batch, const std::vector<int32_t> & toke
}
}

static void normalize(const float * vec, float * out, int n) {
float norm = 0;
for (int i = 0; i < n; i++) {
norm += vec[i] * vec[i];
}
norm = sqrt(norm);
for (int i = 0; i < n; i++) {
out[i] = vec[i] / norm;
}
}

static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd) {
// clear previous kv_cache values (irrelevant for embeddings)
llama_kv_cache_clear(ctx);
Expand All @@ -44,7 +33,6 @@ static void batch_decode(llama_context * ctx, llama_batch & batch, float * outpu
fprintf(stderr, "%s : failed to decode\n", __func__);
}

// normalize on copy
for (int i = 0; i < batch.n_tokens; i++) {
if (!batch.logits[i]) {
continue;
Expand All @@ -61,7 +49,7 @@ static void batch_decode(llama_context * ctx, llama_batch & batch, float * outpu
}

float * out = output + batch.seq_id[i][0] * n_embd;
normalize(embd, out, n_embd);
llama_embd_normalize(embd, out, n_embd);
}
}

Expand Down
8 changes: 7 additions & 1 deletion examples/server/server.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1318,6 +1318,8 @@ struct server_context {

const int n_embd = llama_n_embd(model);

std::vector<float> embd_res(n_embd, 0.0f);

for (int i = 0; i < batch.n_tokens; ++i) {
if (!batch.logits[i] || batch.seq_id[i][0] != slot.id + 1) {
continue;
Expand All @@ -1341,8 +1343,10 @@ struct server_context {
continue;
}

llama_embd_normalize(embd, embd_res.data(), n_embd);

res.data = json {
{"embedding", std::vector<float>(embd, embd + n_embd)},
{"embedding", embd_res},
};
}

Expand Down Expand Up @@ -3345,6 +3349,8 @@ int main(int argc, char ** argv) {
// get the result
server_task_result result = ctx_server.queue_results.recv(id_task);
ctx_server.queue_results.remove_waiting_task_id(id_task);

// append to the responses
responses.push_back(result.data);
}

Expand Down