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

This commit is contained in:
jhqxxx
2026-04-06 15:35:49 +08:00
parent 44d91da650
commit 7142e712f2
34 changed files with 763 additions and 319 deletions
+23 -54
View File
@@ -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 => {
+41
View File
@@ -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(())
}
}
+2
View File
@@ -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;
+42
View File
@@ -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(())
}
}
+2
View File
@@ -6,6 +6,8 @@ use crate::cli::{
};
mod cli;
#[allow(unused)]
mod params;
mod server;
#[tokio::main]
+78
View File
@@ -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<DType>) -> Result<Self> {
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<Vec<u32>> {
let mut token_ids = self.tokenizer.text_encode_vec(text.to_string(), true)?;
token_ids = token_ids
.into_iter()
.filter(|&x| x != 0)
.collect::<Vec<u32>>();
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<Vec<f32>> {
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::<f32>()?;
Ok(embed)
}
}
impl TextEmbedding for AllMiniLML6V2Embedding {
fn embed_texts(&mut self, input: &[String]) -> Result<Vec<Vec<f32>>> {
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)
}
}
+2 -3
View File
@@ -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 {
+29
View File
@@ -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<Vec<Vec<f32>>>;
}
pub enum NormalizeType {
L1,
L2,
ZScore,
MinMax,
MaxAbs,
}
impl NormalizeType {
pub fn normalize(&self, t: &Tensor, dim: usize) -> Result<Tensor> {
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),
}
}
}
+1
View File
@@ -1,5 +1,6 @@
use anyhow::Result;
use candle_core::Tensor;
pub mod embedding;
pub mod generate;
pub mod gguf;
pub mod model_mapping;
+12
View File
@@ -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",
}
}
}
+134 -3
View File
@@ -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<Tensor> {
Ok(t.log()?.affine(1.0 / 10.0_f64.ln(), 0.0)?)
}
pub fn max_abs_normalize(t: &Tensor, dim: usize) -> Result<Tensor> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
// 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<Tensor> {
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<Tensor> {
let peak = t
.to_dtype(DType::F32)?
.abs()?
.max_all()?
.to_scalar::<f32>()?;
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<Tensor> {
// 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<Tensor> {
let x = xs.affine(1.702, 0.0)?;
let x = sigmoid(&x)?;
Ok(xs.mul(&x)?)
}
+2 -2
View File
@@ -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,
},
},
};
@@ -1,9 +1,12 @@
use anyhow::Result;
use candle_core::{D, Device, Tensor};
use crate::utils::{
use crate::{
models::common::modules::log10,
utils::{
audio_utils::{create_hann_window, mel_filter_bank, torch_stft},
tensor_utils::{log10, pad_reflect_last_dim},
tensor_utils::pad_reflect_last_dim,
},
};
pub struct WhisperFeatureExtractor {
@@ -1,9 +1,12 @@
use anyhow::Result;
use candle_core::{D, Device, Tensor};
use crate::utils::{
use crate::{
models::common::modules::z_score_normalize,
utils::{
audio_utils::{create_povey_window, mel_filter_bank, spectrogram},
tensor_utils::{PaddingSide, z_score_normalize},
tensor_utils::PaddingSide,
},
};
pub struct SeamlessM4TFeatureExtractor {
+2 -2
View File
@@ -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 {
+86 -88
View File
@@ -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();
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)
+18 -3
View File
@@ -102,7 +102,11 @@ pub struct Qwen3Model {
impl Qwen3Model {
pub fn new(config: &Qwen3Config, vb: VarBuilder, eos_ids: Vec<u32>) -> Result<Self> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
Ok(self.embed_tokens.forward(input_ids)?)
}
+4 -2
View File
@@ -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,
},
};
+3 -2
View File
@@ -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,
},
};
+66
View File
@@ -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<DType>) -> Result<Self> {
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<Vec<f32>> {
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::<f32>()?;
Ok(norm)
}
}
impl TextEmbedding for Qwen3Embedding {
fn embed_texts(&mut self, input: &[String]) -> Result<Vec<Vec<f32>>> {
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)
}
}
@@ -46,33 +46,3 @@ pub(crate) struct ErrorDetail {
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) code: Option<String>,
}
#[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");
}
}
+22
View File
@@ -0,0 +1,22 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Deserialize)]
pub(crate) struct EmbeddingRequest {
pub model: Option<String>,
pub input: Value,
}
#[derive(Debug, Serialize)]
pub(crate) struct EmbeddingData {
pub object: String,
pub index: usize,
pub embedding: Vec<f32>,
}
#[derive(Debug, Serialize)]
pub(crate) struct EmbeddingResponse {
pub object: String,
pub model: String,
pub data: Vec<EmbeddingData>,
}
+6
View File
@@ -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;
+23
View File
@@ -0,0 +1,23 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Deserialize)]
pub(crate) struct RerankRequest {
pub model: Option<String>,
pub query: String,
pub documents: Vec<String>,
pub top_n: Option<usize>,
}
#[derive(Debug, Serialize)]
struct RerankResult {
index: usize,
relevance_score: f32,
document: String,
}
#[derive(Debug, Serialize)]
struct RerankResponse {
object: String,
model: String,
results: Vec<RerankResult>,
}
+1 -1
View File
@@ -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")]
+2 -2
View File
@@ -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(|| {
+1 -3
View File
@@ -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
///
+77
View File
@@ -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<Vec<String>> {
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 = "<req>")]
pub(crate) async fn embeddings(req: Json<EmbeddingRequest>) -> (Status, Json<Value>) {
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::<Vec<_>>();
let response = EmbeddingResponse {
object: "list".to_string(),
data,
model: model_name,
};
(Status::Ok, Json(serde_json::to_value(response).unwrap()))
}
+4 -1
View File
@@ -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
+2 -1
View File
@@ -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,
};
// 重采样方法枚举
-114
View File
@@ -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<Tensor> {
Ok(res)
}
pub fn quick_gelu(xs: &Tensor) -> Result<Tensor> {
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<Tensor>
Ok(pad_tensor)
}
pub fn log10(t: &Tensor) -> Result<Tensor> {
Ok(t.log()?.affine(1.0 / 10.0_f64.ln(), 0.0)?)
}
pub fn z_score_normalize(t: &Tensor, dim: usize) -> Result<Tensor> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
// 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<Tensor> {
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<Tensor> {
let peak = t
.to_dtype(DType::F32)?
.abs()?
.max_all()?
.to_scalar::<f32>()?;
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<u32>) -> Result<Tensor> {
let max_length = max_length.unwrap_or(length.max_all()?.to_scalar::<u32>()?);
let x = Tensor::arange(0, max_length, length.device())?.unsqueeze(0)?;
@@ -662,17 +559,6 @@ pub fn sequence_mask(length: &Tensor, max_length: Option<u32>) -> Result<Tensor>
Ok(mask)
}
pub fn cosine_similarity(query_vector: &Tensor, matrix: &Tensor) -> Result<Tensor> {
// 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<Tensor> {
if repeats == 1 {
return Ok(t.clone());
+24
View File
@@ -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(())
}
+22
View File
@@ -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(())
}
+19 -1
View File
@@ -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(())
}