add qwen3-embedding and all-minilm-l6-v2

2026-04-06 15:35:49 +08:00
parent 44d91da650
commit 7142e712f2
34 changed files with 763 additions and 319 deletions
@@ -1,3 +1,4 @@
+pub mod all_minilm_l6_v2;
 pub mod bigvgan;
 pub mod campplus;
 pub mod common;
@@ -17,16 +18,22 @@ pub mod qwen2_5vl;
 pub mod qwen3;
 pub mod qwen3_5;
 pub mod qwen3_asr;
+pub mod qwen3_embedding;
 pub mod qwen3vl;
 pub mod rmbg2_0;
 pub mod voxcpm;
 pub mod w2v_bert_2_0;

 use crate::{
-    models::common::model_mapping::WhichModel,
+    models::{
+        all_minilm_l6_v2::AllMiniLML6V2Embedding,
+        common::{embedding::TextEmbedding, model_mapping::WhichModel},
+        qwen3_embedding::Qwen3Embedding,
+    },
    params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse},
 };
 use anyhow::{Result, anyhow};
+use candle_core::{DType, Device};
 use rocket::futures::Stream;

 use crate::models::{
@@ -57,6 +64,7 @@ pub trait GenerateModel {
 }

 pub enum ModelInstance<'a> {
+    AllMiniLML6V2(AllMiniLML6V2Embedding),
    MiniCPM4(MiniCPMGenerateModel<'a>),
    Lfm2(Lfm2GenerateModel<'a>),
    Lfm2VL(Lfm2VLGenerateModel<'a>),
@@ -64,6 +72,7 @@ pub enum ModelInstance<'a> {
    Qwen3(Qwen3GenerateModel<'a>),
    Qwen3_5(Qwen3_5GenerateModel<'a>),
    Qwen3ASR(Qwen3AsrGenerateModel<'a>),
+    Qwen3Embedding(Qwen3Embedding),
    Qwen3VL(Box<Qwen3VLGenerateModel<'a>>),
    DeepSeekOCR(DeepseekOCRGenerateModel),
    HunyuanOCR(HunyuanOCRGenerateModel<'a>),
@@ -78,12 +87,18 @@ pub enum ModelInstance<'a> {
 impl<'a> GenerateModel for ModelInstance<'a> {
    fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
        match self {
+            ModelInstance::AllMiniLML6V2(_) => {
+                Err(anyhow!("embedding model does not support chat completions"))
+            }
            ModelInstance::MiniCPM4(model) => model.generate(mes),
            ModelInstance::Lfm2(model) => model.generate(mes),
            ModelInstance::Lfm2VL(model) => model.generate(mes),
            ModelInstance::Qwen2_5VL(model) => model.generate(mes),
            ModelInstance::Qwen3(model) => model.generate(mes),
            ModelInstance::Qwen3_5(model) => model.generate(mes),
+            ModelInstance::Qwen3Embedding(_) => {
+                Err(anyhow!("embedding model does not support chat completions"))
+            }
            ModelInstance::Qwen3ASR(model) => model.generate(mes),
            ModelInstance::Qwen3VL(model) => model.generate(mes),
            ModelInstance::DeepSeekOCR(model) => model.generate(mes),
@@ -109,12 +124,18 @@ impl<'a> GenerateModel for ModelInstance<'a> {
        >,
    > {
        match self {
+            ModelInstance::AllMiniLML6V2(_) => {
+                Err(anyhow!("embedding model does not support chat completions"))
+            }
            ModelInstance::MiniCPM4(model) => model.generate_stream(mes),
            ModelInstance::Lfm2(model) => model.generate_stream(mes),
            ModelInstance::Lfm2VL(model) => model.generate_stream(mes),
            ModelInstance::Qwen2_5VL(model) => model.generate_stream(mes),
            ModelInstance::Qwen3(model) => model.generate_stream(mes),
            ModelInstance::Qwen3_5(model) => model.generate_stream(mes),
+            ModelInstance::Qwen3Embedding(_) => Err(anyhow!(
+                "embedding model does not support streaming chat completions"
+            )),
            ModelInstance::Qwen3VL(model) => model.generate_stream(mes),
            ModelInstance::Qwen3ASR(model) => model.generate_stream(mes),
            ModelInstance::DeepSeekOCR(model) => model.generate_stream(mes),
@@ -129,16 +150,27 @@ impl<'a> GenerateModel for ModelInstance<'a> {
    }
 }

+impl<'a> ModelInstance<'a> {
+    pub fn embedding(&mut self, input: &[String]) -> Result<Vec<Vec<f32>>> {
+        match self {
+            ModelInstance::Qwen3Embedding(model) => model.embed_texts(input),
+            ModelInstance::AllMiniLML6V2(model) => model.embed_texts(input),
+            _ => Err(anyhow!("current model does not support embeddings")),
+        }
+    }
+}
+
 #[allow(unused)]
 pub fn load_gguf_model<'a>(
    model_type: WhichModel,
    config_path: Option<&str>, // 有些gguf未包含模型其他配置，需额外指定
    gguf_path: &str,
    mmproj_path: Option<&str>,
+    device: Option<&Device>,
 ) -> Result<ModelInstance<'a>> {
    let model = match model_type {
        WhichModel::Qwen3_5Gguf => {
-            let model = Qwen3_5GenerateModel::init_from_gguf(gguf_path, mmproj_path, None)?;
+            let model = Qwen3_5GenerateModel::init_from_gguf(gguf_path, mmproj_path, device)?;
            ModelInstance::Qwen3_5(model)
        }
        _ => {
@@ -149,135 +181,101 @@ pub fn load_gguf_model<'a>(
    Ok(model)
 }

-pub fn load_model<'a>(model_type: WhichModel, path: &str) -> Result<ModelInstance<'a>> {
+pub fn load_model<'a>(
+    model_type: WhichModel,
+    path: &str,
+    device: Option<&Device>,
+    dtype: Option<DType>,
+) -> Result<ModelInstance<'a>> {
    let model = match model_type {
+        WhichModel::AllMiniLML6V2 => {
+            let model = AllMiniLML6V2Embedding::init(path, device, dtype)?;
+            ModelInstance::AllMiniLML6V2(model)
+        }
        WhichModel::MiniCPM4_0_5B => {
-            let model = MiniCPMGenerateModel::init(path, None, None)?;
+            let model = MiniCPMGenerateModel::init(path, device, dtype)?;
            ModelInstance::MiniCPM4(model)
        }
-        WhichModel::LFM2_1_2B => {
-            let model = Lfm2GenerateModel::init(path, None, None)?;
+        WhichModel::LFM2_1_2B | WhichModel::LFM2_5_1_2BInstruct => {
+            let model = Lfm2GenerateModel::init(path, device, dtype)?;
            ModelInstance::Lfm2(model)
        }
-        WhichModel::LFM2_5_1_2BInstruct => {
-            let model = Lfm2GenerateModel::init(path, None, None)?;
-            ModelInstance::Lfm2(model)
-        }
-        WhichModel::LFM2_5VL1_6B => {
-            let model = Lfm2VLGenerateModel::init(path, None, None)?;
+        WhichModel::LFM2_5VL1_6B | WhichModel::LFM2VL1_6B => {
+            let model = Lfm2VLGenerateModel::init(path, device, dtype)?;
            ModelInstance::Lfm2VL(model)
        }
-        WhichModel::LFM2VL1_6B => {
-            let model = Lfm2VLGenerateModel::init(path, None, None)?;
-            ModelInstance::Lfm2VL(model)
-        }
-        WhichModel::Qwen2_5VL3B => {
-            let model = Qwen2_5VLGenerateModel::init(path, None, None)?;
+        WhichModel::Qwen2_5VL3B | WhichModel::Qwen2_5VL7B => {
+            let model = Qwen2_5VLGenerateModel::init(path, device, dtype)?;
            ModelInstance::Qwen2_5VL(model)
        }
-        WhichModel::Qwen2_5VL7B => {
-            let model = Qwen2_5VLGenerateModel::init(path, None, None)?;
-            ModelInstance::Qwen2_5VL(model)
-        }
-        WhichModel::Qwen3_0_6B => {
-            let model = Qwen3GenerateModel::init(path, None, None)?;
+        WhichModel::Qwen3_0_6B | WhichModel::Qwen3_1_7B | WhichModel::Qwen3_4B => {
+            let model = Qwen3GenerateModel::init(path, device, dtype)?;
            ModelInstance::Qwen3(model)
        }
-        WhichModel::Qwen3_1_7B => {
-            let model = Qwen3GenerateModel::init(path, None, None)?;
-            ModelInstance::Qwen3(model)
-        }
-        WhichModel::Qwen3_4B => {
-            let model = Qwen3GenerateModel::init(path, None, None)?;
-            ModelInstance::Qwen3(model)
-        }
-        WhichModel::Qwen3_5_0_8B => {
-            let model = Qwen3_5GenerateModel::init(path, None, None)?;
+        WhichModel::Qwen3_5_0_8B
+        | WhichModel::Qwen3_5_2B
+        | WhichModel::Qwen3_5_4B
+        | WhichModel::Qwen3_5_9B => {
+            let model = Qwen3_5GenerateModel::init(path, device, dtype)?;
            ModelInstance::Qwen3_5(model)
        }
-        WhichModel::Qwen3_5_2B => {
-            let model = Qwen3_5GenerateModel::init(path, None, None)?;
-            ModelInstance::Qwen3_5(model)
-        }
-        WhichModel::Qwen3_5_4B => {
-            let model = Qwen3_5GenerateModel::init(path, None, None)?;
-            ModelInstance::Qwen3_5(model)
-        }
-        WhichModel::Qwen3_5_9B => {
-            let model = Qwen3_5GenerateModel::init(path, None, None)?;
-            ModelInstance::Qwen3_5(model)
-        }
-        WhichModel::Qwen3ASR0_6B => {
-            let model = Qwen3AsrGenerateModel::init(path, None, None)?;
+        WhichModel::Qwen3ASR0_6B | WhichModel::Qwen3ASR1_7B => {
+            let model = Qwen3AsrGenerateModel::init(path, device, dtype)?;
            ModelInstance::Qwen3ASR(model)
        }
-        WhichModel::Qwen3ASR1_7B => {
-            let model = Qwen3AsrGenerateModel::init(path, None, None)?;
-            ModelInstance::Qwen3ASR(model)
+        WhichModel::Qwen3Embedding0_6B
+        | WhichModel::Qwen3Embedding4B
+        | WhichModel::Qwen3Embedding8B => {
+            let model = Qwen3Embedding::init(path, device, dtype)?;
+            ModelInstance::Qwen3Embedding(model)
        }
-        WhichModel::Qwen3VL2B => {
-            let model = Qwen3VLGenerateModel::init(path, None, None)?;
+        WhichModel::Qwen3VL2B
+        | WhichModel::Qwen3VL4B
+        | WhichModel::Qwen3VL8B
+        | WhichModel::Qwen3VL32B => {
+            let model = Qwen3VLGenerateModel::init(path, device, dtype)?;
            ModelInstance::Qwen3VL(Box::new(model))
        }
-        WhichModel::Qwen3VL4B => {
-            let model = Qwen3VLGenerateModel::init(path, None, None)?;
-            ModelInstance::Qwen3VL(Box::new(model))
-        }
-        WhichModel::Qwen3VL8B => {
-            let model = Qwen3VLGenerateModel::init(path, None, None)?;
-            ModelInstance::Qwen3VL(Box::new(model))
-        }
-        WhichModel::Qwen3VL32B => {
-            let model = Qwen3VLGenerateModel::init(path, None, None)?;
-            ModelInstance::Qwen3VL(Box::new(model))
-        }
-        WhichModel::DeepSeekOCR => {
-            let model = DeepseekOCRGenerateModel::init(path, None, None)?;
-            ModelInstance::DeepSeekOCR(model)
-        }
-        WhichModel::DeepSeekOCR2 => {
-            let model = DeepseekOCRGenerateModel::init(path, None, None)?;
+        WhichModel::DeepSeekOCR | WhichModel::DeepSeekOCR2 => {
+            let model = DeepseekOCRGenerateModel::init(path, device, dtype)?;
            ModelInstance::DeepSeekOCR(model)
        }
        WhichModel::HunyuanOCR => {
-            let model = HunyuanOCRGenerateModel::init(path, None, None)?;
+            let model = HunyuanOCRGenerateModel::init(path, device, dtype)?;
            ModelInstance::HunyuanOCR(model)
        }
-        WhichModel::PaddleOCRVL => {
-            let model = PaddleOCRVLGenerateModel::init(path, None, None)?;
-            ModelInstance::PaddleOCRVL(Box::new(model))
-        }
-        WhichModel::PaddleOCRVL1_5 => {
-            let model = PaddleOCRVLGenerateModel::init(path, None, None)?;
+        WhichModel::PaddleOCRVL | WhichModel::PaddleOCRVL1_5 => {
+            let model = PaddleOCRVLGenerateModel::init(path, device, dtype)?;
            ModelInstance::PaddleOCRVL(Box::new(model))
        }
        WhichModel::RMBG2_0 => {
-            let model = RMBG2_0Model::init(path, None, None)?;
+            let model = RMBG2_0Model::init(path, device, dtype)?;
            ModelInstance::RMBG2_0(Box::new(model))
        }
-        WhichModel::VoxCPM => {
-            let model = VoxCPMGenerate::init(path, None, None)?;
-            ModelInstance::VoxCPM(Box::new(model))
-        }
-        WhichModel::VoxCPM1_5 => {
-            let model = VoxCPMGenerate::init(path, None, None)?;
+        WhichModel::VoxCPM | WhichModel::VoxCPM1_5 => {
+            let model = VoxCPMGenerate::init(path, device, dtype)?;
            ModelInstance::VoxCPM(Box::new(model))
        }
        WhichModel::GlmASRNano2512 => {
-            let model = GlmAsrNanoGenerateModel::init(path, None, None)?;
+            let model = GlmAsrNanoGenerateModel::init(path, device, dtype)?;
            ModelInstance::GlmASRNano(model)
        }
        WhichModel::FunASRNano2512 => {
-            let model = FunAsrNanoGenerateModel::init(path, None, None)?;
+            let model = FunAsrNanoGenerateModel::init(path, device, dtype)?;
            ModelInstance::FunASRNano(model)
        }
        WhichModel::GlmOCR => {
-            let model = GlmOcrGenerateModel::init(path, None, None)?;
+            let model = GlmOcrGenerateModel::init(path, device, dtype)?;
            ModelInstance::GlmOCR(model)
        }
        _ => {
            let model_id = model_type.as_string();
-            return Err(anyhow!("model id {model_id} is not safetensor model"));
+            if model_id.to_lowercase().contains("gguf") || model_id.to_lowercase().contains("onnx")
+            {
+                return Err(anyhow!("model id {model_id} is not safetensor model"));
+            } else {
+                return Err(anyhow!("model id {model_id} not impl load_model function"));
+            }
        }
    };
    Ok(model)