From 7142e712f2bf7c94e2188660de34bce4cfb14c01 Mon Sep 17 00:00:00 2001 From: jhqxxx <18280426169@163.com> Date: Mon, 6 Apr 2026 15:35:49 +0800 Subject: [PATCH] add qwen3-embedding and all-minilm-l6-v2 --- src/cli/mod.rs | 77 +++----- src/exec/all_minilm_l6_v2.rs | 41 ++++ src/exec/mod.rs | 2 + src/exec/qwen3_embedding.rs | 42 +++++ src/main.rs | 2 + src/models/all_minilm_l6_v2/mod.rs | 78 ++++++++ src/models/campplus/mod.rs | 5 +- src/models/common/embedding.rs | 29 +++ src/models/common/mod.rs | 1 + src/models/common/model_mapping.rs | 12 ++ src/models/common/modules.rs | 137 +++++++++++++- src/models/deepseek_ocr/model.rs | 4 +- .../feature_extraction_whisper.rs | 9 +- .../seamless_m4t_feature_extractor.rs | 9 +- src/models/mask_gct/model.rs | 4 +- src/models/mod.rs | 176 +++++++++--------- src/models/qwen3/model.rs | 21 ++- src/models/qwen3_5/model.rs | 6 +- src/models/qwen3_asr/processor.rs | 5 +- src/models/qwen3_embedding/mod.rs | 66 +++++++ src/{server/asr_types.rs => params/asr.rs} | 30 --- src/params/embedding.rs | 22 +++ src/params/mod.rs | 6 + src/params/rerank.rs | 23 +++ src/params/shared.rs | 2 +- src/server/api.rs | 4 +- src/server/asr.rs | 4 +- src/server/embedding.rs | 77 ++++++++ src/server/mod.rs | 5 +- src/utils/audio_utils.rs | 3 +- src/utils/tensor_utils.rs | 114 ------------ tests/test_all_minilm_l6_v2.rs | 24 +++ tests/test_qwen3_embedding.rs | 22 +++ tests/weight_test.rs | 20 +- 34 files changed, 763 insertions(+), 319 deletions(-) create mode 100644 src/exec/all_minilm_l6_v2.rs create mode 100644 src/exec/qwen3_embedding.rs create mode 100644 src/models/all_minilm_l6_v2/mod.rs create mode 100644 src/models/common/embedding.rs create mode 100644 src/models/qwen3_embedding/mod.rs rename src/{server/asr_types.rs => params/asr.rs} (57%) create mode 100644 src/params/embedding.rs create mode 100644 src/params/rerank.rs create mode 100644 src/server/embedding.rs create mode 100644 tests/test_all_minilm_l6_v2.rs create mode 100644 tests/test_qwen3_embedding.rs diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 221186a..3cc26f6 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -238,46 +238,28 @@ pub(crate) fn run_run(args: RunArgs) -> anyhow::Result<()> { None => get_default_weight_path(model), }; match model { + WhichModel::AllMiniLML6V2 => { + all_minilm_l6_v2::AllMiniLML6V2Exec::run(&input, output.as_deref(), &weight_path)?; + } WhichModel::MiniCPM4_0_5B => { minicpm4::MiniCPM4Exec::run(&input, output.as_deref(), &weight_path)?; } - WhichModel::LFM2_1_2B => { + WhichModel::LFM2_1_2B | WhichModel::LFM2_5_1_2BInstruct => { lfm2::Lfm2Exec::run(&input, output.as_deref(), &weight_path)?; } - WhichModel::LFM2_5_1_2BInstruct => { - lfm2::Lfm2Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::LFM2_5VL1_6B => { + WhichModel::LFM2_5VL1_6B | WhichModel::LFM2VL1_6B => { lfm2vl::Lfm2VLExec::run(&input, output.as_deref(), &weight_path)?; } - WhichModel::LFM2VL1_6B => { - lfm2vl::Lfm2VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen2_5VL3B => { + WhichModel::Qwen2_5VL3B | WhichModel::Qwen2_5VL7B => { qwen2_5vl::Qwen2_5VLExec::run(&input, output.as_deref(), &weight_path)?; } - WhichModel::Qwen2_5VL7B => { - qwen2_5vl::Qwen2_5VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_0_6B => { + WhichModel::Qwen3_0_6B | WhichModel::Qwen3_1_7B | WhichModel::Qwen3_4B => { qwen3::Qwen3Exec::run(&input, output.as_deref(), &weight_path)?; } - WhichModel::Qwen3_1_7B => { - qwen3::Qwen3Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_4B => { - qwen3::Qwen3Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_5_0_8B => { - qwen3_5::Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_5_2B => { - qwen3_5::Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_5_4B => { - qwen3_5::Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_5_9B => { + WhichModel::Qwen3_5_0_8B + | WhichModel::Qwen3_5_2B + | WhichModel::Qwen3_5_4B + | WhichModel::Qwen3_5_9B => { qwen3_5::Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; } WhichModel::Qwen3_5Gguf => { @@ -288,46 +270,33 @@ pub(crate) fn run_run(args: RunArgs) -> anyhow::Result<()> { path_common.mmproj_path, )?; } - WhichModel::Qwen3ASR0_6B => { + WhichModel::Qwen3ASR0_6B | WhichModel::Qwen3ASR1_7B => { qwen3_asr::Qwen3ASRExec::run(&input, output.as_deref(), &weight_path)?; } - WhichModel::Qwen3ASR1_7B => { - qwen3_asr::Qwen3ASRExec::run(&input, output.as_deref(), &weight_path)?; + WhichModel::Qwen3Embedding0_6B + | WhichModel::Qwen3Embedding4B + | WhichModel::Qwen3Embedding8B => { + qwen3_embedding::Qwen3EmbeddingExec::run(&input, output.as_deref(), &weight_path)?; } - WhichModel::Qwen3VL2B => { + WhichModel::Qwen3VL2B + | WhichModel::Qwen3VL4B + | WhichModel::Qwen3VL8B + | WhichModel::Qwen3VL32B => { qwen3vl::Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; } - WhichModel::Qwen3VL4B => { - qwen3vl::Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3VL8B => { - qwen3vl::Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3VL32B => { - qwen3vl::Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::DeepSeekOCR => { - deepseek_ocr::DeepSeekORExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::DeepSeekOCR2 => { + WhichModel::DeepSeekOCR | WhichModel::DeepSeekOCR2 => { deepseek_ocr::DeepSeekORExec::run(&input, output.as_deref(), &weight_path)?; } WhichModel::HunyuanOCR => { hunyuan_ocr::HunyuanORExec::run(&input, output.as_deref(), &weight_path)?; } - WhichModel::PaddleOCRVL => { - paddleocr_vl::PaddleOVLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::PaddleOCRVL1_5 => { + WhichModel::PaddleOCRVL | WhichModel::PaddleOCRVL1_5 => { paddleocr_vl::PaddleOVLExec::run(&input, output.as_deref(), &weight_path)?; } WhichModel::RMBG2_0 => { rmbg2_0::RMBG2_0Exec::run(&input, output.as_deref(), &weight_path)?; } - WhichModel::VoxCPM => { - voxcpm::VoxCPMExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::VoxCPM1_5 => { + WhichModel::VoxCPM | WhichModel::VoxCPM1_5 => { voxcpm::VoxCPMExec::run(&input, output.as_deref(), &weight_path)?; } WhichModel::GlmASRNano2512 => { diff --git a/src/exec/all_minilm_l6_v2.rs b/src/exec/all_minilm_l6_v2.rs new file mode 100644 index 0000000..ffc0993 --- /dev/null +++ b/src/exec/all_minilm_l6_v2.rs @@ -0,0 +1,41 @@ +use std::time::Instant; + +use crate::{ + exec::ExecModel, + models::{all_minilm_l6_v2::AllMiniLML6V2Embedding, common::embedding::TextEmbedding}, + utils::get_file_path, +}; +use anyhow::Result; + +pub struct AllMiniLML6V2Exec; + +impl ExecModel for AllMiniLML6V2Exec { + fn run(input: &[String], output: Option<&str>, weight_path: &str) -> Result<()> { + let input_text = &input[0]; + let target_text = if input_text.starts_with("file://") { + let path = get_file_path(input_text)?; + std::fs::read_to_string(path)? + } else { + input_text.clone() + }; + + let i_start = Instant::now(); + let mut model = AllMiniLML6V2Embedding::init(weight_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + + let i_start = Instant::now(); + let result = model.embed_texts(&[target_text])?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in generate is: {:?}", i_duration); + + println!("Result: {:?}", result); + + if let Some(out) = output { + std::fs::write(out, format!("{:?}", result))?; + println!("Output saved to: {}", out); + } + + Ok(()) + } +} diff --git a/src/exec/mod.rs b/src/exec/mod.rs index 4939774..9681442 100644 --- a/src/exec/mod.rs +++ b/src/exec/mod.rs @@ -3,6 +3,7 @@ //! This module provides model-specific exec implementations for the `run` subcommand. //! Each model has its own exec module that handles input/output parsing and model invocation. +pub mod all_minilm_l6_v2; pub mod deepseek_ocr; pub mod fun_asr_nano; pub mod glm_asr_nano; @@ -16,6 +17,7 @@ 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; diff --git a/src/exec/qwen3_embedding.rs b/src/exec/qwen3_embedding.rs new file mode 100644 index 0000000..741f7ca --- /dev/null +++ b/src/exec/qwen3_embedding.rs @@ -0,0 +1,42 @@ +use std::time::Instant; + +use crate::{ + exec::ExecModel, + models::{common::embedding::TextEmbedding, qwen3_embedding::Qwen3Embedding}, + utils::get_file_path, +}; +use anyhow::Result; + +pub struct Qwen3EmbeddingExec; + +impl ExecModel for Qwen3EmbeddingExec { + fn run(input: &[String], output: Option<&str>, weight_path: &str) -> Result<()> { + let input_text = &input[0]; + let target_text = if input_text.starts_with("file://") { + // let path = &input[7..]; + let path = get_file_path(input_text)?; + std::fs::read_to_string(path)? + } else { + input_text.clone() + }; + + let i_start = Instant::now(); + let mut model = Qwen3Embedding::init(weight_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + + let i_start = Instant::now(); + let result = model.embed_texts(&[target_text])?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in generate is: {:?}", i_duration); + + println!("Result: {:?}", result); + + if let Some(out) = output { + std::fs::write(out, format!("{:?}", result))?; + println!("Output saved to: {}", out); + } + + Ok(()) + } +} diff --git a/src/main.rs b/src/main.rs index b8649e2..0615222 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,6 +6,8 @@ use crate::cli::{ }; mod cli; +#[allow(unused)] +mod params; mod server; #[tokio::main] diff --git a/src/models/all_minilm_l6_v2/mod.rs b/src/models/all_minilm_l6_v2/mod.rs new file mode 100644 index 0000000..4a4eeb0 --- /dev/null +++ b/src/models/all_minilm_l6_v2/mod.rs @@ -0,0 +1,78 @@ +use crate::{ + models::common::embedding::{NormalizeType, TextEmbedding}, + tokenizer::TokenizerModel, + utils::{find_type_files, get_device, get_dtype}, +}; +use anyhow::{Result, anyhow}; +use candle_core::{DType, Device, Tensor}; +use candle_nn::VarBuilder; +use candle_transformers::models::bert::{BertModel, Config as BertConfig}; + +pub struct AllMiniLML6V2Embedding { + tokenizer: TokenizerModel, + model: BertModel, + device: Device, + normalize: NormalizeType, +} + +impl AllMiniLML6V2Embedding { + pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { + let tokenizer = TokenizerModel::init(path)?; + let config_path = path.to_string() + "/config.json"; + let cfg: BertConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; + let device = get_device(device); + let dtype = get_dtype(dtype, "float32"); + let model_list = find_type_files(path, "safetensors")?; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, &device)? }; + let model = BertModel::load(vb, &cfg)?; + Ok(Self { + tokenizer, + model, + device, + normalize: NormalizeType::L2, + }) + } + + fn prepare_token_ids(&self, text: &str) -> Result> { + let mut token_ids = self.tokenizer.text_encode_vec(text.to_string(), true)?; + token_ids = token_ids + .into_iter() + .filter(|&x| x != 0) + .collect::>(); + if token_ids.is_empty() { + return Err(anyhow!("embedding tokenized input cannot be empty")); + } + Ok(token_ids) + } + fn embed_one(&mut self, text: &str) -> Result> { + let token_ids = self.prepare_token_ids(text)?; + let seq_len = token_ids.len(); + let input_ids = Tensor::from_slice(&token_ids, (1, seq_len), &self.device)?; + let token_type_ids = Tensor::zeros((1, seq_len), DType::U32, &self.device)?; + let attention_mask = Tensor::ones((1, seq_len), DType::U32, &self.device)?; + let hidden = self + .model + .forward(&input_ids, &token_type_ids, Some(&attention_mask))? + .to_dtype(DType::F32)?; + let hidden = hidden.mean(1)?; + let embed = self + .normalize + .normalize(&hidden, hidden.rank() - 1)? + .squeeze(0)?; + let embed = embed.to_vec1::()?; + Ok(embed) + } +} + +impl TextEmbedding for AllMiniLML6V2Embedding { + fn embed_texts(&mut self, input: &[String]) -> Result>> { + if input.is_empty() { + return Err(anyhow!("embedding input cannot be empty")); + } + let mut out = Vec::with_capacity(input.len()); + for text in input { + out.push(self.embed_one(text)?); + } + Ok(out) + } +} diff --git a/src/models/campplus/mod.rs b/src/models/campplus/mod.rs index a60d343..779cfd9 100644 --- a/src/models/campplus/mod.rs +++ b/src/models/campplus/mod.rs @@ -2,9 +2,8 @@ use anyhow::Result; use candle_core::{D, Tensor}; use candle_nn::{BatchNorm, Conv1d, Conv2d, Module, ModuleT, VarBuilder, ops::sigmoid}; -use crate::{ - models::common::modules::{get_batch_norm, get_conv1d, get_conv2d}, - utils::tensor_utils::{pool1d, statistics_pooling}, +use crate::models::common::modules::{ + get_batch_norm, get_conv1d, get_conv2d, pool1d, statistics_pooling, }; pub struct Shortcut { diff --git a/src/models/common/embedding.rs b/src/models/common/embedding.rs new file mode 100644 index 0000000..5b2da4e --- /dev/null +++ b/src/models/common/embedding.rs @@ -0,0 +1,29 @@ +use anyhow::Result; +use candle_core::Tensor; + +use crate::models::common::modules::{ + l1_normalize, l2_normalize, max_abs_normalize, min_max_normalize, z_score_normalize, +}; +pub trait TextEmbedding { + fn embed_texts(&mut self, input: &[String]) -> Result>>; +} + +pub enum NormalizeType { + L1, + L2, + ZScore, + MinMax, + MaxAbs, +} + +impl NormalizeType { + pub fn normalize(&self, t: &Tensor, dim: usize) -> Result { + match self { + NormalizeType::L1 => l1_normalize(t, dim), + NormalizeType::L2 => l2_normalize(t, dim), + NormalizeType::ZScore => z_score_normalize(t, dim), + NormalizeType::MinMax => min_max_normalize(t, dim), + NormalizeType::MaxAbs => max_abs_normalize(t, dim), + } + } +} diff --git a/src/models/common/mod.rs b/src/models/common/mod.rs index 19c19a0..33aa1ec 100644 --- a/src/models/common/mod.rs +++ b/src/models/common/mod.rs @@ -1,5 +1,6 @@ use anyhow::Result; use candle_core::Tensor; +pub mod embedding; pub mod generate; pub mod gguf; pub mod model_mapping; diff --git a/src/models/common/model_mapping.rs b/src/models/common/model_mapping.rs index 291ac1f..ec747ee 100644 --- a/src/models/common/model_mapping.rs +++ b/src/models/common/model_mapping.rs @@ -2,6 +2,8 @@ use clap::ValueEnum; #[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] pub enum WhichModel { + #[value(name = "sentence-transformers/all-MiniLM-L6-v2")] + AllMiniLML6V2, #[value(name = "LiquidAI/LFM2-1.2B")] LFM2_1_2B, #[value(name = "LiquidAI/LFM2.5-1.2B-Instruct")] @@ -36,6 +38,12 @@ pub enum WhichModel { Qwen3ASR0_6B, #[value(name = "Qwen/Qwen3-ASR-1.7B")] Qwen3ASR1_7B, + #[value(name = "Qwen/Qwen3-Embedding-0.6B")] + Qwen3Embedding0_6B, + #[value(name = "Qwen/Qwen3-Embedding-4B")] + Qwen3Embedding4B, + #[value(name = "Qwen/Qwen3-Embedding-8B")] + Qwen3Embedding8B, #[value(name = "Qwen/Qwen3-VL-2B-Instruct")] Qwen3VL2B, #[value(name = "Qwen/Qwen3-VL-4B-Instruct")] @@ -153,6 +161,10 @@ impl WhichModel { WhichModel::RMBG2_0 => "image", // TTS models WhichModel::VoxCPM | WhichModel::VoxCPM1_5 => "tts", + WhichModel::Qwen3Embedding0_6B + | WhichModel::Qwen3Embedding4B + | WhichModel::Qwen3Embedding8B + | WhichModel::AllMiniLML6V2 => "embedding", } } } diff --git a/src/models/common/modules.rs b/src/models/common/modules.rs index 684fcbc..cad29e3 100644 --- a/src/models/common/modules.rs +++ b/src/models/common/modules.rs @@ -1,5 +1,5 @@ -use anyhow::Result; -use candle_core::{D, IndexOp, Tensor}; +use anyhow::{Result, anyhow}; +use candle_core::{D, DType, IndexOp, Tensor}; use candle_nn::{ Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig, ConvTranspose1d, ConvTranspose1dConfig, Embedding, LayerNorm, LayerNormConfig, Linear, Module, @@ -9,7 +9,7 @@ use candle_nn::{ use crate::{ position_embed::rope::{RoPE, apply_rotary_pos_emb, apply_rotary_pos_emb_roformer}, - utils::tensor_utils::{prepare_causal_attention_mask, repeat_kv}, + utils::tensor_utils::{pad_replicate_last_dim, prepare_causal_attention_mask, repeat_kv}, }; #[derive(Debug, Clone)] @@ -1327,3 +1327,134 @@ pub fn conv1d_depthwise(input: &Tensor, weight: &Tensor, bias: Option<&Tensor>) } } } + +pub fn log10(t: &Tensor) -> Result { + Ok(t.log()?.affine(1.0 / 10.0_f64.ln(), 0.0)?) +} + +pub fn max_abs_normalize(t: &Tensor, dim: usize) -> Result { + let rank = t.rank(); + if dim >= rank { + return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank))); + } + Ok(t.broadcast_div(&t.abs()?.max_keepdim(dim)?)?) +} + +pub fn min_max_normalize(t: &Tensor, dim: usize) -> Result { + let rank = t.rank(); + if dim >= rank { + return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank))); + } + let t_min = t.min_keepdim(dim)?; + Ok(t.broadcast_sub(&t_min)? + .broadcast_div(&t.max_keepdim(dim)?.sub(&t_min)?)?) +} + +pub fn z_score_normalize(t: &Tensor, dim: usize) -> Result { + let rank = t.rank(); + if dim >= rank { + return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank))); + } + Ok(t.broadcast_sub(&t.mean_keepdim(dim)?)? + .broadcast_div(&t.var_keepdim(dim)?.sqrt()?)?) +} + +pub fn l2_normalize(t: &Tensor, dim: usize) -> Result { + let rank = t.rank(); + if dim >= rank { + return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank))); + } + let l2_norm = t.sqr()?.sum_keepdim(dim)?.affine(1.0, 1e-6)?.sqrt()?; + Ok(t.broadcast_div(&l2_norm)?) +} + +pub fn l1_normalize(t: &Tensor, dim: usize) -> Result { + let rank = t.rank(); + if dim >= rank { + return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank))); + } + let l1_norm = t.abs()?.sum_keepdim(dim)?; + Ok(t.broadcast_div(&l1_norm)?) +} + +pub fn pool1d(xs: &Tensor, pool_size: usize, ceil_mode: bool, stype: &str) -> Result { + // xs: (bs, c, dim) + // ceil_mode: 是否保留不完整窗口，为true时通过pad实现 + if pool_size == 0 { + return Err(anyhow!("pool_size must be greater than 0")); + } + let (bs, c, dim) = xs.dims3()?; + let xs_reshape = if ceil_mode { + let remain = dim % pool_size; + if remain > 0 { + let pad = pool_size - remain; + let xs_pad = pad_replicate_last_dim(xs, (0, pad))?; + xs_pad.reshape((bs, c, (), pool_size))? + } else { + xs.reshape((bs, c, (), pool_size))? + } + } else { + let remain = dim % pool_size; + if remain > 0 { + let xs_del = xs.narrow(D::Minus1, 0, dim - remain)?; + xs_del.reshape((bs, c, (), pool_size))? + } else { + xs.reshape((bs, c, (), pool_size))? + } + }; + let xs_pool = match stype { + "avg" => xs_reshape.mean(D::Minus1)?, + "max" => xs_reshape.max(D::Minus1)?, + "min" => xs_reshape.min(D::Minus1)?, + _ => { + return Err(anyhow!( + "unsupported pool type: {}, supported types are: avg, max, min", + stype + )); + } + }; + Ok(xs_pool) +} + +pub fn statistics_pooling(xs: &Tensor, dim: D, keepdim: bool) -> Result { + let mean = xs.mean(dim)?; + let std = xs.var(dim)?.sqrt()?; + let mut stats = Tensor::cat(&[mean, std], D::Minus1)?; + if keepdim { + stats = stats.unsqueeze(dim)?; + } + Ok(stats) +} +pub fn float_range_normalize(t: &Tensor) -> Result { + let peak = t + .to_dtype(DType::F32)? + .abs()? + .max_all()? + .to_scalar::()?; + if peak == 0.0 { + return Ok(t.clone()); + } + let mut t = t.clone(); + if peak > 1.0 { + t = t.affine(1.0 / peak as f64, 0.0)?; + } + t = t.clamp(-1.0, 1.0)?; + Ok(t) +} + +pub fn cosine_similarity(query_vector: &Tensor, matrix: &Tensor) -> Result { + // query_vector: (n, dim) + // matrix: (m, dim) + let query_norm = l2_normalize(query_vector, query_vector.rank() - 1)?; + let matrix_norm = l2_normalize(matrix, matrix.rank() - 1)?; + let similarity = query_norm + .matmul(&matrix_norm.transpose(D::Minus1, D::Minus2)?)? + .squeeze(D::Minus1)?; + Ok(similarity) +} + +pub fn quick_gelu(xs: &Tensor) -> Result { + let x = xs.affine(1.702, 0.0)?; + let x = sigmoid(&x)?; + Ok(xs.mul(&x)?) +} diff --git a/src/models/deepseek_ocr/model.rs b/src/models/deepseek_ocr/model.rs index 6955e1e..5e9ded8 100644 --- a/src/models/deepseek_ocr/model.rs +++ b/src/models/deepseek_ocr/model.rs @@ -16,7 +16,7 @@ use crate::{ InferenceModel, modules::{ GateUpDownMLP, NaiveAttention, QKVCatAttention, TwoLinearMLP, - eager_attention_forward, get_conv2d, get_layer_norm, + eager_attention_forward, get_conv2d, get_layer_norm, quick_gelu, }, }, deepseek_ocr::config::{DeepseekOCRConfig, DeepseekV2Config}, @@ -27,7 +27,7 @@ use crate::{ interpolate::{interpolate_bicubic, interpolate_linear_1d}, tensor_utils::{ attn_masked_fill, index_select_2d, masked_scatter_dim0, nonzero, onehot, - prepare_causal_attention_mask, quick_gelu, topk, + prepare_causal_attention_mask, topk, }, }, }; diff --git a/src/models/feature_extractor/feature_extraction_whisper.rs b/src/models/feature_extractor/feature_extraction_whisper.rs index 507005e..2a67f3b 100644 --- a/src/models/feature_extractor/feature_extraction_whisper.rs +++ b/src/models/feature_extractor/feature_extraction_whisper.rs @@ -1,9 +1,12 @@ use anyhow::Result; use candle_core::{D, Device, Tensor}; -use crate::utils::{ - audio_utils::{create_hann_window, mel_filter_bank, torch_stft}, - tensor_utils::{log10, pad_reflect_last_dim}, +use crate::{ + models::common::modules::log10, + utils::{ + audio_utils::{create_hann_window, mel_filter_bank, torch_stft}, + tensor_utils::pad_reflect_last_dim, + }, }; pub struct WhisperFeatureExtractor { diff --git a/src/models/feature_extractor/seamless_m4t_feature_extractor.rs b/src/models/feature_extractor/seamless_m4t_feature_extractor.rs index 12e23e9..c9315ee 100644 --- a/src/models/feature_extractor/seamless_m4t_feature_extractor.rs +++ b/src/models/feature_extractor/seamless_m4t_feature_extractor.rs @@ -1,9 +1,12 @@ use anyhow::Result; use candle_core::{D, Device, Tensor}; -use crate::utils::{ - audio_utils::{create_povey_window, mel_filter_bank, spectrogram}, - tensor_utils::{PaddingSide, z_score_normalize}, +use crate::{ + models::common::modules::z_score_normalize, + utils::{ + audio_utils::{create_povey_window, mel_filter_bank, spectrogram}, + tensor_utils::PaddingSide, + }, }; pub struct SeamlessM4TFeatureExtractor { diff --git a/src/models/mask_gct/model.rs b/src/models/mask_gct/model.rs index bc3ae5c..bd2bb04 100644 --- a/src/models/mask_gct/model.rs +++ b/src/models/mask_gct/model.rs @@ -6,10 +6,10 @@ use candle_nn::{ use crate::{ models::{ - common::modules::{WNConv1d, conv1d_depthwise, get_conv1d, get_layer_norm}, + common::modules::{WNConv1d, conv1d_depthwise, get_conv1d, get_layer_norm, l2_normalize}, mask_gct::config::SemanticCodec, }, - utils::{interpolate::interpolate_nearest_1d, tensor_utils::l2_normalize}, + utils::interpolate::interpolate_nearest_1d, }; pub struct ConvNeXtBlock { diff --git a/src/models/mod.rs b/src/models/mod.rs index aad8123..b7b18e9 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -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>), 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 { 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>> { + 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> { 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> { +pub fn load_model<'a>( + model_type: WhichModel, + path: &str, + device: Option<&Device>, + dtype: Option, +) -> Result> { 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) diff --git a/src/models/qwen3/model.rs b/src/models/qwen3/model.rs index ef82def..8703806 100644 --- a/src/models/qwen3/model.rs +++ b/src/models/qwen3/model.rs @@ -102,7 +102,11 @@ pub struct Qwen3Model { impl Qwen3Model { pub fn new(config: &Qwen3Config, vb: VarBuilder, eos_ids: Vec) -> Result { - let vb = vb.pp("model"); + let vb = if vb.contains_tensor("model.embed_tokens.weight") { + vb.pp("model") + } else { + vb + }; let vocab_size = config.vocab_size; let embed_tokens = embedding(vocab_size, config.hidden_size, vb.pp("embed_tokens"))?; let mut layers = vec![]; @@ -133,6 +137,17 @@ impl Qwen3Model { input_ids: Option<&Tensor>, inputs_embeds: Option<&Tensor>, seqlen_offset: usize, + ) -> Result { + let hidden_state = self.forward_hidden(input_ids, inputs_embeds, seqlen_offset)?; + let logits = self.lm_head.forward(&hidden_state)?; + Ok(logits) + } + + pub fn forward_hidden( + &mut self, + input_ids: Option<&Tensor>, + inputs_embeds: Option<&Tensor>, + seqlen_offset: usize, ) -> Result { if input_ids.is_none() && inputs_embeds.is_none() { return Err(anyhow::anyhow!( @@ -170,9 +185,9 @@ impl Qwen3Model { } hidden_states = self.norm.forward(&hidden_states)?; let hidden_state = hidden_states.narrow(1, seq_len - 1, 1)?; - let logits = self.lm_head.forward(&hidden_state)?; - Ok(logits) + Ok(hidden_state) } + pub fn embedding_token_id(&self, input_ids: &Tensor) -> Result { Ok(self.embed_tokens.forward(input_ids)?) } diff --git a/src/models/qwen3_5/model.rs b/src/models/qwen3_5/model.rs index 7b0ffa9..49429e0 100644 --- a/src/models/qwen3_5/model.rs +++ b/src/models/qwen3_5/model.rs @@ -12,14 +12,16 @@ use crate::{ common::{ InferenceModel, gguf::{GateUpDownMLPGguf, Gguf, ProjKind, QuantizedLinear}, - modules::{conv1d_depthwise, eager_attention_forward, get_conv1d, softplus}, + modules::{ + conv1d_depthwise, eager_attention_forward, get_conv1d, l2_normalize, softplus, + }, }, qwen3_5::config::{Qwen3_5Config, Qwen3_5TextConfig}, qwen3vl::model::Qwen3VLVisionModel, }, position_embed::rope::{Qwen3VLTextRotaryEmbedding, glm_asr_apply_rotary_pos_emb}, utils::tensor_utils::{ - get_equal_mask, get_vision_next_indices, l2_normalize, masked_scatter_dim0, nonzero_index, + get_equal_mask, get_vision_next_indices, masked_scatter_dim0, nonzero_index, prepare_causal_attention_mask, repeat_interleave, split_tensor, zero_index, }, }; diff --git a/src/models/qwen3_asr/processor.rs b/src/models/qwen3_asr/processor.rs index 1c7c41b..08499b8 100644 --- a/src/models/qwen3_asr/processor.rs +++ b/src/models/qwen3_asr/processor.rs @@ -1,4 +1,6 @@ -use crate::params::chat::ChatCompletionParameters; +use crate::{ + models::common::modules::float_range_normalize, params::chat::ChatCompletionParameters, +}; use anyhow::Result; use candle_core::{Device, Tensor}; @@ -10,7 +12,6 @@ use crate::{ utils::{ audio_utils::{extract_audios, split_audio_into_chunks}, capitalize_first_letter, - tensor_utils::float_range_normalize, }, }; diff --git a/src/models/qwen3_embedding/mod.rs b/src/models/qwen3_embedding/mod.rs new file mode 100644 index 0000000..04bc9e6 --- /dev/null +++ b/src/models/qwen3_embedding/mod.rs @@ -0,0 +1,66 @@ +use crate::{ + models::{ + common::embedding::{NormalizeType, TextEmbedding}, + qwen3::{config::Qwen3Config, model::Qwen3Model}, + }, + tokenizer::TokenizerModel, + utils::{find_type_files, get_device, get_dtype}, +}; +use anyhow::{Result, anyhow}; +use candle_core::{DType, Device}; +use candle_nn::VarBuilder; + +pub struct Qwen3Embedding { + tokenizer: TokenizerModel, + model: Qwen3Model, + device: Device, + normalize: NormalizeType, +} + +impl Qwen3Embedding { + pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { + let tokenizer = TokenizerModel::init(path)?; + let config_path = path.to_string() + "/config.json"; + let cfg: Qwen3Config = serde_json::from_slice(&std::fs::read(config_path)?)?; + let device = get_device(device); + let dtype = get_dtype(dtype, cfg.torch_dtype.as_str()); + let model_list = find_type_files(path, "safetensors")?; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, &device)? }; + let model = Qwen3Model::new(&cfg, vb, vec![])?; + Ok(Self { + tokenizer, + model, + device, + normalize: NormalizeType::L2, + }) + } + + fn embed_one(&mut self, text: &str) -> Result> { + let input_ids = self.tokenizer.text_encode(text.to_string(), &self.device)?; + let hidden = self + .model + .forward_hidden(Some(&input_ids), None, 0)? + .squeeze(0)? + .to_dtype(DType::F32)?; + let norm = self + .normalize + .normalize(&hidden, hidden.rank() - 1)? + .squeeze(0)?; + let norm = norm.to_vec1::()?; + Ok(norm) + } +} + +impl TextEmbedding for Qwen3Embedding { + fn embed_texts(&mut self, input: &[String]) -> Result>> { + if input.is_empty() { + return Err(anyhow!("embedding input cannot be empty")); + } + let mut out = Vec::with_capacity(input.len()); + for text in input { + out.push(self.embed_one(text)?); + self.model.clear_kv_cache(); + } + Ok(out) + } +} diff --git a/src/server/asr_types.rs b/src/params/asr.rs similarity index 57% rename from src/server/asr_types.rs rename to src/params/asr.rs index 4ef26df..3fade9a 100644 --- a/src/server/asr_types.rs +++ b/src/params/asr.rs @@ -46,33 +46,3 @@ pub(crate) struct ErrorDetail { #[serde(skip_serializing_if = "Option::is_none")] pub(crate) code: Option, } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_transcription_response_serialization() { - let response = TranscriptionResponse { - text: "Hello, world!".to_string(), - }; - let json = serde_json::to_string(&response).unwrap(); - assert_eq!(json, r#"{"text":"Hello, world!"}"#); - } - - #[test] - fn test_error_response_serialization() { - let error = ErrorResponse { - error: ErrorDetail { - message: "Invalid audio file".to_string(), - error_type: "invalid_request_error".to_string(), - code: Some("invalid_audio".to_string()), - }, - }; - let json = serde_json::to_string(&error).unwrap(); - let parsed: serde_json::Value = serde_json::from_str(&json).unwrap(); - assert_eq!(parsed["error"]["message"], "Invalid audio file"); - assert_eq!(parsed["error"]["type"], "invalid_request_error"); - assert_eq!(parsed["error"]["code"], "invalid_audio"); - } -} diff --git a/src/params/embedding.rs b/src/params/embedding.rs new file mode 100644 index 0000000..8833f72 --- /dev/null +++ b/src/params/embedding.rs @@ -0,0 +1,22 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Debug, Deserialize)] +pub(crate) struct EmbeddingRequest { + pub model: Option, + pub input: Value, +} + +#[derive(Debug, Serialize)] +pub(crate) struct EmbeddingData { + pub object: String, + pub index: usize, + pub embedding: Vec, +} + +#[derive(Debug, Serialize)] +pub(crate) struct EmbeddingResponse { + pub object: String, + pub model: String, + pub data: Vec, +} diff --git a/src/params/mod.rs b/src/params/mod.rs index 399aa97..6e1083a 100644 --- a/src/params/mod.rs +++ b/src/params/mod.rs @@ -1,2 +1,8 @@ +#[allow(unused)] +pub mod asr; pub mod chat; +#[allow(unused)] +pub mod embedding; +#[allow(unused)] +pub mod rerank; pub mod shared; diff --git a/src/params/rerank.rs b/src/params/rerank.rs new file mode 100644 index 0000000..303ab70 --- /dev/null +++ b/src/params/rerank.rs @@ -0,0 +1,23 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize)] +pub(crate) struct RerankRequest { + pub model: Option, + pub query: String, + pub documents: Vec, + pub top_n: Option, +} + +#[derive(Debug, Serialize)] +struct RerankResult { + index: usize, + relevance_score: f32, + document: String, +} + +#[derive(Debug, Serialize)] +struct RerankResponse { + object: String, + model: String, + results: Vec, +} diff --git a/src/params/shared.rs b/src/params/shared.rs index d675487..27a450d 100644 --- a/src/params/shared.rs +++ b/src/params/shared.rs @@ -179,7 +179,7 @@ pub enum FinishReason { EndTurn, /// The finish reason is unspecified. [Gemini] #[serde(rename = "FINISH_REASON_UNSPECIFIED ")] - FinishReasonUnspecified, + Unspecified, #[serde(rename = "MALFORMED_FUNCTION_CALL")] MalformedFunctionCall, #[serde(rename = "OTHER")] diff --git a/src/server/api.rs b/src/server/api.rs index 7dee106..38a0d83 100644 --- a/src/server/api.rs +++ b/src/server/api.rs @@ -43,7 +43,7 @@ pub fn init( if let Some(gguf_path) = gguf { let gguf_path = string_to_static_str(gguf_path); let mmproj_path = mmproj.map(string_to_static_str); - load_gguf_model(model_type, None, gguf_path, mmproj_path)? + load_gguf_model(model_type, None, gguf_path, mmproj_path, None)? } else { return Err(anyhow!("gguf model need gguf model path")); } @@ -51,7 +51,7 @@ pub fn init( return Err(anyhow!("onnx comming soon but now not support")); } else { let model_path = string_to_static_str(path); - load_model(model_type, model_path)? + load_model(model_type, model_path, None, None)? }; MODEL.get_or_init(|| { diff --git a/src/server/asr.rs b/src/server/asr.rs index 26be045..cfd8d81 100644 --- a/src/server/asr.rs +++ b/src/server/asr.rs @@ -11,10 +11,8 @@ use rocket::http::Status; use rocket::serde::json::Json; use rocket::{form::Form, post}; +use crate::params::asr::{ErrorDetail, ErrorResponse, TranscriptionRequest, TranscriptionResponse}; use crate::server::api::MODEL; -use crate::server::asr_types::{ - ErrorDetail, ErrorResponse, TranscriptionRequest, TranscriptionResponse, -}; /// Handle audio transcription requests /// diff --git a/src/server/embedding.rs b/src/server/embedding.rs new file mode 100644 index 0000000..3b5c33d --- /dev/null +++ b/src/server/embedding.rs @@ -0,0 +1,77 @@ +use rocket::{http::Status, post, serde::json::Json}; +use serde_json::Value; + +use crate::{ + params::embedding::{EmbeddingData, EmbeddingRequest, EmbeddingResponse}, + server::api::MODEL, +}; + +fn parse_embedding_input(input: &Value) -> anyhow::Result> { + match input { + Value::String(s) => Ok(vec![s.clone()]), + Value::Array(arr) => { + let mut out = Vec::with_capacity(arr.len()); + for v in arr { + let s = v.as_str().ok_or_else(|| { + anyhow::anyhow!("embedding input array must contain only strings") + })?; + out.push(s.to_string()); + } + if out.is_empty() { + return Err(anyhow::anyhow!("embedding input cannot be empty")); + } + Ok(out) + } + _ => Err(anyhow::anyhow!( + "embedding input must be a string or an array of strings" + )), + } +} + +#[post("/embeddings", data = "")] +pub(crate) async fn embeddings(req: Json) -> (Status, Json) { + let texts = match parse_embedding_input(&req.input) { + Ok(v) => v, + Err(e) => { + return ( + Status::BadRequest, + Json(serde_json::json!({ "error": e.to_string() })), + ); + } + }; + let model_ref = match MODEL.get().cloned() { + Some(v) => v, + None => { + return ( + Status::ServiceUnavailable, + Json(serde_json::json!({ "error": "model not init" })), + ); + } + }; + let mut guard = model_ref.write().await; + let embeddings = match guard.instance.embedding(&texts) { + Ok(v) => v, + Err(e) => { + return ( + Status::BadRequest, + Json(serde_json::json!({ "error": e.to_string() })), + ); + } + }; + let model_name = guard.which_model.as_string(); + let data = embeddings + .into_iter() + .enumerate() + .map(|(index, embedding)| EmbeddingData { + object: "embedding".to_string(), + index, + embedding, + }) + .collect::>(); + let response = EmbeddingResponse { + object: "list".to_string(), + data, + model: model_name, + }; + (Status::Ok, Json(serde_json::to_value(response).unwrap())) +} diff --git a/src/server/mod.rs b/src/server/mod.rs index d9a81a0..3e26d6b 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -10,7 +10,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; // ASR (Automatic Speech Recognition) API module pub(crate) mod api; pub(crate) mod asr; -pub(crate) mod asr_types; +pub(crate) mod embedding; pub(crate) mod process; pub(crate) async fn start_http_server( @@ -61,6 +61,9 @@ pub(crate) async fn start_http_server( builder = builder.mount("/audio", routes![api::speech, asr::transcriptions]); // /v1/audio/transcriptions (OpenAI standard ASR transcription endpoint) builder = builder.mount("/v1/audio", routes![asr::transcriptions]); + // /embeddings and /v1/embeddings (OpenAI-compatible embeddings endpoint) + builder = builder.mount("/", routes![embedding::embeddings]); + builder = builder.mount("/v1", routes![embedding::embeddings]); // Health check and model info endpoints builder = builder.mount("/", routes![api::health, api::models]); // Shutdown endpoint diff --git a/src/utils/audio_utils.rs b/src/utils/audio_utils.rs index 4ad37b0..eb9ac7a 100644 --- a/src/utils/audio_utils.rs +++ b/src/utils/audio_utils.rs @@ -4,6 +4,7 @@ use std::path::{Path, PathBuf}; use std::thread; use std::{f64::consts::PI, io::Cursor}; +use crate::models::common::modules::log10; use crate::params::chat::{ ChatCompletionParameters, ChatCompletionResponse, ChatMessage, ChatMessageContent, ChatMessageContentPart, @@ -33,7 +34,7 @@ use symphonia::core::probe::Hint; use crate::utils::get_default_save_dir; use crate::utils::tensor_utils::{ - linspace, log10, pad_reflect_last_dim, pad_replicate_last_dim, split_tensor, + linspace, pad_reflect_last_dim, pad_replicate_last_dim, split_tensor, }; // 重采样方法枚举 diff --git a/src/utils/tensor_utils.rs b/src/utils/tensor_utils.rs index e491af3..fd86e7d 100644 --- a/src/utils/tensor_utils.rs +++ b/src/utils/tensor_utils.rs @@ -2,7 +2,6 @@ use std::f32; use anyhow::{Result, anyhow}; use candle_core::{D, DType, Device, IndexOp, Tensor, shape::Dim}; -use candle_nn::ops::sigmoid; pub enum PaddingSide { Left, @@ -463,12 +462,6 @@ pub fn index_select_2d(t: &Tensor, index: &Tensor) -> Result { Ok(res) } -pub fn quick_gelu(xs: &Tensor) -> Result { - let x = xs.affine(1.702, 0.0)?; - let x = sigmoid(&x)?; - Ok(xs.mul(&x)?) -} - pub fn topk(weight: &Tensor, topk: usize) -> Result<(Tensor, Tensor)> { let topk_idx = weight .arg_sort_last_dim(false)? @@ -558,102 +551,6 @@ pub fn pad_replicate_last_dim(t: &Tensor, pad: (usize, usize)) -> Result Ok(pad_tensor) } -pub fn log10(t: &Tensor) -> Result { - Ok(t.log()?.affine(1.0 / 10.0_f64.ln(), 0.0)?) -} - -pub fn z_score_normalize(t: &Tensor, dim: usize) -> Result { - let rank = t.rank(); - if dim >= rank { - return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank))); - } - Ok(t.broadcast_sub(&t.mean_keepdim(dim)?)? - .broadcast_div(&t.var_keepdim(dim)?.sqrt()?)?) -} - -pub fn l2_normalize(t: &Tensor, dim: usize) -> Result { - let rank = t.rank(); - if dim >= rank { - return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank))); - } - let l2_norm = t.sqr()?.sum_keepdim(dim)?.affine(1.0, 1e-6)?.sqrt()?; - Ok(t.broadcast_div(&l2_norm)?) -} - -pub fn l1_normalize(t: &Tensor, dim: usize) -> Result { - let rank = t.rank(); - if dim >= rank { - return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank))); - } - let l1_norm = t.abs()?.sum_keepdim(dim)?; - Ok(t.broadcast_div(&l1_norm)?) -} - -pub fn pool1d(xs: &Tensor, pool_size: usize, ceil_mode: bool, stype: &str) -> Result { - // xs: (bs, c, dim) - // ceil_mode: 是否保留不完整窗口，为true时通过pad实现 - if pool_size == 0 { - return Err(anyhow!("pool_size must be greater than 0")); - } - let (bs, c, dim) = xs.dims3()?; - let xs_reshape = if ceil_mode { - let remain = dim % pool_size; - if remain > 0 { - let pad = pool_size - remain; - let xs_pad = pad_replicate_last_dim(xs, (0, pad))?; - xs_pad.reshape((bs, c, (), pool_size))? - } else { - xs.reshape((bs, c, (), pool_size))? - } - } else { - let remain = dim % pool_size; - if remain > 0 { - let xs_del = xs.narrow(D::Minus1, 0, dim - remain)?; - xs_del.reshape((bs, c, (), pool_size))? - } else { - xs.reshape((bs, c, (), pool_size))? - } - }; - let xs_pool = match stype { - "avg" => xs_reshape.mean(D::Minus1)?, - "max" => xs_reshape.max(D::Minus1)?, - "min" => xs_reshape.min(D::Minus1)?, - _ => { - return Err(anyhow!( - "unsupported pool type: {}, supported types are: avg, max, min", - stype - )); - } - }; - Ok(xs_pool) -} - -pub fn statistics_pooling(xs: &Tensor, dim: D, keepdim: bool) -> Result { - let mean = xs.mean(dim)?; - let std = xs.var(dim)?.sqrt()?; - let mut stats = Tensor::cat(&[mean, std], D::Minus1)?; - if keepdim { - stats = stats.unsqueeze(dim)?; - } - Ok(stats) -} -pub fn float_range_normalize(t: &Tensor) -> Result { - let peak = t - .to_dtype(DType::F32)? - .abs()? - .max_all()? - .to_scalar::()?; - if peak == 0.0 { - return Ok(t.clone()); - } - let mut t = t.clone(); - if peak > 1.0 { - t = t.affine(1.0 / peak as f64, 0.0)?; - } - t = t.clamp(-1.0, 1.0)?; - Ok(t) -} - pub fn sequence_mask(length: &Tensor, max_length: Option) -> Result { let max_length = max_length.unwrap_or(length.max_all()?.to_scalar::()?); let x = Tensor::arange(0, max_length, length.device())?.unsqueeze(0)?; @@ -662,17 +559,6 @@ pub fn sequence_mask(length: &Tensor, max_length: Option) -> Result Ok(mask) } -pub fn cosine_similarity(query_vector: &Tensor, matrix: &Tensor) -> Result { - // query_vector: (n, dim) - // matrix: (m, dim) - let query_norm = l2_normalize(query_vector, query_vector.rank() - 1)?; - let matrix_norm = l2_normalize(matrix, matrix.rank() - 1)?; - let similarity = query_norm - .matmul(&matrix_norm.transpose(D::Minus1, D::Minus2)?)? - .squeeze(D::Minus1)?; - Ok(similarity) -} - pub fn repeat_interleave(t: &Tensor, repeats: usize, dim: usize) -> Result { if repeats == 1 { return Ok(t.clone()); diff --git a/tests/test_all_minilm_l6_v2.rs b/tests/test_all_minilm_l6_v2.rs new file mode 100644 index 0000000..5e37d18 --- /dev/null +++ b/tests/test_all_minilm_l6_v2.rs @@ -0,0 +1,24 @@ +use aha::models::{ + all_minilm_l6_v2::AllMiniLML6V2Embedding, common::embedding::TextEmbedding, +}; +use anyhow::Result; +use std::time::Instant; + +#[test] +fn all_minilm_l6_v2_embedding() -> Result<()> { + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda --test test_all_minilm_l6_v2 all_minilm_l6_v2_embedding -r -- --nocapture + + let save_dir = + aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/sentence-transformers/all-MiniLM-L6-v2/", save_dir); + + let i_start = Instant::now(); + let mut model = AllMiniLML6V2Embedding::init(&model_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + let input_texts = ["test ALL_MINILM_L6_V2 embedding".to_string()]; + let result = model.embed_texts(&input_texts)?; + println!("result: {:?}", result); + + Ok(()) +} diff --git a/tests/test_qwen3_embedding.rs b/tests/test_qwen3_embedding.rs new file mode 100644 index 0000000..8ef2b6e --- /dev/null +++ b/tests/test_qwen3_embedding.rs @@ -0,0 +1,22 @@ +use aha::models::{common::embedding::TextEmbedding, qwen3_embedding::Qwen3Embedding}; +use anyhow::Result; +use std::time::Instant; + +#[test] +fn qwen3_embedding() -> Result<()> { + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda --test test_qwen3_embedding qwen3_embedding -r -- --nocapture + + let save_dir = + aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/Qwen/Qwen3-Embedding-0.6B/", save_dir); + + let i_start = Instant::now(); + let mut model = Qwen3Embedding::init(&model_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + let input_texts = ["test Qwen3-Embedding-0.6B".to_string()]; + let result = model.embed_texts(&input_texts)?; + println!("result: {:?}", result); + + Ok(()) +} diff --git a/tests/weight_test.rs b/tests/weight_test.rs index f80389d..ebe4d4c 100644 --- a/tests/weight_test.rs +++ b/tests/weight_test.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use aha::utils::{find_type_files, get_device}; use anyhow::Result; -use candle_core::{Device, pickle::read_all_with_key, safetensors}; +use candle_core::{Device, pickle::read_all_with_key, quantized::gguf_file, safetensors}; use candle_nn::VarBuilder; #[test] @@ -310,3 +310,21 @@ fn lfm2vl_weight() -> Result<()> { println!("model_list: {:?}", model_list); Ok(()) } + +#[test] +fn gguf_weight() -> Result<()> { + // cargo test -F cuda --test weight_test gguf_weight -r -- --nocapture + let gguf_path = "/home/jhq/.aha/Qwen/Qwen3-Embedding-0.6B-GGUF/Qwen3-Embedding-0.6B-f16.gguf"; + let mut model_file = std::fs::File::open(gguf_path)?; + let model = gguf_file::Content::read(&mut model_file)?; + for (key, value) in model.tensor_infos { + println!("{key}: {:#?}", value); + } + // for (key, value) in model.metadata { + // if key.contains("tokeni") { + // continue; + // } + // println!("{key}: {:#?}", value); + // } + Ok(()) +}