From 1b382420d89eba6c3f6898915f9ad5ab52ff361c Mon Sep 17 00:00:00 2001 From: Arun Babu Neelicattu Date: Fri, 3 Jul 2026 19:13:55 +0200 Subject: [PATCH] feat(telemetry): add embedding token count usage Introduce embed_with_usage to AutoEmbeddingModel to return prompt token counts. Implement for Gemma_Embedding, avoiding double-tokenization, and update RestHandler to report total batch usage. Update standalone test suite to verify token counts. --- .../modeling_gemma_embedding.cpp | 8 +++++++- .../AutoEmbeddingModel/auto_embedding_model.hpp | 17 +++++++++++++++++ .../modeling_gemma_embedding.hpp | 1 + src/server/rest_handler.cpp | 9 ++++++--- src/test/gemma_embedding/test.cpp | 8 +++++++- 5 files changed, 38 insertions(+), 5 deletions(-) diff --git a/src/common/AutoEmbeddingModel/modeling_gemma_embedding.cpp b/src/common/AutoEmbeddingModel/modeling_gemma_embedding.cpp index ba10efb8..3e49a205 100644 --- a/src/common/AutoEmbeddingModel/modeling_gemma_embedding.cpp +++ b/src/common/AutoEmbeddingModel/modeling_gemma_embedding.cpp @@ -18,6 +18,11 @@ void Gemma_Embedding::load_model(std::string model_path, json model_info, bool e } std::vector Gemma_Embedding::embed(std::string& text, embedding_task_type_t task_type) { + size_t dummy_prompt_tokens; + return this->embed_with_usage(text, task_type, dummy_prompt_tokens); +} + +std::vector Gemma_Embedding::embed_with_usage(std::string& text, embedding_task_type_t task_type, size_t& prompt_tokens) { std::string task_prefix = this->_get_task_prefix(task_type); std::string full_text = "" + task_prefix + text + ""; std::vector tokens = this->tokenizer->encode(full_text); @@ -26,10 +31,11 @@ std::vector Gemma_Embedding::embed(std::string& text, embedding_task_type tokens.resize(2048); tokens[2047] = 1; } + prompt_tokens = tokens.size(); buffer y = this->embedding_model_impl->embed(tokens); std::vector result(y.size()); - for (int i = 0; i < y.size(); i++) { + for (size_t i = 0; i < y.size(); i++) { result[i] = static_cast(y[i]); } return result; diff --git a/src/include/AutoEmbeddingModel/auto_embedding_model.hpp b/src/include/AutoEmbeddingModel/auto_embedding_model.hpp index 73e2c049..4679e269 100644 --- a/src/include/AutoEmbeddingModel/auto_embedding_model.hpp +++ b/src/include/AutoEmbeddingModel/auto_embedding_model.hpp @@ -82,6 +82,23 @@ class AutoEmbeddingModel { virtual void load_model(std::string model_path, json model_info, bool enable_preemption) {} virtual std::vector embed(std::string& text, embedding_task_type_t task_type) = 0; + + /// \brief Embed the text and return prompt token count + /// \param text the text + /// \param task_type the task type + /// \param prompt_tokens [out] output variable to store the count of tokens processed + /// \return the embedding result vector + /// \note Subclasses should override this method to perform single-pass tokenization + /// and ensure model-specific formatting (e.g. prefixes/suffixes) is accounted for. + /// Default fallback implementation will tokenise raw text twice. + virtual std::vector embed_with_usage(std::string& text, embedding_task_type_t task_type, size_t& prompt_tokens) { + if (tokenizer) { + prompt_tokens = tokenizer->encode(text).size(); + } else { + prompt_tokens = 0; + } + return embed(text, task_type); + } }; diff --git a/src/include/AutoEmbeddingModel/modeling_gemma_embedding.hpp b/src/include/AutoEmbeddingModel/modeling_gemma_embedding.hpp index 53369492..fe1eb041 100644 --- a/src/include/AutoEmbeddingModel/modeling_gemma_embedding.hpp +++ b/src/include/AutoEmbeddingModel/modeling_gemma_embedding.hpp @@ -57,4 +57,5 @@ class Gemma_Embedding : public AutoEmbeddingModel{ void load_model(std::string model_path, json model_info, bool enable_preemption = false) override; std::vector embed(std::string& text, embedding_task_type_t task_type) override; + std::vector embed_with_usage(std::string& text, embedding_task_type_t task_type, size_t& prompt_tokens) override; }; \ No newline at end of file diff --git a/src/server/rest_handler.cpp b/src/server/rest_handler.cpp index 6920f9e1..2b294da8 100644 --- a/src/server/rest_handler.cpp +++ b/src/server/rest_handler.cpp @@ -843,10 +843,13 @@ void RestHandler::handle_embeddings(const json& request, json response; if (this->embed) { json embedding_data = json::array(); + size_t total_prompt_tokens = 0; #ifndef FASTFLOWLM_LINUX_LIMITED_MODELS for (size_t i = 0; i < inputs.size(); ++i) { std::cout << "Embedding input[" << i << "]: " << "\n" << inputs[i] << std::endl; - std::vector embedding_result = this->auto_embedding_engine->embed(inputs[i], embedding_task_type_t::task_query); + size_t prompt_tokens = 0; + std::vector embedding_result = this->auto_embedding_engine->embed_with_usage(inputs[i], embedding_task_type_t::task_query, prompt_tokens); + total_prompt_tokens += prompt_tokens; embedding_data.push_back({ {"object", "embedding"}, {"embedding", embedding_result}, @@ -862,8 +865,8 @@ void RestHandler::handle_embeddings(const json& request, {"data", embedding_data}, {"model", model}, {"usage", { - {"prompt_tokens", 0}, - {"total_tokens", 0} + {"prompt_tokens", total_prompt_tokens}, + {"total_tokens", total_prompt_tokens} }} }; } diff --git a/src/test/gemma_embedding/test.cpp b/src/test/gemma_embedding/test.cpp index 3a201ac4..25552ed0 100644 --- a/src/test/gemma_embedding/test.cpp +++ b/src/test/gemma_embedding/test.cpp @@ -113,13 +113,19 @@ int main(int argc, char* argv[]) { std::string text = "Alice's Adventures in Wonderland ALICE'S ADVENTURES IN WONDERLAND Lewis Carroll THE MILLENNIUM FULCRUM EDITION 3.0 CHAPTER I Down the Rabbit-HoleAlice was beginning to get very tired of sitting by her sister on the bank, and of having nothing to do: once or twice she had peeped into the book her sister was reading, but it had no pictures or conversations in it, and what is the use of a book,'thought Alice without pictures or conversation?' So she was considering in her own mind (as well as she could, for the hot day made her feel very sleepy and stupid), whether the pleasure of making a daisy-chain would be worth the trouble of getting up and picking the daisies, when suddenly a White Rabbit with pink eyes ran close by her. There was nothing so VERY remarkable in that; nor did Alice think it so VERY much out of the way to hear the Rabbit say to itself, `Oh dear! Oh dear! I shall be late!'"; time_utils::time_point start_time = time_utils::now(); - auto y = embedding->embed(text, task_query); + size_t prompt_tokens = 0; + auto y = embedding->embed_with_usage(text, task_query, prompt_tokens); buffer y_bf16(y.size()); for (int i = 0; i < y.size(); i++){ y_bf16[i] = (bf16)y[i]; } time_utils::time_point end_time = time_utils::now(); std::cout << "Time: " << time_utils::duration_ms(start_time, end_time).first << "ms" << std::endl; + std::cout << "Usage (prompt_tokens): " << prompt_tokens << std::endl; + if (prompt_tokens == 0) { + std::cerr << "Error: prompt_tokens is 0!" << std::endl; + return 1; + } utils::print_matrix(y_bf16, 128); buffer ref_bf16 = buffer(y_bf16.size());