1#include "arg.h"
2#include "common.h"
3#include "log.h"
4#include "llama.h"
5
6#include <ctime>
7#include <algorithm>
8
9#if defined(_MSC_VER)
10#pragma warning(disable: 4244 4267) // possible loss of data
11#endif
12
13static std::vector<std::string> split_lines(const std::string & s, const std::string & separator = "\n") {
14 std::vector<std::string> lines;
15 size_t start = 0;
16 size_t end = s.find(str: separator);
17
18 while (end != std::string::npos) {
19 lines.push_back(x: s.substr(pos: start, n: end - start));
20 start = end + separator.length();
21 end = s.find(str: separator, pos: start);
22 }
23
24 lines.push_back(x: s.substr(pos: start)); // Add the last part
25
26 return lines;
27}
28
29static void batch_add_seq(llama_batch & batch, const std::vector<int32_t> & tokens, llama_seq_id seq_id) {
30 size_t n_tokens = tokens.size();
31 for (size_t i = 0; i < n_tokens; i++) {
32 common_batch_add(batch, id: tokens[i], pos: i, seq_ids: { seq_id }, logits: true);
33 }
34}
35
36static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd, int embd_norm) {
37 const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);
38
39 // clear previous kv_cache values (irrelevant for embeddings)
40 llama_memory_clear(mem: llama_get_memory(ctx), data: true);
41
42 // run model
43 LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq);
44 if (llama_decode(ctx, batch) < 0) {
45 LOG_ERR("%s : failed to process\n", __func__);
46 }
47
48 for (int i = 0; i < batch.n_tokens; i++) {
49 if (!batch.logits[i]) {
50 continue;
51 }
52
53 const float * embd = nullptr;
54 int embd_pos = 0;
55
56 if (pooling_type == LLAMA_POOLING_TYPE_NONE) {
57 // try to get token embeddings
58 embd = llama_get_embeddings_ith(ctx, i);
59 embd_pos = i;
60 GGML_ASSERT(embd != NULL && "failed to get token embeddings");
61 } else {
62 // try to get sequence embeddings - supported only when pooling_type is not NONE
63 embd = llama_get_embeddings_seq(ctx, seq_id: batch.seq_id[i][0]);
64 embd_pos = batch.seq_id[i][0];
65 GGML_ASSERT(embd != NULL && "failed to get sequence embeddings");
66 }
67
68 float * out = output + embd_pos * n_embd;
69 common_embd_normalize(inp: embd, out, n: n_embd, embd_norm);
70 }
71}
72
73// plain, pipe-friendly output: one embedding per line
74static void print_raw_embeddings(const float * emb,
75 int n_embd_count,
76 int n_embd,
77 const llama_model * model,
78 enum llama_pooling_type pooling_type,
79 int embd_normalize) {
80 const uint32_t n_cls_out = llama_model_n_cls_out(model);
81 const bool is_rank = (pooling_type == LLAMA_POOLING_TYPE_RANK);
82 const int cols = is_rank ? std::min<int>(a: n_embd, b: (int) n_cls_out) : n_embd;
83
84 for (int j = 0; j < n_embd_count; ++j) {
85 for (int i = 0; i < cols; ++i) {
86 if (embd_normalize == 0) {
87 LOG("%1.0f%s", emb[j * n_embd + i], (i + 1 < cols ? " " : ""));
88 } else {
89 LOG("%1.7f%s", emb[j * n_embd + i], (i + 1 < cols ? " " : ""));
90 }
91 }
92 LOG("\n");
93 }
94}
95
96int main(int argc, char ** argv) {
97 common_params params;
98
99 if (!common_params_parse(argc, argv, params, ex: LLAMA_EXAMPLE_EMBEDDING)) {
100 return 1;
101 }
102
103 common_init();
104
105 params.embedding = true;
106
107 // if the number of prompts that would be encoded is known in advance, it's more efficient to specify the
108 // --parallel argument accordingly. for convenience, if not specified, we fallback to unified KV cache
109 // in order to support any number of prompts
110 if (params.n_parallel == 1) {
111 LOG_INF("%s: n_parallel == 1 -> unified KV cache is enabled\n", __func__);
112 params.kv_unified = true;
113 }
114
115 // utilize the full context
116 if (params.n_batch < params.n_ctx) {
117 LOG_WRN("%s: setting batch size to %d\n", __func__, params.n_ctx);
118 params.n_batch = params.n_ctx;
119 }
120
121 // for non-causal models, batch size must be equal to ubatch size
122 if (params.attention_type != LLAMA_ATTENTION_TYPE_CAUSAL) {
123 params.n_ubatch = params.n_batch;
124 }
125
126 // get max number of sequences per batch
127 const int n_seq_max = llama_max_parallel_sequences();
128
129 llama_backend_init();
130 llama_numa_init(numa: params.numa);
131
132 // load the model
133 common_init_result llama_init = common_init_from_params(params);
134
135 llama_model * model = llama_init.model.get();
136 llama_context * ctx = llama_init.context.get();
137
138 if (model == NULL) {
139 LOG_ERR("%s: unable to load model\n", __func__);
140 return 1;
141 }
142
143 const llama_vocab * vocab = llama_model_get_vocab(model);
144
145 const int n_ctx_train = llama_model_n_ctx_train(model);
146 const int n_ctx = llama_n_ctx(ctx);
147
148 const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);
149
150 if (llama_model_has_encoder(model) && llama_model_has_decoder(model)) {
151 LOG_ERR("%s: computing embeddings in encoder-decoder models is not supported\n", __func__);
152 return 1;
153 }
154
155 if (n_ctx > n_ctx_train) {
156 LOG_WRN("%s: warning: model was trained on only %d context tokens (%d specified)\n",
157 __func__, n_ctx_train, n_ctx);
158 }
159
160 // print system information
161 {
162 LOG_INF("\n");
163 LOG_INF("%s\n", common_params_get_system_info(params).c_str());
164 }
165
166 // split the prompt into lines
167 std::vector<std::string> prompts = split_lines(s: params.prompt, separator: params.embd_sep);
168
169 // max batch size
170 const uint64_t n_batch = params.n_batch;
171
172 // get added sep and eos token, if any
173 const std::string added_sep_token = llama_vocab_get_add_sep(vocab) ? llama_vocab_get_text(vocab, token: llama_vocab_sep(vocab)) : "";
174 const std::string added_eos_token = llama_vocab_get_add_eos(vocab) ? llama_vocab_get_text(vocab, token: llama_vocab_eos(vocab)) : "";
175 const char * rerank_prompt = llama_model_chat_template(model, name: "rerank");
176
177 // tokenize the prompts and trim
178 std::vector<std::vector<int32_t>> inputs;
179 for (const auto & prompt : prompts) {
180 std::vector<llama_token> inp;
181
182 // split classification pairs and insert expected separator tokens
183 if (pooling_type == LLAMA_POOLING_TYPE_RANK && prompt.find(str: params.cls_sep) != std::string::npos) {
184 std::vector<std::string> pairs = split_lines(s: prompt, separator: params.cls_sep);
185 if (rerank_prompt != nullptr) {
186 const std::string query = pairs[0];
187 const std::string doc = pairs[1];
188 std::string final_prompt = rerank_prompt;
189 string_replace_all(s&: final_prompt, search: "{query}" , replace: query);
190 string_replace_all(s&: final_prompt, search: "{document}", replace: doc );
191 inp = common_tokenize(vocab, text: final_prompt, add_special: true, parse_special: true);
192 } else {
193 std::string final_prompt;
194 for (size_t i = 0; i < pairs.size(); i++) {
195 final_prompt += pairs[i];
196 if (i != pairs.size() - 1) {
197 if (!added_eos_token.empty()) {
198 final_prompt += added_eos_token;
199 }
200 if (!added_sep_token.empty()) {
201 final_prompt += added_sep_token;
202 }
203 }
204 }
205 inp = common_tokenize(ctx, text: final_prompt, add_special: true, parse_special: true);
206 }
207 } else {
208 inp = common_tokenize(ctx, text: prompt, add_special: true, parse_special: true);
209 }
210 if (inp.size() > n_batch) {
211 LOG_ERR("%s: number of tokens in input line (%lld) exceeds batch size (%lld), increase batch size and re-run\n",
212 __func__, (long long int) inp.size(), (long long int) n_batch);
213 return 1;
214 }
215 inputs.push_back(x: inp);
216 }
217
218 // check if the last token is SEP/EOS
219 // it should be automatically added by the tokenizer when 'tokenizer.ggml.add_eos_token' is set to 'true'
220 for (auto & inp : inputs) {
221 if (inp.empty() || (inp.back() != llama_vocab_sep(vocab) && inp.back() != llama_vocab_eos(vocab))) {
222 LOG_WRN("%s: last token in the prompt is not SEP or EOS\n", __func__);
223 LOG_WRN("%s: 'tokenizer.ggml.add_eos_token' should be set to 'true' in the GGUF header\n", __func__);
224 }
225 }
226
227 // tokenization stats
228 if (params.verbose_prompt) {
229 for (int i = 0; i < (int) inputs.size(); i++) {
230 LOG_INF("%s: prompt %d: '%s'\n", __func__, i, prompts[i].c_str());
231 LOG_INF("%s: number of tokens in prompt = %zu\n", __func__, inputs[i].size());
232 for (int j = 0; j < (int) inputs[i].size(); j++) {
233 LOG("%6d -> '%s'\n", inputs[i][j], common_token_to_piece(ctx, inputs[i][j]).c_str());
234 }
235 LOG("\n\n");
236 }
237 }
238
239 // initialize batch
240 const int n_prompts = prompts.size();
241 struct llama_batch batch = llama_batch_init(n_tokens: n_batch, embd: 0, n_seq_max: 1);
242
243 // count number of embeddings
244 int n_embd_count = 0;
245 if (pooling_type == LLAMA_POOLING_TYPE_NONE) {
246 for (int k = 0; k < n_prompts; k++) {
247 n_embd_count += inputs[k].size();
248 }
249 } else {
250 n_embd_count = n_prompts;
251 }
252
253 // allocate output
254 const int n_embd = llama_model_n_embd(model);
255 std::vector<float> embeddings(n_embd_count * n_embd, 0);
256 float * emb = embeddings.data();
257
258 // break into batches
259 int e = 0; // number of embeddings already stored
260 int s = 0; // number of prompts in current batch
261 for (int k = 0; k < n_prompts; k++) {
262 // clamp to n_batch tokens
263 auto & inp = inputs[k];
264
265 const uint64_t n_toks = inp.size();
266
267 // encode if at capacity
268 if (batch.n_tokens + n_toks > n_batch || s >= n_seq_max) {
269 float * out = emb + e * n_embd;
270 batch_decode(ctx, batch, output: out, n_seq: s, n_embd, embd_norm: params.embd_normalize);
271 e += pooling_type == LLAMA_POOLING_TYPE_NONE ? batch.n_tokens : s;
272 s = 0;
273 common_batch_clear(batch);
274 }
275
276 // add to batch
277 batch_add_seq(batch, tokens: inp, seq_id: s);
278 s += 1;
279 }
280
281 // final batch
282 float * out = emb + e * n_embd;
283 batch_decode(ctx, batch, output: out, n_seq: s, n_embd, embd_norm: params.embd_normalize);
284
285 if (params.embd_out.empty()) {
286 LOG("\n");
287
288 if (pooling_type == LLAMA_POOLING_TYPE_NONE) {
289 for (int j = 0; j < n_embd_count; j++) {
290 LOG("embedding %d: ", j);
291 for (int i = 0; i < std::min(a: 3, b: n_embd); i++) {
292 if (params.embd_normalize == 0) {
293 LOG("%6.0f ", emb[j * n_embd + i]);
294 } else {
295 LOG("%9.6f ", emb[j * n_embd + i]);
296 }
297 }
298 LOG(" ... ");
299 for (int i = n_embd - 3; i < n_embd; i++) {
300 if (params.embd_normalize == 0) {
301 LOG("%6.0f ", emb[j * n_embd + i]);
302 } else {
303 LOG("%9.6f ", emb[j * n_embd + i]);
304 }
305 }
306 LOG("\n");
307 }
308 } else if (pooling_type == LLAMA_POOLING_TYPE_RANK) {
309 const uint32_t n_cls_out = llama_model_n_cls_out(model);
310 std::vector<std::string> cls_out_labels;
311
312 for (uint32_t i = 0; i < n_cls_out; i++) {
313 const char * label = llama_model_cls_label(model, i);
314 const std::string label_i(label == nullptr ? "" : label);
315 cls_out_labels.emplace_back(args: label_i.empty() ? std::to_string(val: i) : label_i);
316 }
317
318 for (int j = 0; j < n_embd_count; j++) {
319 for (uint32_t i = 0; i < n_cls_out; i++) {
320 // NOTE: if you change this log - update the tests in ci/run.sh
321 if (n_cls_out == 1) {
322 LOG("rerank score %d: %8.3f\n", j, emb[j * n_embd]);
323 } else {
324 LOG("rerank score %d: %8.3f [%s]\n", j, emb[j * n_embd + i], cls_out_labels[i].c_str());
325 }
326 }
327 }
328 } else {
329 // print the first part of the embeddings or for a single prompt, the full embedding
330 for (int j = 0; j < n_prompts; j++) {
331 LOG("embedding %d: ", j);
332 for (int i = 0; i < (n_prompts > 1 ? std::min(a: 16, b: n_embd) : n_embd); i++) {
333 if (params.embd_normalize == 0) {
334 LOG("%6.0f ", emb[j * n_embd + i]);
335 } else {
336 LOG("%9.6f ", emb[j * n_embd + i]);
337 }
338 }
339 LOG("\n");
340 }
341
342 // print cosine similarity matrix
343 if (n_prompts > 1) {
344 LOG("\n");
345 LOG("cosine similarity matrix:\n\n");
346 for (int i = 0; i < n_prompts; i++) {
347 LOG("%6.6s ", prompts[i].c_str());
348 }
349 LOG("\n");
350 for (int i = 0; i < n_prompts; i++) {
351 for (int j = 0; j < n_prompts; j++) {
352 float sim = common_embd_similarity_cos(embd1: emb + i * n_embd, embd2: emb + j * n_embd, n: n_embd);
353 LOG("%6.2f ", sim);
354 }
355 LOG("%1.10s", prompts[i].c_str());
356 LOG("\n");
357 }
358 }
359 }
360 }
361
362 if (params.embd_out == "json" || params.embd_out == "json+" || params.embd_out == "array") {
363 const bool notArray = params.embd_out != "array";
364
365 LOG(notArray ? "{\n \"object\": \"list\",\n \"data\": [\n" : "[");
366 for (int j = 0;;) { // at least one iteration (one prompt)
367 if (notArray) LOG(" {\n \"object\": \"embedding\",\n \"index\": %d,\n \"embedding\": ",j);
368 LOG("[");
369 for (int i = 0;;) { // at least one iteration (n_embd > 0)
370 LOG(params.embd_normalize == 0 ? "%1.0f" : "%1.7f", emb[j * n_embd + i]);
371 i++;
372 if (i < n_embd) LOG(","); else break;
373 }
374 LOG(notArray ? "]\n }" : "]");
375 j++;
376 if (j < n_embd_count) LOG(notArray ? ",\n" : ","); else break;
377 }
378 LOG(notArray ? "\n ]" : "]\n");
379
380 if (params.embd_out == "json+" && n_prompts > 1) {
381 LOG(",\n \"cosineSimilarity\": [\n");
382 for (int i = 0;;) { // at least two iteration (n_embd_count > 1)
383 LOG(" [");
384 for (int j = 0;;) { // at least two iteration (n_embd_count > 1)
385 float sim = common_embd_similarity_cos(embd1: emb + i * n_embd, embd2: emb + j * n_embd, n: n_embd);
386 LOG("%6.2f", sim);
387 j++;
388 if (j < n_embd_count) LOG(", "); else break;
389 }
390 LOG(" ]");
391 i++;
392 if (i < n_embd_count) LOG(",\n"); else break;
393 }
394 LOG("\n ]");
395 }
396
397 if (notArray) LOG("\n}\n");
398 } else if (params.embd_out == "raw") {
399 print_raw_embeddings(emb, n_embd_count, n_embd, model, pooling_type, embd_normalize: params.embd_normalize);
400 }
401
402 LOG("\n");
403 llama_perf_context_print(ctx);
404
405 // clean up
406 llama_batch_free(batch);
407 llama_backend_free();
408
409 return 0;
410}
411