qwen3 embeddings support (#3529)

* Add support for last token pooling mode

Co-authored-by: Trawinski, Dariusz <[email protected]>

---------

Co-authored-by: Damian Kalinowski <[email protected]>
Co-authored-by: Trawinski, Dariusz <[email protected]>

Author: 3 people

demos/common/export_models/export_model.py

@@ -65,6 +65,7 @@ def add_common_arguments(parser):
 parser_embeddings_ov = subparsers.add_parser('embeddings_ov', help='export model for embeddings endpoint with directory structure aligned with OpenVINO tools')
 add_common_arguments(parser_embeddings_ov)
 parser_embeddings_ov.add_argument('--skip_normalize', default=True, action='store_false', help='Skip normalize the embeddings.', dest='normalize')
+parser_embeddings_ov.add_argument('--pooling', default="CLS", choices=["CLS", "LAST"], help='Embeddings pooling mode', dest='pooling')
 parser_embeddings_ov.add_argument('--truncate', default=False, action='store_true', help='Truncate the prompts to fit to the embeddings model', dest='truncate')
 parser_embeddings_ov.add_argument('--num_streams', default=1,type=int, help='The number of parallel execution streams to use for the model. Use at least 2 on 2 socket CPU systems.', dest='num_streams')
 
@@ -139,6 +140,8 @@ def add_common_arguments(parser):
     [type.googleapis.com / mediapipe.EmbeddingsCalculatorOVOptions]: {
       models_path: "{{model_path}}",
       normalize_embeddings: {% if not normalize %}false{% else %}true{% endif%},
+      {%- if pooling %}
+      pooling: {{pooling}},{% endif %}
       target_device: "{{target_device|default("CPU", true)}}"
     }
   }

demos/embeddings/README.md

