add qwen3 rag models

Author: foldl

README.md

@@ -13,6 +13,7 @@ pure C++ implementation based on [@ggerganov](https://github.com/ggerganov)'s [g
 
 **What's New:**
 
+* 2025-06-06: Qwen-3 Embedding & Reranker
 * 2025-06-03: Kimi-VL
 * 2025-05-28: Gemma3 fully supported
 * 2025-05-23: [I can see](./docs/multimodal.md): Fuyu

convert.py

@@ -201,6 +201,8 @@ class ModelType(Enum):
     OrpheusTTS              = 0x10000106
     OuteTTSLlaMA            = 0x10000107
     OuteTTSQwen3            = 0x10000108
+    QWen3_Embedding         = 0x10000109
+    QWen3_ReRanker          = 0x1000010A
 
     LlaMAMulti    = 0x20000001
 
@@ -4489,6 +4491,7 @@ class QWen3Converter(BaseConverter):
     MODEL_TYPE = ModelType.QWen3
 
     layer_is_sparse = []
+    has_lm_head = True
 
     @staticmethod
     def dump_config(f, config, ggml_type):
@@ -4591,13 +4594,33 @@ def get_weight_names(config):
             "model.norm.weight"
         ]
 
-        if not config.tie_word_embeddings:
+        if QWen3Converter.has_lm_head and (not config.tie_word_embeddings):
             weight_names += [
                 "lm_head.weight"
             ]
 
         return weight_names
 
+class QWen3EmbConverter(BaseConverter):
+    MODEL_TYPE = ModelType.QWen3_Embedding
+
+    @classmethod
+    def state_dict_pp(cls, config, state_dict):
+        r = {}
+        for name in state_dict:
+            r['model.' + name] = state_dict[name]
+
+        return r
+
+    @staticmethod
+    def dump_config(f, config, ggml_type):
+        QWen3Converter.dump_config(f, config, ggml_type)
+
+    @staticmethod
+    def get_weight_names(config):
+        QWen3Converter.has_lm_head = False
+        return QWen3Converter.get_weight_names(config)
+
 def permute2(weights: torch.Tensor, n_head: int, partial_rotary_factor: float) -> torch.Tensor:
     hidden_size = weights.shape[0]
     head_dim = hidden_size // n_head
@@ -7126,6 +7149,11 @@ def main():
     elif arch == 'deepseek-r1-distill-qwen3':
         QWen3Converter.MODEL_TYPE = ModelType.DeepSeek_R1_Distill_QWen3
         QWen3Converter.convert(config, model_files, vocab, ggml_type, args.save_path)
+    elif arch == 'qwen3-embedding':
+        QWen3EmbConverter.convert(config, model_files, vocab, ggml_type, args.save_path)
+    elif arch == 'qwen3-reranker':
+        QWen3Converter.MODEL_TYPE = ModelType.QWen3_ReRanker
+        QWen3Converter.convert(config, model_files, vocab, ggml_type, args.save_path)
     elif arch == 'reka-flash-3':
         assert config.rope_scaling is None, 'config.rope_scaling must be null'
         assert not config.tie_word_embeddings, 'config.tie_word_embeddings must be false'

docs/models.md

@@ -313,10 +313,13 @@ Please use `--format completion` for these models.
 
         Note: Only dense embedding is implemented.
 
