From 073723447fd03170d3face2f948d6830aa99a938 Mon Sep 17 00:00:00 2001 From: jhqxxx <18280426169@163.com> Date: Mon, 2 Feb 2026 21:30:52 +0800 Subject: [PATCH] merge main --- src/chat_template/mod.rs | 37 ++++-- src/models/feature_extractor/config.rs | 21 ++++ .../feature_extraction_whisper.rs | 113 ++++++++++++++++++ src/models/feature_extractor/mod.rs | 4 +- .../seamless_m4t_feature_extractor.rs | 6 +- src/models/glm_asr_nano/config.rs | 17 +-- src/models/glm_asr_nano/processor.rs | 70 +++++++---- src/models/index_tts2/processor.rs | 2 +- src/models/mod.rs | 1 + src/models/qwen3_asr/model.rs | 6 +- src/utils/audio_utils.rs | 21 +++- src/utils/mod.rs | 30 +++++ src/utils/tensor_utils.rs | 16 +++ 13 files changed, 285 insertions(+), 59 deletions(-) create mode 100644 src/models/feature_extractor/config.rs create mode 100644 src/models/feature_extractor/feature_extraction_whisper.rs diff --git a/src/chat_template/mod.rs b/src/chat_template/mod.rs index 0cbf6d1..3a4c3d2 100644 --- a/src/chat_template/mod.rs +++ b/src/chat_template/mod.rs @@ -6,16 +6,33 @@ use crate::utils::string_to_static_str; pub fn get_template(path: String) -> Result { let tokenizer_config_file = path.clone() + "/tokenizer_config.json"; - assert!( - std::path::Path::new(&tokenizer_config_file).exists(), - "tokenizer_config.json not exists in model path" - ); - let tokenizer_config: serde_json::Value = - serde_json::from_slice(&std::fs::read(tokenizer_config_file)?) - .map_err(|e| anyhow!(format!("load tokenizer_config file error:{}", e)))?; - let chat_template = tokenizer_config["chat_template"] - .as_str() - .ok_or(anyhow!(format!("chat_template to str error")))?; + let chat_template = if std::path::Path::new(&tokenizer_config_file).exists() { + let tokenizer_config: serde_json::Value = + serde_json::from_slice(&std::fs::read(tokenizer_config_file)?) + .map_err(|e| anyhow!(format!("load tokenizer_config file error:{}", e)))?; + let chat_template = tokenizer_config["chat_template"] + .as_str() + .map(|s| s.to_string()); + match chat_template { + Some(tem) => Some(tem), + None => { + let chat_template_file = path.clone() + "/chat_template.json"; + if std::path::Path::new(&chat_template_file).exists() { + let chat_template_config: serde_json::Value = + serde_json::from_slice(&std::fs::read(chat_template_file)?) + .map_err(|e| anyhow!(format!("load chat_template file error:{}", e)))?; + chat_template_config["chat_template"] + .as_str() + .map(|s| s.to_string()) + } else { + None + } + } + } + } else { + None + }; + let chat_template = chat_template.ok_or(anyhow!(format!("chat_template is none")))?; // 修复模板中的问题行 let fixed_template = chat_template .replace( diff --git a/src/models/feature_extractor/config.rs b/src/models/feature_extractor/config.rs new file mode 100644 index 0000000..b79a94d --- /dev/null +++ b/src/models/feature_extractor/config.rs @@ -0,0 +1,21 @@ +use serde::Deserialize; + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct FeatureExtractor { + pub chunk_length: usize, + pub dither: f64, + pub feature_size: usize, + pub hop_length: usize, + pub n_fft: usize, + pub n_samples: usize, + pub nb_max_frames: usize, + pub padding_side: String, + pub padding_value: f32, + pub return_attention_mask: bool, + #[serde(default = "default_sampling_rate")] + pub sampling_rate: usize, +} + +fn default_sampling_rate() -> usize { + 16000 +} \ No newline at end of file diff --git a/src/models/feature_extractor/feature_extraction_whisper.rs b/src/models/feature_extractor/feature_extraction_whisper.rs new file mode 100644 index 0000000..f105fe3 --- /dev/null +++ b/src/models/feature_extractor/feature_extraction_whisper.rs @@ -0,0 +1,113 @@ +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}, +}; + +pub struct WhisperFeatureExtractor { + feature_size: usize, + hop_length: usize, + chunk_length: usize, + n_samples: usize, + n_fft: usize, + dither: f64, + padding_value: f32, + sampling_rate: usize, + mel_filters: Tensor, + window: Tensor, +} + +impl WhisperFeatureExtractor { + pub fn new( + feature_size: usize, + hop_length: usize, + chunk_length: usize, + n_fft: usize, + dither: f64, + padding_value: f32, + sampling_rate: usize, + device: &Device, + ) -> Result { + let window = create_hann_window(n_fft, candle_core::DType::F32, device)?; + let window = window.unsqueeze(0)?.unsqueeze(0)?; + let mel_filters = mel_filter_bank( + 1 + n_fft / 2, + feature_size, + 0.0, + 8000.0, + sampling_rate as f32, + Some("slaney"), + crate::utils::audio_utils::MelScale::Slaney, + false, + device, + )? + .t()?; + let n_samples = chunk_length * sampling_rate; + Ok(Self { + feature_size, + hop_length, + chunk_length, + n_samples, + n_fft, + dither, + padding_value, + sampling_rate, + mel_filters, + window, + }) + } + + pub fn call( + &self, + raw_speech: &Tensor, + sampling_rate: usize, + // do_normalize: bool, + return_attention_mask: bool, + ) -> Result<(Tensor, Option)> { + // raw_speech: 重采样后的音频，shape: (bs, raw_len) + // sampling_rate: 音频采样率，验证是否与模型的预处理采样率一致 + if sampling_rate != self.sampling_rate { + return Err(anyhow::anyhow!( + "The model feature extractor was trained sampling rate {} not equal to audio sample rate {}", + self.sampling_rate, + sampling_rate + )); + } + + let input_features = self.extract_fbank_features(raw_speech)?; + let mask_len = input_features.dim(2)?; + let mask = if return_attention_mask { + let mask = Tensor::new(1u32, input_features.device())?.broadcast_as((1, mask_len))?; + Some(mask) + } else { + None + }; + Ok((input_features, mask)) + } + + pub fn extract_fbank_features(&self, waveform: &Tensor) -> Result { + let mut waveform = waveform.clone(); + if self.dither != 0.0 { + waveform = waveform.add(&waveform.randn_like(0.0, 1.0)?.affine(self.dither, 0.0)?)?; + } + let pad = self.n_fft / 2; + let waveform = pad_reflect_last_dim(&waveform, (pad, pad))?; + let (_, samples) = waveform.dims2()?; + + let magnitudes = torch_stft(&waveform, self.n_fft, self.hop_length, &self.window)? + .transpose(D::Minus1, D::Minus2)?; + let n_frames = (samples - self.n_fft) / self.hop_length + 1; + let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?; + let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?; + let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?; + // let ln_spec = mel_spec.log()?; + // let log10_spec = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?; + let log10_spec = log10(&mel_spec)?; + let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?; + let log10_spec = log10_spec.broadcast_maximum(&max_val)?; + let log_spec = log10_spec.affine(1.0, 4.0)?.affine(1.0 / 4.0, 0.0)?; + Ok(log_spec) + } +} diff --git a/src/models/feature_extractor/mod.rs b/src/models/feature_extractor/mod.rs index 5fb1e5a..80016f5 100644 --- a/src/models/feature_extractor/mod.rs +++ b/src/models/feature_extractor/mod.rs @@ -1 +1,3 @@ -pub mod seamless_m4t_feature_extractor; \ No newline at end of file +pub mod seamless_m4t_feature_extractor; +pub mod feature_extraction_whisper; +pub mod config; \ No newline at end of file diff --git a/src/models/feature_extractor/seamless_m4t_feature_extractor.rs b/src/models/feature_extractor/seamless_m4t_feature_extractor.rs index aa371cb..a399fb7 100644 --- a/src/models/feature_extractor/seamless_m4t_feature_extractor.rs +++ b/src/models/feature_extractor/seamless_m4t_feature_extractor.rs @@ -7,7 +7,7 @@ use crate::utils::{ }; pub struct SeamlessM4TFeatureExtractor { - feature_size: usize, + // feature_size: usize, num_mel_bins: usize, padding_side: PaddingSide, padding_value: f32, @@ -19,7 +19,7 @@ pub struct SeamlessM4TFeatureExtractor { impl SeamlessM4TFeatureExtractor { pub fn new( - feature_size: usize, + // feature_size: usize, num_mel_bins: usize, padding_side: PaddingSide, padding_value: f32, @@ -40,7 +40,7 @@ impl SeamlessM4TFeatureExtractor { )?; let window = create_povey_window(400, candle_core::DType::F32, device)?; Ok(Self { - feature_size, + // feature_size, num_mel_bins, padding_side, padding_value, diff --git a/src/models/glm_asr_nano/config.rs b/src/models/glm_asr_nano/config.rs index dc29a7a..ea90498 100644 --- a/src/models/glm_asr_nano/config.rs +++ b/src/models/glm_asr_nano/config.rs @@ -1,6 +1,8 @@ use candle_nn::Activation; use serde::Deserialize; +use crate::models::feature_extractor::config::FeatureExtractor; + #[derive(Debug, Clone, PartialEq, Deserialize)] pub struct GlmAsrNanoProcessorConfig { pub audio_token: String, @@ -9,20 +11,7 @@ pub struct GlmAsrNanoProcessorConfig { pub max_audio_len: usize, } -#[derive(Debug, Clone, PartialEq, Deserialize)] -pub struct FeatureExtractor { - pub chunk_length: usize, - pub dither: f32, - pub feature_size: usize, - pub hop_length: usize, - pub n_fft: usize, - pub n_samples: usize, - pub nb_max_frames: usize, - pub padding_side: String, - pub padding_value: f32, - pub return_attention_mask: bool, - pub sampling_rate: usize, -} + #[derive(Debug, Clone, PartialEq, Deserialize)] pub struct GlmAsrNanoConfig { diff --git a/src/models/glm_asr_nano/processor.rs b/src/models/glm_asr_nano/processor.rs index 52783fa..55f0b30 100644 --- a/src/models/glm_asr_nano/processor.rs +++ b/src/models/glm_asr_nano/processor.rs @@ -5,7 +5,10 @@ use anyhow::Result; use candle_core::{D, DType, Device, IndexOp, Tensor}; use crate::{ - models::glm_asr_nano::config::GlmAsrNanoProcessorConfig, + models::{ + feature_extractor::feature_extraction_whisper::WhisperFeatureExtractor, + glm_asr_nano::config::GlmAsrNanoProcessorConfig, + }, utils::{ audio_utils::{ apply_stft, create_hann_window, extract_audios, extract_frames, mel_filter_bank, @@ -28,6 +31,7 @@ pub struct GlmAsrNanoProcessor { max_audio_len: usize, // default_transcription_prompt: String, device: Device, + whisper_feature_extrator: WhisperFeatureExtractor, } impl GlmAsrNanoProcessor { @@ -67,6 +71,16 @@ impl GlmAsrNanoProcessor { device, )? .t()?; + let whisper_feature_extrator = WhisperFeatureExtractor::new( + processor_cfg.feature_extractor.feature_size, + processor_cfg.feature_extractor.hop_length, + processor_cfg.feature_extractor.chunk_length, + processor_cfg.feature_extractor.n_fft, + processor_cfg.feature_extractor.dither, + processor_cfg.feature_extractor.padding_value, + processor_cfg.feature_extractor.sampling_rate, + device, + )?; Ok(Self { sampling_rate, chunk_length, @@ -80,34 +94,35 @@ impl GlmAsrNanoProcessor { max_audio_len, // default_transcription_prompt, device: device.clone(), + whisper_feature_extrator, }) } - pub fn extract_fbank_features(&self, waveform: &Tensor) -> Result { - let pad = self.n_fft / 2; - let waveform = pad_reflect_last_dim(waveform, (pad, pad))?; - let (_, samples) = waveform.dims2()?; + // pub fn extract_fbank_features(&self, waveform: &Tensor) -> Result { + // let pad = self.n_fft / 2; + // let waveform = pad_reflect_last_dim(waveform, (pad, pad))?; + // let (_, samples) = waveform.dims2()?; - // // (bs, n_frames, n_fft) - // let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?; - // // 应用汉明窗口 - // let result = frames.broadcast_mul(&self.window)?; - // // 傅立叶变换 - // let magnitudes = apply_stft(&result)?.transpose(D::Minus1, D::Minus2)?; - let magnitudes = torch_stft(&waveform, self.n_fft, self.hop_length, &self.window)? - .transpose(D::Minus1, D::Minus2)?; - let n_frames = (samples - self.n_fft) / self.hop_length + 1; - let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?; - let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?; - let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?; - // let ln_spec = mel_spec.log()?; - // let log10_spec = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?; - let log10_spec = log10(&mel_spec)?; - let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?; - let log10_spec = log10_spec.broadcast_maximum(&max_val)?; - let log_spec = log10_spec.affine(1.0, 4.0)?.affine(1.0 / 4.0, 0.0)?; - Ok(log_spec) - } + // // // (bs, n_frames, n_fft) + // // let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?; + // // // 应用汉明窗口 + // // let result = frames.broadcast_mul(&self.window)?; + // // // 傅立叶变换 + // // let magnitudes = apply_stft(&result)?.transpose(D::Minus1, D::Minus2)?; + // let magnitudes = torch_stft(&waveform, self.n_fft, self.hop_length, &self.window)? + // .transpose(D::Minus1, D::Minus2)?; + // let n_frames = (samples - self.n_fft) / self.hop_length + 1; + // let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?; + // let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?; + // let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?; + // // let ln_spec = mel_spec.log()?; + // // let log10_spec = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?; + // let log10_spec = log10(&mel_spec)?; + // let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?; + // let log10_spec = log10_spec.broadcast_maximum(&max_val)?; + // let log_spec = log10_spec.affine(1.0, 4.0)?.affine(1.0 / 4.0, 0.0)?; + // Ok(log_spec) + // } pub fn feature_extractor(&self, raw_speech: Vec) -> Result<(Tensor, Tensor)> { let mut pad_audio = vec![]; @@ -132,7 +147,10 @@ impl GlmAsrNanoProcessor { } let input_features = Tensor::cat(&pad_audio, 0)?; let input_features_mask = Tensor::new(input_features_mask, input_features.device())?; - let input_features = self.extract_fbank_features(&input_features)?; + // let input_features = self.extract_fbank_features(&input_features)?; + let (input_features, _) = + self.whisper_feature_extrator + .call(&input_features, self.sampling_rate, false)?; let (_, audio_len) = input_features_mask.dims2()?; let mask_idx: Vec = (0..audio_len) .step_by(self.hop_length) diff --git a/src/models/index_tts2/processor.rs b/src/models/index_tts2/processor.rs index 7b92580..e10e9e8 100644 --- a/src/models/index_tts2/processor.rs +++ b/src/models/index_tts2/processor.rs @@ -45,7 +45,7 @@ impl IndexTTS2Processor { dtype: DType, ) -> Result { let feature_extractor = SeamlessM4TFeatureExtractor::new( - 80, + // 80, 80, crate::utils::tensor_utils::PaddingSide::Right, 1.0, diff --git a/src/models/mod.rs b/src/models/mod.rs index 0299553..8505a39 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -11,6 +11,7 @@ pub mod minicpm4; pub mod paddleocr_vl; pub mod qwen2_5vl; pub mod qwen3; +pub mod qwen3_asr; pub mod qwen3vl; pub mod rmbg2_0; pub mod voxcpm; diff --git a/src/models/qwen3_asr/model.rs b/src/models/qwen3_asr/model.rs index d3f971f..13f4ba6 100644 --- a/src/models/qwen3_asr/model.rs +++ b/src/models/qwen3_asr/model.rs @@ -11,7 +11,7 @@ use crate::{ qwen3::{config::Qwen3Config, model::Qwen3Model}, }, position_embed::sinusoidal_pe::SinusoidalPositionEncoderCat, - utils::tensor_utils::{get_equal_mask, mask_filled, masked_scatter_dim0}, + utils::tensor_utils::{get_equal_mask, attn_masked_fill, masked_scatter_dim0}, }; pub struct MultiHeadedAttentionSANM { @@ -124,9 +124,9 @@ impl MultiHeadedAttentionSANM { }; // mask: rank = 2 let mask = get_equal_mask(&mask, 0)?; - let scores = mask_filled(scores, &mask, f32::NEG_INFINITY)?; + let scores = attn_masked_fill(scores, &mask, f32::NEG_INFINITY)?; let attn = softmax_last_dim(&scores)?; - mask_filled(&attn, &mask, 0.0)? + attn_masked_fill(&attn, &mask, 0.0)? } else { softmax_last_dim(scores)? }; diff --git a/src/utils/audio_utils.rs b/src/utils/audio_utils.rs index 3b89a72..1ec8ace 100644 --- a/src/utils/audio_utils.rs +++ b/src/utils/audio_utils.rs @@ -32,7 +32,7 @@ use symphonia::core::meta::MetadataOptions; 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}; +use crate::utils::tensor_utils::{linspace, log10, pad_reflect_last_dim, pad_replicate_last_dim, split_tensor}; // 重采样方法枚举 #[derive(Debug, Clone, Copy)] @@ -1606,3 +1606,22 @@ pub fn spectrogram( } Ok(spectrogram) } + +pub fn split_audio_into_chunks(wav: &Tensor, sr: usize, max_chunk_sec: f32) -> Result> { + // wav: (1, len) + let total_len = wav.dim(1)?; + let total_sec = total_len as f32 / sr as f32; + let mut wavs = vec![]; + if total_sec <= max_chunk_sec { + wavs.push(wav.clone()); + } else { + let max_len = (max_chunk_sec * sr as f32).round() as usize; + let split_len = total_len / max_len; + let mut splits = vec![max_len; split_len]; + let remain_len = total_len % max_len; + splits.push(remain_len); + let split_wav = split_tensor(wav, &splits, 1)?; + wavs.extend_from_slice(&split_wav); + } + Ok(wavs) +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index a1bf32a..fad6422 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -712,6 +712,25 @@ pub fn extract_user_text(mes: &ChatCompletionParameters) -> Result { Ok(ret) } +pub fn extract_user_text_vec(mes: &ChatCompletionParameters) -> Result> { + let mut ret = vec![]; + for chat_mes in mes.messages.clone() { + if let ChatMessage::User { content, .. } = chat_mes.clone() + && let ChatMessageContent::ContentPart(part_vec) = content + { + for part in part_vec { + if let ChatMessageContentPart::Text(text_part) = part { + let text = text_part.text; + if text.chars().count() > 0 { + ret.push(text); + } + } + } + } + } + Ok(ret) +} + pub fn get_default_save_dir() -> Option { home_dir().map(|mut path| { path.push(".aha"); @@ -772,3 +791,14 @@ pub fn get_file_path(file: &str) -> Result { }; Ok(path) } + +pub fn capitalize_first_letter(input: &str) -> String { + if input.is_empty() { + return input.to_string(); + } + + let mut chars = input.chars(); + let first_char = chars.next().unwrap().to_uppercase().collect::(); + let remaining = chars.as_str().to_lowercase(); + format!("{}{}", first_char, remaining) +} diff --git a/src/utils/tensor_utils.rs b/src/utils/tensor_utils.rs index ce0a9cf..fbb0deb 100644 --- a/src/utils/tensor_utils.rs +++ b/src/utils/tensor_utils.rs @@ -1043,3 +1043,19 @@ pub fn statistics_pooling(xs: &Tensor, dim: D, keepdim: bool) -> Result } Ok(stats) } +pub fn float_range_normalize(t: &Tensor) -> Result { + let peak = t + .to_dtype(DType::F32)? + .abs()? + .max_all()? + .to_scalar::()?; + if peak == 0.0 { + return Ok(t.clone()); + } + let mut t = t.clone(); + if peak > 1.0 { + t = t.affine(1.0 / peak as f64, 0.0)?; + } + t = t.clamp(-1.0, 1.0)?; + Ok(t) +}