@@ -57,15 +57,16 @@ models
 └── config.json
 
 ```
-> **Note** The actual models support version management and can be automatically swapped to newer version when new model is uploaded in newer version folder.
-> In case you trained the pytorch model it can be exported like below:
-> `python export_model.py embeddings_ov --source_model <pytorch model> --model_name Alibaba-NLP/gte-large-en-v1.5 --precision int8 --config_file_path models/config.json --version 2`
 
-The default configuration of the `EmbeddingsCalculator` should work in most cases but the parameters can be tuned inside the `node_options` section in the `graph.pbtxt` file. Runtime configuration for both models can be tuned in `subconfig.json` file. They can be set automatically via export parameters in the `export_model.py` script.
+The default configuration of the `EmbeddingsCalculatorOV` should work in most cases but the parameters can be tuned inside the `node_options` section in the `graph.pbtxt` file. They can be set automatically via export parameters in the `export_model.py` script.
 
 For example:
-`python export_model.py embeddings_ov --source_model Alibaba-NLP/gte-large-en-v1.5 --precision int8 --num_streams 2 --skip_normalize --config_file_path models/config.json`
+`python export_model.py embeddings_ov --source_model Alibaba-NLP/gte-large-en-v1.5 --weight-format int8 --skip_normalize --config_file_path models/config.json`
 
+> **Note:** By default OVMS returns first token embeddings as sequence embeddings (called CLS pooling). It can be changed using `--pooling` option if needed by the model. Supported values are CLS and LAST. For example:
+```console
+python export_model.py embeddings_ov --source_model Qwen/Qwen3-Embedding-0.6B --weight-format fp16 --pooling LAST --config_file_path models/config.json`
+```
 
 ## Tested models
 All models supported by [optimum-intel](https://github.com/huggingface/optimum-intel) should be compatible. In serving validation are included Hugging Face models:
@@ -240,7 +241,7 @@ The script [compare_results.py](./compare_results.py) can assist with such exper
 ```bash
 popd
 cd model_server/demos/embeddings
-python compare_results.py --model Alibaba-NLP/gte-large-en-v1.5 --service_url http://localhost:8000/v3/embeddings --input "hello world" --input "goodbye world"
+python compare_results.py --model Alibaba-NLP/gte-large-en-v1.5 --service_url http://localhost:8000/v3/embeddings --pooling CLS --input "hello world" --input "goodbye world"
 
 input ['hello world', 'goodbye world']
 HF Duration: 50.626 ms NewModel

demos/embeddings/compare_results.py

@@ -28,6 +28,8 @@
                     dest='model_name')
 parser.add_argument('--input', default=[], help='List of strings to query. default: []',
                     dest='input', action='append')
+parser.add_argument('--pooling', default="CLS", choices=["CLS", "LAST"], help='Embeddings pooling mode', dest='pooling')
+
 args = vars(parser.parse_args())
 
 model_id = args['model_name']
@@ -43,8 +45,14 @@ def run_model():
         start_time = datetime.datetime.now()
         input = tokenizer(text, padding=True, truncation=True, return_tensors='pt')
         model_output = model_pt(**input)
-        embeddings = model_output.last_hidden_state[:, 0]
-        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+        if args['pooling'] == "LAST":
+            sequence_lengths = input['attention_mask'].sum(dim=1) - 1
+            batch_size = model_output.last_hidden_state.shape[0]
+            embeddings = model_output.last_hidden_state[torch.arange(batch_size, device=model_output.last_hidden_state.device), sequence_lengths]
+            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+        else:
+            embeddings = model_output.last_hidden_state[:, 0]
+            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
         end_time = datetime.datetime.now()
         duration = (end_time - start_time).total_seconds() * 1000
         print("HF Duration:", duration, "ms", type(model_pt).__name__)

src/capi_frontend/server_settings.hpp

@@ -114,7 +114,7 @@ struct EmbeddingsGraphSettingsImpl {
     std::string modelName = "";
     uint32_t numStreams = 1;
     std::string normalize = "true";
-    std::string meanPooling = "false";
+    std::string pooling = "CLS";
 };
 
 struct RerankGraphSettingsImpl {

src/embeddings/BUILD

@@ -29,6 +29,14 @@ ovms_cc_library(
     alwayslink = 1,
 )
 
+ovms_cc_library(
+    name = "embeddings_servable",
+    hdrs = ["embeddings_servable.hpp"],
+    deps = ["//src:sidepacket_servable",],
+    visibility = ["//visibility:public"],
+    alwayslink = 1,
+)
+
 mediapipe_proto_library(
     name = "embeddings_calculator_proto", # embeddings_calculator_cc_proto - just mediapipe stuff with mediapipe_proto_library adding nonvisible target
     srcs = ["embeddings_calculator.proto"],
@@ -81,11 +89,12 @@ ovms_cc_library(
         "//src:libovmslogging",
         "//src:libovmsprofiler",
         "embeddings_calculator_ov_cc_proto",
-        ":embeddings_api",
+        ":embeddings_servable",
         "//src:sidepacket_servable",
         "//src:model_metric_reporter",
         "//src:executingstreamidguard",
         "//src:libovms_execution_context",
+        ":embeddings_api",
     ],
     visibility = ["//visibility:public"],
     alwayslink = 1,

src/embeddings/embeddings_api.cpp

@@ -162,7 +162,7 @@ void EmbeddingsHandler::setPromptTokensUsage(int promptTokens) {
 
 #pragma warning(push)
 #pragma warning(disable : 4267)
-absl::Status EmbeddingsHandler::parseResponse(StringBuffer& buffer, const ov::Tensor& embeddingsTensor, const bool normalizeEmbeddings) {
+absl::Status EmbeddingsHandler::parseResponse(StringBuffer& buffer, const ov::Tensor& embeddingsTensor, const bool normalizeEmbeddings, const PoolingMode poolingMode, const std::optional<ov::Tensor>& attentionMask) {
     Writer<StringBuffer> writer(buffer);
     writer.StartObject();
 
@@ -171,15 +171,42 @@ absl::Status EmbeddingsHandler::parseResponse(StringBuffer& buffer, const ov::Te
 
     writer.String("data");
     writer.StartArray();
-    // TODO: mean pooling
 
     ov::Shape outputShape = embeddingsTensor.get_shape();
     if (outputShape.size() != 3) {
         return absl::InvalidArgumentError("Invalid embeddings tensor shape");
     }
     size_t batchSize = outputShape[0];
     for (size_t batchIterator = 0; batchIterator < batchSize; batchIterator++) {
-        size_t stride = batchIterator * outputShape[1] * outputShape[2];
+        size_t stride;
+        if (poolingMode == PoolingMode::LAST) {
+            size_t attendedTokens = 0;
+            if (!attentionMask.has_value()) {
+                return absl::InvalidArgumentError("Last token pooling mode requires attention mask");
+            }
+            auto maxNumberOfTokens = attentionMask->get_shape()[1];
+            if (attentionMask->get_element_type() == ov::element::Type_t::i64) {
+                for (int i = 0; i < maxNumberOfTokens; i++) {
+                    attendedTokens += reinterpret_cast<int64_t*>(attentionMask->data())[i + batchIterator * maxNumberOfTokens];
+                }
+            } else if (attentionMask->get_element_type() == ov::element::Type_t::i32) {
+                for (int i = 0; i < maxNumberOfTokens; i++) {
+                    attendedTokens += reinterpret_cast<int32_t*>(attentionMask->data())[i + batchIterator * maxNumberOfTokens];
+                }
+            } else if (attentionMask->get_element_type() == ov::element::Type_t::i8) {
+                for (int i = 0; i < maxNumberOfTokens; i++) {
+                    attendedTokens += reinterpret_cast<uint8_t*>(attentionMask->data())[i + batchIterator * maxNumberOfTokens];
+                }
+            } else {
+                return absl::InternalError("Attention mask element type invalid.");
+            }
+            if (!(attendedTokens <= outputShape[1])) {
+                return absl::InternalError("Embeddings output and attention mask shape mismatch");
+            }
+            stride = batchIterator * outputShape[1] * outputShape[2] + (attendedTokens - 1) * outputShape[2];
+        } else {
+            stride = batchIterator * outputShape[1] * outputShape[2];
+        }
         size_t size = outputShape[2];
         float* dataPtr = reinterpret_cast<float*>(embeddingsTensor.data()) + stride;
         float* dataPtrEnd = dataPtr + size;

src/embeddings/embeddings_api.hpp

@@ -36,6 +36,11 @@
 
 namespace ovms {
 
+enum class PoolingMode {
+    CLS,
+    LAST
+};
+
 struct EmbeddingsRequest {
     enum class EncodingFormat {
         FLOAT,
@@ -60,7 +65,7 @@ class EmbeddingsHandler {
     EmbeddingsRequest::EncodingFormat getEncodingFormat() const;
 
     absl::Status parseRequest();
-    absl::Status parseResponse(rapidjson::StringBuffer& buffer, const ov::Tensor& embeddingsTensor, const bool normalizeEmbeddings);
+    absl::Status parseResponse(rapidjson::StringBuffer& buffer, const ov::Tensor& embeddingsTensor, const bool normalizeEmbeddings, const PoolingMode poolingMode = PoolingMode::CLS, const std::optional<ov::Tensor>& attentionMask = std::nullopt);
     void setPromptTokensUsage(int promptTokens);
 };
 }  // namespace ovms

src/embeddings/embeddings_calculator.proto

@@ -26,5 +26,4 @@ message EmbeddingsCalculatorOptions {
     optional EmbeddingsCalculatorOptions ext = 1134737;
     }
     optional bool normalize_embeddings = 1 [default = true];
-    optional bool mean_pooling = 2 [default = false];
 }

src/embeddings/embeddings_calculator_ov.cc

@@ -40,7 +40,7 @@
 #include "../model_metric_reporter.hpp"
 #include "embeddings_api.hpp"
 #include "src/embeddings/embeddings_calculator_ov.pb.h"
-#include "../sidepacket_servable.hpp"
+#include "embeddings_servable.hpp"
 
 using namespace rapidjson;
 using namespace ovms;
@@ -63,7 +63,7 @@ class EmbeddingsCalculatorOV : public CalculatorBase {
     mediapipe::Timestamp timestamp{0};
 
 protected:
-    std::shared_ptr<ovms::SidepacketServable> embeddings_session{nullptr};
+    std::shared_ptr<ovms::EmbeddingsServable> embeddings_session{nullptr};
 
 public:
     static absl::Status GetContract(CalculatorContract* cc) {
@@ -148,10 +148,12 @@ class EmbeddingsCalculatorOV : public CalculatorBase {
                     for (int i = 0; i < tokens.attention_mask.get_size(); i++) {
                         attendedTokens += reinterpret_cast<int32_t*>(tokens.attention_mask.data())[i];
                     }
-                } else {
+                } else if (tokens.attention_mask.get_element_type() == ov::element::Type_t::i8) {
                     for (int i = 0; i < tokens.attention_mask.get_byte_size(); i++) {
                         attendedTokens += reinterpret_cast<uint8_t*>(tokens.attention_mask.data())[i];
                     }
+                } else {
+                    return absl::InternalError("Attention mask element type invalid.");
                 }
                 handler.setPromptTokensUsage(attendedTokens);
             } else if (auto tokenized_documents = std::get_if<std::vector<std::vector<int64_t>>>(&input)) {
@@ -241,7 +243,13 @@ class EmbeddingsCalculatorOV : public CalculatorBase {
 
         auto parseResponseStartTime = std::chrono::high_resolution_clock::now();
         StringBuffer buffer;
-        status = handler.parseResponse(buffer, embeddingsTensor, cc->Options<EmbeddingsCalculatorOVOptions>().normalize_embeddings());
+        PoolingMode mode;
+        if (cc->Options<EmbeddingsCalculatorOVOptions>().pooling() == mediapipe::EmbeddingsCalculatorOVOptions::LAST) {
+            mode = PoolingMode::LAST;
+        } else {
+            mode = PoolingMode::CLS;
+        }
+        status = handler.parseResponse(buffer, embeddingsTensor, cc->Options<EmbeddingsCalculatorOVOptions>().normalize_embeddings(), mode, tokens.attention_mask);
         if (!status.ok()) {
             return status;
         }

src/embeddings/embeddings_calculator_ov.proto

@@ -27,7 +27,11 @@ message EmbeddingsCalculatorOVOptions {
     }
     required string models_path = 1;
     optional bool normalize_embeddings = 2 [default = true];
-    optional bool mean_pooling = 3 [default = false];
-    optional string target_device = 4 [default = "CPU"];
-    optional string plugin_config = 5 [default = ""];
+    optional string target_device = 3 [default = "CPU"];
+    optional string plugin_config = 4 [default = ""];
+    enum Pooling {
+      CLS = 0;
+      LAST = 1;
+    }
+    optional Pooling pooling = 5 [default = CLS];
 }