+    * Qwen-3 Embedding: [0.6B](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B/tree/b22da495047858cce924d27d76261e96be6febc0), [4B](https://huggingface.co/Qwen/Qwen3-Embedding-4B/tree/636cd9bf47d976946cdbb2b0c3ca0cb2f8eea5ff), [8B](https://huggingface.co/Qwen/Qwen3-Embedding-8B/commit/4e423935c619ae4df87b646a3ce949610c66241c)
+
 * QA Ranking (`XLMRobertaForSequenceClassification`)
     * [x] [BCE-ReRanker](https://huggingface.co/maidalun1020/bce-reranker-base_v1)
     * [x] [BGE-ReRanker-M3](https://huggingface.co/BAAI/bge-reranker-v2-m3) (`-a BGE-Reranker-M3`)
     * [x] [MiniCPM-Reranker-Light](https://huggingface.co/openbmb/MiniCPM-Reranker-Light)
+    * [x] Qwen-3 Reranker: [0.6B](https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/tree/ad4c588e592307dad69ff0fabc1b3ca5ea8e9f76), [4B](https://huggingface.co/Qwen/Qwen3-Reranker-4B/tree/57906229d41697e4494d50ca5859598cf86154a1), [8B](https://huggingface.co/Qwen/Qwen3-Reranker-8B/tree/d678ef8b29dd0eb9d784473da5d5169b21ec948a)
 
 ## LoRA Models
 

models/qwen.cpp

@@ -793,23 +793,23 @@ namespace v3
     class ConditionalGeneration : public BaseModelForConditionalGeneration
     {
     public:
-        typedef BaseModelForConditionalGeneration Base;
         typedef HeterogeneousModel ModelClass;
     public:
-        ConditionalGeneration(const Config &config, const RuntimeConfig &runtime_config, ModelType type = ModelType::MODEL_TYPE_QWEN3)
+        ConditionalGeneration(const Config &config, const RuntimeConfig &runtime_config, ModelType type = ModelType::MODEL_TYPE_QWEN3, const bool skip_lm_head = false, int extra_tensors = 0)
             : BaseModelForConditionalGeneration(type, config, runtime_config, 4096 * 4),
               config(config)
         {
             const size_t tensor_ovhd = ggml_tensor_overhead();
             const int sparse_layers = get_sparse_layer_num();
             const size_t num_tensors = 3 + (config.tie_word_embeddings ? -1 : 0)
                                          + (config.num_hidden_layers - sparse_layers) * 14
-                                         + sparse_layers * (14 + 1);
+                                         + sparse_layers * (14 + 1)
+                                         + extra_tensors;
             const size_t ctx_size = num_tensors * tensor_ovhd;
             w_ctx_.gctx = GGMLContext({.mem_size = ctx_size, .mem_buffer = nullptr, .no_alloc = true});
             w_ctx_.dtype = config.dtype;
 
-            if (config.tie_word_embeddings)
+            if (skip_lm_head || config.tie_word_embeddings)
             {
                 transformer = new ModelClass(&w_ctx_, config.num_hidden_layers, config.hidden_size,
                     create_embedding<Embedding>(&w_ctx_, config),
@@ -837,18 +837,18 @@ namespace v3
             {
                 if (config.layer_is_sparse[i])
                 {
-                    auto layer = (QWen3MoEBlock128_8 *)Base::get_typed_transformer<ModelClass>()->get_layer(i);
+                    auto layer = (QWen3MoEBlock128_8 *)get_typed_transformer<ModelClass>()->get_layer(i);
                     layer->attention.freq_base = config.rope_theta;
                     layer->mlp.norm_topk_prob = config.norm_topk_prob != 0;
                 }
                 else
                 {
-                    auto layer = (QWen3Block *)Base::get_typed_transformer<ModelClass>()->get_layer(i);
+                    auto layer = (QWen3Block *)get_typed_transformer<ModelClass>()->get_layer(i);
                     layer->attention.freq_base = config.rope_theta;
                 }
             }
 
-            CHATLLM_CHECK(w_ctx_.get_used_mem() == w_ctx_.get_mem_size())
+            CHATLLM_CHECK(w_ctx_.get_used_mem() + extra_tensors * ggml_tensor_overhead() == w_ctx_.get_mem_size())
                 << "corrupted model weights: " << w_ctx_.get_used_mem() / ggml_tensor_overhead() << " != " << w_ctx_.get_mem_size() / ggml_tensor_overhead();
         }
 
@@ -913,4 +913,151 @@ namespace ds_r1_distill_v3
     };
 
     typedef v3::ConditionalGeneration ConditionalGeneration;
+}
+
+
+namespace v3_emb
+{
+    typedef v3::Config Config;
+
+    class Tokenizer : public v3::Tokenizer
+    {
+    public:
+        Tokenizer(const BaseConfig &config)
+            : v3::Tokenizer(config)
+        {
+            task = "Given a web search query, retrieve relevant passages that answer the query";
+        }
+
+        void encode_embedding(const std::string &text, std::vector<int> &ids, EmbeddingPurpose purpose) const override;
+
+    public:
+        std::string task;
+    };
+
+    void Tokenizer::encode_embedding(const std::string &text, std::vector<int> &ids, EmbeddingPurpose purpose) const
+    {
+        std::ostringstream oss;
+        switch (purpose)
+        {
+        case EmbeddingPurpose::Query:
+            oss << "Instruct: " << task << "\nQuery:" << text;
+            BaseTokenizer::encode(oss.str(), ids);
+            break;
+
+        default:
+            BaseTokenizer::encode(text, ids);
+            break;
+        }
+        ids.push_back(eos_token_id);
+    }
+
+
+    class ConditionalGeneration : public v3::ConditionalGeneration
+    {
+    public:
+        ConditionalGeneration(const Config &config, const RuntimeConfig &runtime_config, ModelType type = ModelType::MODEL_TYPE_QWEN3_Embedding, const bool skip_lm_head = true, int extra_tensors = 0)
+            : v3::ConditionalGeneration(config, runtime_config, type, skip_lm_head, extra_tensors)
+        {
+            dynamic_cast<HeterogeneousModel *>(transformer)->set_final_steps(std::make_unique<EmbeddingLastTokenFinalSteps>());
+        }
+
+        void set_additional_args(const std::map<std::string, std::string> &args) override
+        {
+            Tokenizer *tok = dynamic_cast<Tokenizer *>(tokenizer);
+            auto it = args.find("task");
+            if (it != args.end())
+            {
+                tok->task = it->second;
+            }
+        }
+    };
+}
+
+namespace v3_ranker
+{
+    typedef v3::Config Config;
+
+    class Tokenizer : public v3_emb::Tokenizer
+    {
+    public:
+        Tokenizer(const BaseConfig &config)
+            : v3_emb::Tokenizer(config)
+        {
+        }
+
+        size_t load(tokenizer::DataReader *buffer, int n_vocab) override;
+
+        void encode_qa(const std::string &q, const std::string &a, std::vector<int> &ids) const override;
+    public:
+        int yes_token_id;
+        int no_token_id;
+    };
+
+    size_t Tokenizer::load(tokenizer::DataReader *buffer, int n_vocab)
+    {
+        size_t size = v3_emb::Tokenizer::load(buffer, n_vocab);
+
+        yes_token_id = tp->PieceToId("yes");
+         no_token_id = tp->PieceToId("no");
+
+        return size;
+    }
+
+    void Tokenizer::encode_qa(const std::string &q, const std::string &a, std::vector<int> &ids) const
+    {
+        std::ostringstream oss;
+        oss << "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n<|im_start|>user\n";
+        oss << "<Instruct>: " << task << "\n<Query>: " << q << "\n<Document>: " << a;
+        oss << "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
+
+        BaseTokenizer::encode(oss.str(), ids);
+    }
+
+    class FinalSteps : public LMFinalSteps
+    {
+    public:
+        ggml::tensor *forward(HeterogeneousModel *model, ComputeContext *ctx, ggml::tensor *input_ids, ggml::tensor *hidden_states) override;
+    public:
+        ggml::tensor *yes_no_ids;
+    };
+
+    ggml::tensor *FinalSteps::forward(HeterogeneousModel *model, ComputeContext *ctx, ggml::tensor *input_ids, ggml::tensor *hidden_states)
+    {
+        ggml::tensor *logits = LMFinalSteps::forward(model, ctx, input_ids, hidden_states);
+        logits = ggml::reshape_2d(ctx, logits, 1, ggml::get_dim(logits, 0));
+        logits = ggml::get_rows(ctx, logits, yes_no_ids);
+        logits = ggml::reshape_1d(ctx, logits, 2);
+        logits = ggml::soft_max(ctx, logits);
+        logits = ggml::view_1d(ctx, logits, 1, 0);
+        return logits;
+    }
+
+    class ConditionalGeneration : public v3_emb::ConditionalGeneration
+    {
+    public:
+        ConditionalGeneration(const Config &config, const RuntimeConfig &runtime_config)
+            : v3_emb::ConditionalGeneration(config, runtime_config, MODEL_TYPE_QWEN3_ReRanker, false, 1)
+        {
+            dynamic_cast<HeterogeneousModel *>(transformer)->set_final_steps(std::make_unique<FinalSteps>());
+
+            FinalSteps *steps = dynamic_cast<FinalSteps *>(dynamic_cast<HeterogeneousModel *>(transformer)->get_final_steps());
+            steps->yes_no_ids = ggml::new_tensor_1d(&w_ctx_, ggml::type::GGML_TYPE_I32, 2);
+            w_ctx_.get_allocator()->alloc(steps->yes_no_ids);
+            yes_no_ids = steps->yes_no_ids;
+        }
+
+        void set_tokenizer(BaseTokenizer *tokenizer) override
+        {
+            v3::ConditionalGeneration::set_tokenizer(tokenizer);
+
+            Tokenizer *tok = dynamic_cast<Tokenizer *>(tokenizer);
+            int ids[2];
+            ids[0] = tok->yes_token_id;
+            ids[1] = tok->no_token_id;
+            Backend::write_tensor_data(yes_no_ids, ids, 0, sizeof(ids));
+        }
+    protected:
+        ggml::tensor *yes_no_ids = nullptr;
+    };
 }

src/chat.cpp

@@ -613,6 +613,11 @@ namespace chatllm
         tp->Encode(input, &ids);
     }
 
+    void BaseTokenizer::encode_embedding(const std::string &text, std::vector<int> &ids, EmbeddingPurpose purpose) const
+    {
+        encode(text, ids);
+    }
+
     std::vector<int> BaseTokenizer::encode(const std::string &text) const
     {
         std::vector<int> ids;
@@ -1851,11 +1856,11 @@ namespace chatllm
         tokenizer->encode(input, result);
     }
 
-    void Pipeline::text_embedding(const std::string &input, const GenerationConfig &gen_config, std::vector<float> &result)
+    void Pipeline::text_embedding(const std::string &input, const GenerationConfig &gen_config, std::vector<float> &result, BaseTokenizer::EmbeddingPurpose purpose)
     {
         if (!modelobj.loaded) return;
         std::vector<int> input_ids;
-        tokenizer->encode(input, input_ids);
+        tokenizer->encode_embedding(input, input_ids, purpose);
         model->text_embedding(gen_config, input_ids, result);
     }
 

src/chat.h

@@ -250,6 +250,13 @@ namespace chatllm
             int emb_vec_number;
             std::vector<float> data;
         };
+
+        enum EmbeddingPurpose
+        {
+            Document,
+            Query,
+        };
+
     public:
         BaseTokenizer(const BaseConfig &config,
                         BaseHistoryEncoder *chat_encoder,
@@ -268,6 +275,7 @@ namespace chatllm
         virtual void encode_external_text_completion(const std::string &text, std::vector<int> &ids) const;
 
         virtual void encode_qa(const std::string &q, const std::string &a, std::vector<int> &ids) const;
+        virtual void encode_embedding(const std::string &text, std::vector<int> &ids, EmbeddingPurpose purpose) const;
 
         virtual std::string decode(const std::vector<int> &ids) const;
 
@@ -1320,7 +1328,7 @@ namespace chatllm
         void set_extending_method(ExtendingMethod method);
         virtual void set_additional_args(const std::map<std::string, std::string> &args);
 
-        void text_embedding(const std::string &input, const GenerationConfig &gen_config, std::vector<float> &result);
+        void text_embedding(const std::string &input, const GenerationConfig &gen_config, std::vector<float> &result, BaseTokenizer::EmbeddingPurpose purpose = BaseTokenizer::EmbeddingPurpose::Document);
         void text_tokenize(const std::string &input, const GenerationConfig &gen_config, std::vector<int> &result);
         float qa_rank(const std::string &q, const std::string &a, const GenerationConfig &gen_config);
 

src/layers.h

@@ -125,7 +125,7 @@ namespace chatllm
         ggml::tensor *norm_inplace(ComputeContext *ctx, ggml::tensor *a, float eps);
         ggml::tensor *rms_norm_inplace(ComputeContext *ctx, ggml::tensor *a, float eps);
         ggml::tensor *rms_norm(ComputeContext *ctx, ggml::tensor *a, float eps);
-        ggml::tensor *simple_norm(ComputeContext *ctx, ggml::tensor *a, float eps);
+        ggml::tensor *simple_norm(ComputeContext *ctx, ggml::tensor *a, float eps); // p=2 normalization
 
         ggml::tensor *rope(ComputeContext *ctx, ggml::tensor *a, ggml::tensor *b, int n_dims, int mode);
         ggml::tensor *rope_ext(ComputeContext *ctx, ggml::tensor *a, ggml::tensor *b, ggml::tensor *c,

src/main.cpp

@@ -811,17 +811,18 @@ static void run_tts(Args &args, chatllm::Pipeline &pipeline, TextStreamer &strea
 static void run_text_embedding(Args &args, chatllm::Pipeline &pipeline, TextStreamer &streamer, const chatllm::GenerationConfig &gen_config)
 {
     std::vector<float> result;
+    chatllm::BaseTokenizer::EmbeddingPurpose purpose = chatllm::BaseTokenizer::EmbeddingPurpose::Document;
 
     if (!args.interactive)
     {
-        pipeline.text_embedding(args.prompt, gen_config, result);
+        pipeline.text_embedding(args.prompt, gen_config, result, purpose);
         print_embedding(result, streamer.cout);
         return;
     }
 
     while (1)
     {
-        streamer.cout << "Input > " << std::flush;
+        streamer.cout << "Input " << (purpose == chatllm::BaseTokenizer::EmbeddingPurpose::Document ? "Doc" : "Query") <<  " > " << std::flush;
         std::string input;
         if (!get_utf8_line(input, args.multi_line))
         {
@@ -831,11 +832,13 @@ static void run_text_embedding(Args &args, chatllm::Pipeline &pipeline, TextStre
         if (input.empty()) continue;
 
         result.clear();
-        pipeline.text_embedding(input, gen_config, result);
+        pipeline.text_embedding(input, gen_config, result, purpose);
         streamer.cout << "      > ";
 
         print_embedding(result, streamer.cout);
 
+        purpose = purpose == chatllm::BaseTokenizer::EmbeddingPurpose::Document ?
+                    chatllm::BaseTokenizer::EmbeddingPurpose::Query : chatllm::BaseTokenizer::EmbeddingPurpose::Document;
     }
     streamer.cout << "Bye\n";
 }