merge main

src/chat_template/mod.rs

@@ -6,16 +6,33 @@ use crate::utils::string_to_static_str;
 
 pub fn get_template(path: String) -> Result<String> {
     let tokenizer_config_file = path.clone() + "/tokenizer_config.json";
-    assert!(
+    let chat_template = if std::path::Path::new(&tokenizer_config_file).exists() {
-        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")))?;
+            .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(

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
+}

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<Self> {
+        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<Tensor>)> {
+        // 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<Tensor> {
+        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)
+    }
+}

src/models/feature_extractor/mod.rs

@@ -1 +1,3 @@
 pub mod seamless_m4t_feature_extractor;
+pub mod feature_extraction_whisper;
+pub mod config;

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,

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 {

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<Tensor> {
+    // pub fn extract_fbank_features(&self, waveform: &Tensor) -> Result<Tensor> {
-        let pad = self.n_fft / 2;
+    //     let pad = self.n_fft / 2;
-        let waveform = pad_reflect_last_dim(waveform, (pad, pad))?;
+    //     let waveform = pad_reflect_last_dim(waveform, (pad, pad))?;
-        let (_, samples) = waveform.dims2()?;
+    //     let (_, samples) = waveform.dims2()?;
 
-        // // (bs, n_frames, n_fft)
+    //     // // (bs, n_frames, n_fft)
-        // let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?;
+    //     // let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?;
-        // // 应用汉明窗口
+    //     // // 应用汉明窗口
-        // let result = frames.broadcast_mul(&self.window)?;
+    //     // let result = frames.broadcast_mul(&self.window)?;
-        // // 傅立叶变换
+    //     // // 傅立叶变换
-        // let magnitudes = apply_stft(&result)?.transpose(D::Minus1, D::Minus2)?;
+    //     // let magnitudes = apply_stft(&result)?.transpose(D::Minus1, D::Minus2)?;
-        let magnitudes = torch_stft(&waveform, self.n_fft, self.hop_length, &self.window)?
+    //     let magnitudes = torch_stft(&waveform, self.n_fft, self.hop_length, &self.window)?
-            .transpose(D::Minus1, D::Minus2)?;
+    //         .transpose(D::Minus1, D::Minus2)?;
-        let n_frames = (samples - self.n_fft) / self.hop_length + 1;
+    //     let n_frames = (samples - self.n_fft) / self.hop_length + 1;
-        let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?;
+    //     let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?;
-        let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?;
+    //     let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?;
-        let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?;
+    //     let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?;
-        // let ln_spec = mel_spec.log()?;
+    //     // 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 = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?;
-        let log10_spec = log10(&mel_spec)?;
+    //     let log10_spec = log10(&mel_spec)?;
-        let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?;
+    //     let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?;
-        let log10_spec = log10_spec.broadcast_maximum(&max_val)?;
+    //     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)?;
+    //     let log_spec = log10_spec.affine(1.0, 4.0)?.affine(1.0 / 4.0, 0.0)?;
-        Ok(log_spec)
+    //     Ok(log_spec)
-    }
+    // }
 
     pub fn feature_extractor(&self, raw_speech: Vec<Tensor>) -> 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<u32> = (0..audio_len)
             .step_by(self.hop_length)

src/models/index_tts2/processor.rs

@@ -45,7 +45,7 @@ impl IndexTTS2Processor {
         dtype: DType,
     ) -> Result<Self> {
         let feature_extractor = SeamlessM4TFeatureExtractor::new(
-            80,
+            // 80,
             80,
             crate::utils::tensor_utils::PaddingSide::Right,
             1.0,

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;

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)?
         };

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<Vec<Tensor>> {
+    // 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)
+}

src/utils/mod.rs

@@ -712,6 +712,25 @@ pub fn extract_user_text(mes: &ChatCompletionParameters) -> Result<String> {
     Ok(ret)
 }
 
+pub fn extract_user_text_vec(mes: &ChatCompletionParameters) -> Result<Vec<String>> {
+    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<String> {
     home_dir().map(|mut path| {
         path.push(".aha");
@@ -772,3 +791,14 @@ pub fn get_file_path(file: &str) -> Result<PathBuf> {
     };
     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::<String>();
+    let remaining = chars.as_str().to_lowercase();
+    format!("{}{}", first_char, remaining)
+}

src/utils/tensor_utils.rs

@@ -1043,3 +1043,19 @@ pub fn statistics_pooling(xs: &Tensor, dim: D, keepdim: bool) -> Result<Tensor>
     }
     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)
+}