Skip to content

Commit 888f99a

Browse files
authored
debug : include LLAMA_POOLING_TYPE_UNSPECIFIED in pooling check (ggml-org#18692)
* debug : include LLAMA_POOLING_TYPE_UNSPECIFIED in pooling check This commit updates the pooling check in the debug example to also include LLAMA_POOLING_TYPE_UNSPECIFIED and not just LLAMA_POOLING_TYPE_NONE. * debug : normalize both pooled and token embeddings This commit updates debug.cpp to normalize embeddings for both pooled and non-pooled outputs. For pooled embeddings, normalization is applied to the single vector, and for non-pooled embeddings, normalization is applied to each token embedding vector individually. The motivation for this is to enable non-pooled embeddings to be normalized which was not possible previously.
1 parent eda9e0f commit 888f99a

1 file changed

Lines changed: 34 additions & 16 deletions

File tree

examples/debug/debug.cpp

Lines changed: 34 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -57,11 +57,21 @@ struct callback_data {
5757
}
5858
};
5959

60+
static bool has_pooling(llama_context * ctx) {
61+
switch (llama_pooling_type(ctx)) {
62+
case LLAMA_POOLING_TYPE_NONE:
63+
case LLAMA_POOLING_TYPE_UNSPECIFIED:
64+
return false;
65+
default:
66+
return true;
67+
}
68+
}
69+
6070
struct output_data {
6171
float * data_ptr = nullptr;
6272
int data_size = 0;
6373
std::string type_suffix;
64-
std::vector<float> storage;
74+
std::vector<float> embd_norm;
6575
std::string prompt;
6676
std::vector<llama_token> tokens;
6777

@@ -73,24 +83,32 @@ struct output_data {
7383
prompt = params.prompt;
7484

7585
if (params.embedding) {
76-
const int n_embd = llama_model_n_embd_out(model);
77-
const bool pooling_enabled = llama_pooling_type(ctx) != LLAMA_POOLING_TYPE_NONE;
78-
const int n_embd_count = pooling_enabled ? 1 : tokens.size();
79-
const int n_embeddings = n_embd * n_embd_count;
80-
81-
float * embeddings;
82-
if (pooling_enabled) {
83-
embeddings = llama_get_embeddings_seq(ctx, 0);
84-
storage.resize(n_embeddings);
85-
common_embd_normalize(embeddings, storage.data(), n_embeddings, params.embd_normalize);
86-
embeddings = storage.data();
87-
} else {
88-
embeddings = llama_get_embeddings(ctx);
86+
const int n_embd = llama_model_n_embd_out(model);
87+
const bool pooling = has_pooling(ctx);
88+
const int n_embd_count = pooling ? 1 : tokens.size();
89+
const int n_floats = n_embd * n_embd_count;
90+
91+
float * embd_raw = pooling ? llama_get_embeddings_seq(ctx, 0) : llama_get_embeddings(ctx);
92+
if (embd_raw == nullptr) {
93+
throw std::runtime_error("failed to get embeddings from the model");
8994
}
9095

91-
data_ptr = embeddings;
92-
data_size = n_embeddings;
96+
LOG_DBG("pooling_enabled: %s\n", pooling ? "true" : "false");
97+
LOG_DBG("n_embd: %d\n", n_embd);
98+
LOG_DBG("n_floats: %d\n", n_floats);
99+
LOG_DBG("n_embd_count: %d\n", n_embd_count);
100+
101+
data_ptr = embd_raw;
102+
data_size = n_floats;
93103
type_suffix = "-embeddings";
104+
105+
if (params.embd_normalize >= 0) {
106+
embd_norm.resize(n_floats);
107+
for (int i = 0; i < n_embd_count; i++) {
108+
common_embd_normalize(embd_raw+i*n_embd, embd_norm.data()+i*n_embd, n_embd, params.embd_normalize);
109+
}
110+
data_ptr = embd_norm.data();
111+
}
94112
} else {
95113
const float * logits = llama_get_logits_ith(ctx, tokens.size() - 1);
96114
const int n_logits = llama_vocab_n_tokens(vocab);

0 commit comments

Comments
 (0)