add qwen3-embedding and all-minilm-l6-v2
This commit is contained in:
+23
-54
@@ -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 => {
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,8 @@ use crate::cli::{
|
||||
};
|
||||
|
||||
mod cli;
|
||||
#[allow(unused)]
|
||||
mod params;
|
||||
mod server;
|
||||
|
||||
#[tokio::main]
|
||||
|
||||
@@ -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,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 {
|
||||
|
||||
@@ -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,5 +1,6 @@
|
||||
use anyhow::Result;
|
||||
use candle_core::Tensor;
|
||||
pub mod embedding;
|
||||
pub mod generate;
|
||||
pub mod gguf;
|
||||
pub mod model_mapping;
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)?)
|
||||
}
|
||||
|
||||
@@ -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::{
|
||||
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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
+87
-89
@@ -1,3 +1,4 @@
|
||||
pub mod all_minilm_l6_v2;
|
||||
pub mod bigvgan;
|
||||
pub mod campplus;
|
||||
pub mod common;
|
||||
@@ -17,16 +18,22 @@ pub mod qwen2_5vl;
|
||||
pub mod qwen3;
|
||||
pub mod qwen3_5;
|
||||
pub mod qwen3_asr;
|
||||
pub mod qwen3_embedding;
|
||||
pub mod qwen3vl;
|
||||
pub mod rmbg2_0;
|
||||
pub mod voxcpm;
|
||||
pub mod w2v_bert_2_0;
|
||||
|
||||
use crate::{
|
||||
models::common::model_mapping::WhichModel,
|
||||
models::{
|
||||
all_minilm_l6_v2::AllMiniLML6V2Embedding,
|
||||
common::{embedding::TextEmbedding, model_mapping::WhichModel},
|
||||
qwen3_embedding::Qwen3Embedding,
|
||||
},
|
||||
params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse},
|
||||
};
|
||||
use anyhow::{Result, anyhow};
|
||||
use candle_core::{DType, Device};
|
||||
use rocket::futures::Stream;
|
||||
|
||||
use crate::models::{
|
||||
@@ -57,6 +64,7 @@ pub trait GenerateModel {
|
||||
}
|
||||
|
||||
pub enum ModelInstance<'a> {
|
||||
AllMiniLML6V2(AllMiniLML6V2Embedding),
|
||||
MiniCPM4(MiniCPMGenerateModel<'a>),
|
||||
Lfm2(Lfm2GenerateModel<'a>),
|
||||
Lfm2VL(Lfm2VLGenerateModel<'a>),
|
||||
@@ -64,6 +72,7 @@ pub enum ModelInstance<'a> {
|
||||
Qwen3(Qwen3GenerateModel<'a>),
|
||||
Qwen3_5(Qwen3_5GenerateModel<'a>),
|
||||
Qwen3ASR(Qwen3AsrGenerateModel<'a>),
|
||||
Qwen3Embedding(Qwen3Embedding),
|
||||
Qwen3VL(Box<Qwen3VLGenerateModel<'a>>),
|
||||
DeepSeekOCR(DeepseekOCRGenerateModel),
|
||||
HunyuanOCR(HunyuanOCRGenerateModel<'a>),
|
||||
@@ -78,12 +87,18 @@ pub enum ModelInstance<'a> {
|
||||
impl<'a> GenerateModel for ModelInstance<'a> {
|
||||
fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
|
||||
match self {
|
||||
ModelInstance::AllMiniLML6V2(_) => {
|
||||
Err(anyhow!("embedding model does not support chat completions"))
|
||||
}
|
||||
ModelInstance::MiniCPM4(model) => model.generate(mes),
|
||||
ModelInstance::Lfm2(model) => model.generate(mes),
|
||||
ModelInstance::Lfm2VL(model) => model.generate(mes),
|
||||
ModelInstance::Qwen2_5VL(model) => model.generate(mes),
|
||||
ModelInstance::Qwen3(model) => model.generate(mes),
|
||||
ModelInstance::Qwen3_5(model) => model.generate(mes),
|
||||
ModelInstance::Qwen3Embedding(_) => {
|
||||
Err(anyhow!("embedding model does not support chat completions"))
|
||||
}
|
||||
ModelInstance::Qwen3ASR(model) => model.generate(mes),
|
||||
ModelInstance::Qwen3VL(model) => model.generate(mes),
|
||||
ModelInstance::DeepSeekOCR(model) => model.generate(mes),
|
||||
@@ -109,12 +124,18 @@ impl<'a> GenerateModel for ModelInstance<'a> {
|
||||
>,
|
||||
> {
|
||||
match self {
|
||||
ModelInstance::AllMiniLML6V2(_) => {
|
||||
Err(anyhow!("embedding model does not support chat completions"))
|
||||
}
|
||||
ModelInstance::MiniCPM4(model) => model.generate_stream(mes),
|
||||
ModelInstance::Lfm2(model) => model.generate_stream(mes),
|
||||
ModelInstance::Lfm2VL(model) => model.generate_stream(mes),
|
||||
ModelInstance::Qwen2_5VL(model) => model.generate_stream(mes),
|
||||
ModelInstance::Qwen3(model) => model.generate_stream(mes),
|
||||
ModelInstance::Qwen3_5(model) => model.generate_stream(mes),
|
||||
ModelInstance::Qwen3Embedding(_) => Err(anyhow!(
|
||||
"embedding model does not support streaming chat completions"
|
||||
)),
|
||||
ModelInstance::Qwen3VL(model) => model.generate_stream(mes),
|
||||
ModelInstance::Qwen3ASR(model) => model.generate_stream(mes),
|
||||
ModelInstance::DeepSeekOCR(model) => model.generate_stream(mes),
|
||||
@@ -129,16 +150,27 @@ impl<'a> GenerateModel for ModelInstance<'a> {
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> ModelInstance<'a> {
|
||||
pub fn embedding(&mut self, input: &[String]) -> Result<Vec<Vec<f32>>> {
|
||||
match self {
|
||||
ModelInstance::Qwen3Embedding(model) => model.embed_texts(input),
|
||||
ModelInstance::AllMiniLML6V2(model) => model.embed_texts(input),
|
||||
_ => Err(anyhow!("current model does not support embeddings")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(unused)]
|
||||
pub fn load_gguf_model<'a>(
|
||||
model_type: WhichModel,
|
||||
config_path: Option<&str>, // 有些gguf未包含模型其他配置,需额外指定
|
||||
gguf_path: &str,
|
||||
mmproj_path: Option<&str>,
|
||||
device: Option<&Device>,
|
||||
) -> Result<ModelInstance<'a>> {
|
||||
let model = match model_type {
|
||||
WhichModel::Qwen3_5Gguf => {
|
||||
let model = Qwen3_5GenerateModel::init_from_gguf(gguf_path, mmproj_path, None)?;
|
||||
let model = Qwen3_5GenerateModel::init_from_gguf(gguf_path, mmproj_path, device)?;
|
||||
ModelInstance::Qwen3_5(model)
|
||||
}
|
||||
_ => {
|
||||
@@ -149,135 +181,101 @@ pub fn load_gguf_model<'a>(
|
||||
Ok(model)
|
||||
}
|
||||
|
||||
pub fn load_model<'a>(model_type: WhichModel, path: &str) -> Result<ModelInstance<'a>> {
|
||||
pub fn load_model<'a>(
|
||||
model_type: WhichModel,
|
||||
path: &str,
|
||||
device: Option<&Device>,
|
||||
dtype: Option<DType>,
|
||||
) -> Result<ModelInstance<'a>> {
|
||||
let model = match model_type {
|
||||
WhichModel::AllMiniLML6V2 => {
|
||||
let model = AllMiniLML6V2Embedding::init(path, device, dtype)?;
|
||||
ModelInstance::AllMiniLML6V2(model)
|
||||
}
|
||||
WhichModel::MiniCPM4_0_5B => {
|
||||
let model = MiniCPMGenerateModel::init(path, None, None)?;
|
||||
let model = MiniCPMGenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::MiniCPM4(model)
|
||||
}
|
||||
WhichModel::LFM2_1_2B => {
|
||||
let model = Lfm2GenerateModel::init(path, None, None)?;
|
||||
WhichModel::LFM2_1_2B | WhichModel::LFM2_5_1_2BInstruct => {
|
||||
let model = Lfm2GenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::Lfm2(model)
|
||||
}
|
||||
WhichModel::LFM2_5_1_2BInstruct => {
|
||||
let model = Lfm2GenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Lfm2(model)
|
||||
}
|
||||
WhichModel::LFM2_5VL1_6B => {
|
||||
let model = Lfm2VLGenerateModel::init(path, None, None)?;
|
||||
WhichModel::LFM2_5VL1_6B | WhichModel::LFM2VL1_6B => {
|
||||
let model = Lfm2VLGenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::Lfm2VL(model)
|
||||
}
|
||||
WhichModel::LFM2VL1_6B => {
|
||||
let model = Lfm2VLGenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Lfm2VL(model)
|
||||
}
|
||||
WhichModel::Qwen2_5VL3B => {
|
||||
let model = Qwen2_5VLGenerateModel::init(path, None, None)?;
|
||||
WhichModel::Qwen2_5VL3B | WhichModel::Qwen2_5VL7B => {
|
||||
let model = Qwen2_5VLGenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::Qwen2_5VL(model)
|
||||
}
|
||||
WhichModel::Qwen2_5VL7B => {
|
||||
let model = Qwen2_5VLGenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Qwen2_5VL(model)
|
||||
}
|
||||
WhichModel::Qwen3_0_6B => {
|
||||
let model = Qwen3GenerateModel::init(path, None, None)?;
|
||||
WhichModel::Qwen3_0_6B | WhichModel::Qwen3_1_7B | WhichModel::Qwen3_4B => {
|
||||
let model = Qwen3GenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::Qwen3(model)
|
||||
}
|
||||
WhichModel::Qwen3_1_7B => {
|
||||
let model = Qwen3GenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Qwen3(model)
|
||||
}
|
||||
WhichModel::Qwen3_4B => {
|
||||
let model = Qwen3GenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Qwen3(model)
|
||||
}
|
||||
WhichModel::Qwen3_5_0_8B => {
|
||||
let model = Qwen3_5GenerateModel::init(path, None, None)?;
|
||||
WhichModel::Qwen3_5_0_8B
|
||||
| WhichModel::Qwen3_5_2B
|
||||
| WhichModel::Qwen3_5_4B
|
||||
| WhichModel::Qwen3_5_9B => {
|
||||
let model = Qwen3_5GenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::Qwen3_5(model)
|
||||
}
|
||||
WhichModel::Qwen3_5_2B => {
|
||||
let model = Qwen3_5GenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Qwen3_5(model)
|
||||
}
|
||||
WhichModel::Qwen3_5_4B => {
|
||||
let model = Qwen3_5GenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Qwen3_5(model)
|
||||
}
|
||||
WhichModel::Qwen3_5_9B => {
|
||||
let model = Qwen3_5GenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Qwen3_5(model)
|
||||
}
|
||||
WhichModel::Qwen3ASR0_6B => {
|
||||
let model = Qwen3AsrGenerateModel::init(path, None, None)?;
|
||||
WhichModel::Qwen3ASR0_6B | WhichModel::Qwen3ASR1_7B => {
|
||||
let model = Qwen3AsrGenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::Qwen3ASR(model)
|
||||
}
|
||||
WhichModel::Qwen3ASR1_7B => {
|
||||
let model = Qwen3AsrGenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Qwen3ASR(model)
|
||||
WhichModel::Qwen3Embedding0_6B
|
||||
| WhichModel::Qwen3Embedding4B
|
||||
| WhichModel::Qwen3Embedding8B => {
|
||||
let model = Qwen3Embedding::init(path, device, dtype)?;
|
||||
ModelInstance::Qwen3Embedding(model)
|
||||
}
|
||||
WhichModel::Qwen3VL2B => {
|
||||
let model = Qwen3VLGenerateModel::init(path, None, None)?;
|
||||
WhichModel::Qwen3VL2B
|
||||
| WhichModel::Qwen3VL4B
|
||||
| WhichModel::Qwen3VL8B
|
||||
| WhichModel::Qwen3VL32B => {
|
||||
let model = Qwen3VLGenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::Qwen3VL(Box::new(model))
|
||||
}
|
||||
WhichModel::Qwen3VL4B => {
|
||||
let model = Qwen3VLGenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Qwen3VL(Box::new(model))
|
||||
}
|
||||
WhichModel::Qwen3VL8B => {
|
||||
let model = Qwen3VLGenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Qwen3VL(Box::new(model))
|
||||
}
|
||||
WhichModel::Qwen3VL32B => {
|
||||
let model = Qwen3VLGenerateModel::init(path, None, None)?;
|
||||
ModelInstance::Qwen3VL(Box::new(model))
|
||||
}
|
||||
WhichModel::DeepSeekOCR => {
|
||||
let model = DeepseekOCRGenerateModel::init(path, None, None)?;
|
||||
ModelInstance::DeepSeekOCR(model)
|
||||
}
|
||||
WhichModel::DeepSeekOCR2 => {
|
||||
let model = DeepseekOCRGenerateModel::init(path, None, None)?;
|
||||
WhichModel::DeepSeekOCR | WhichModel::DeepSeekOCR2 => {
|
||||
let model = DeepseekOCRGenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::DeepSeekOCR(model)
|
||||
}
|
||||
WhichModel::HunyuanOCR => {
|
||||
let model = HunyuanOCRGenerateModel::init(path, None, None)?;
|
||||
let model = HunyuanOCRGenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::HunyuanOCR(model)
|
||||
}
|
||||
WhichModel::PaddleOCRVL => {
|
||||
let model = PaddleOCRVLGenerateModel::init(path, None, None)?;
|
||||
ModelInstance::PaddleOCRVL(Box::new(model))
|
||||
}
|
||||
WhichModel::PaddleOCRVL1_5 => {
|
||||
let model = PaddleOCRVLGenerateModel::init(path, None, None)?;
|
||||
WhichModel::PaddleOCRVL | WhichModel::PaddleOCRVL1_5 => {
|
||||
let model = PaddleOCRVLGenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::PaddleOCRVL(Box::new(model))
|
||||
}
|
||||
WhichModel::RMBG2_0 => {
|
||||
let model = RMBG2_0Model::init(path, None, None)?;
|
||||
let model = RMBG2_0Model::init(path, device, dtype)?;
|
||||
ModelInstance::RMBG2_0(Box::new(model))
|
||||
}
|
||||
WhichModel::VoxCPM => {
|
||||
let model = VoxCPMGenerate::init(path, None, None)?;
|
||||
ModelInstance::VoxCPM(Box::new(model))
|
||||
}
|
||||
WhichModel::VoxCPM1_5 => {
|
||||
let model = VoxCPMGenerate::init(path, None, None)?;
|
||||
WhichModel::VoxCPM | WhichModel::VoxCPM1_5 => {
|
||||
let model = VoxCPMGenerate::init(path, device, dtype)?;
|
||||
ModelInstance::VoxCPM(Box::new(model))
|
||||
}
|
||||
WhichModel::GlmASRNano2512 => {
|
||||
let model = GlmAsrNanoGenerateModel::init(path, None, None)?;
|
||||
let model = GlmAsrNanoGenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::GlmASRNano(model)
|
||||
}
|
||||
WhichModel::FunASRNano2512 => {
|
||||
let model = FunAsrNanoGenerateModel::init(path, None, None)?;
|
||||
let model = FunAsrNanoGenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::FunASRNano(model)
|
||||
}
|
||||
WhichModel::GlmOCR => {
|
||||
let model = GlmOcrGenerateModel::init(path, None, None)?;
|
||||
let model = GlmOcrGenerateModel::init(path, device, dtype)?;
|
||||
ModelInstance::GlmOCR(model)
|
||||
}
|
||||
_ => {
|
||||
let model_id = model_type.as_string();
|
||||
return Err(anyhow!("model id {model_id} is not safetensor model"));
|
||||
if model_id.to_lowercase().contains("gguf") || model_id.to_lowercase().contains("onnx")
|
||||
{
|
||||
return Err(anyhow!("model id {model_id} is not safetensor model"));
|
||||
} else {
|
||||
return Err(anyhow!("model id {model_id} not impl load_model function"));
|
||||
}
|
||||
}
|
||||
};
|
||||
Ok(model)
|
||||
|
||||
@@ -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)?)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
};
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
};
|
||||
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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>,
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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
|
||||
///
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
// 重采样方法枚举
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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
@@ -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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user