diff --git a/CHANGELOG.md b/CHANGELOG.md index 196cfc7..b859a3e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Changelog +## [Unreleased] - 2025-02-04 + +### Added +- Support for Qwen3-ASR model ## [Unreleased] diff --git a/download_and_run.sh b/download_and_run.sh index 814667f..428638b 100755 --- a/download_and_run.sh +++ b/download_and_run.sh @@ -17,7 +17,9 @@ show_help() { echo " minicpm4-0.5b" echo " qwen2.5vl-3b" echo " qwen2.5vl-7b" - echo " qwen3-0.6b" + echo " qwen3-0.6b" + echo " qwen3asr-0.6b" + echo " qwen3asr-1.7b" echo " qwen3vl-2b" echo " qwen3vl-4b" echo " qwen3vl-8b" @@ -56,6 +58,12 @@ case $MODEL_ALIAS in "qwen3-0.6b") MODEL_ID="Qwen/Qwen3-0.6B" ;; + "qwen3asr-0.6b") + MODEL_ID="Qwen/Qwen3-ASR-0.6B" + ;; + "qwen3asr-1.7b") + MODEL_ID="Qwen/Qwen3-ASR-1.7B" + ;; "qwen3vl-2b") MODEL_ID="Qwen/Qwen3-VL-2B-Instruct" ;; diff --git a/src/exec/mod.rs b/src/exec/mod.rs index 67d9735..5eb4708 100644 --- a/src/exec/mod.rs +++ b/src/exec/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/exec/qwen3_asr.rs b/src/exec/qwen3_asr.rs new file mode 100644 index 0000000..e75ed19 --- /dev/null +++ b/src/exec/qwen3_asr.rs @@ -0,0 +1,67 @@ +//! Fun-ASR-Nano-2512 exec implementation for CLI `run` subcommand + +use std::time::Instant; + +use anyhow::{Ok, Result}; + +use crate::exec::ExecModel; +use crate::models::qwen3_asr::generate::Qwen3AsrGenerateModel; +use crate::models::{GenerateModel}; + +pub struct Qwen3ASRExec; + +impl ExecModel for Qwen3ASRExec { + fn run(input: &[String], output: Option<&str>, weight_path: &str) -> Result<()> { + + let i_start = Instant::now(); + let mut model = Qwen3AsrGenerateModel::init(weight_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + + // Create ChatCompletionParameters for ASR + let url = &input[1]; + let input_url = if url.starts_with("http://") + || url.starts_with("https://") + || url.starts_with("file://") + { + url.clone() + } else { + format!("file://{}", url) + }; + + let message = format!( + r#"{{ + "model": "fun-asr-nano", + "messages": [ + {{ + "role": "user", + "content": [ + {{ + "type": "audio", + "audio_url": {{ + "url": "{}" + }} + }} + ] + }} + ] + }}"#, + input_url + ); + let mes = serde_json::from_str(&message)?; + + let i_start = Instant::now(); + let res = model.generate(mes)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in generate is: {:?}", i_duration); + + println!("Result: {:?}", res); + + if let Some(out) = output { + std::fs::write(out, format!("{:?}", res))?; + println!("Output saved to: {}", out); + } + + Ok(()) + } +} diff --git a/src/main.rs b/src/main.rs index 1cbc4be..2a459c4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,15 +1,15 @@ -use std::{net::IpAddr, str::FromStr, time::Duration}; - -use aha::{models::WhichModel, utils::{download_model, get_default_save_dir}}; +use std::{net::IpAddr, str::FromStr}; +use aha::{ + models::WhichModel, + utils::{download_model, get_default_save_dir}, +}; use clap::{Args, Parser, Subcommand, ValueEnum}; -use modelscope::ModelScope; use rocket::{ Config, data::{ByteUnit, Limits}, routes, }; -use tokio::time::sleep; use crate::api::init; mod api; @@ -149,6 +149,8 @@ fn get_model_id(model: WhichModel) -> &'static str { WhichModel::Qwen2_5vl3B => "Qwen/Qwen2.5-VL-3B-Instruct", WhichModel::Qwen2_5vl7B => "Qwen/Qwen2.5-VL-7B-Instruct", WhichModel::Qwen3_0_6B => "Qwen/Qwen3-0.6B", + WhichModel::Qwen3ASR0_6B => "Qwen/Qwen3-ASR-0.6B", + WhichModel::Qwen3ASR1_7B => "Qwen/Qwen3-ASR-1.7B", WhichModel::Qwen3vl2B => "Qwen/Qwen3-VL-2B-Instruct", WhichModel::Qwen3vl4B => "Qwen/Qwen3-VL-4B-Instruct", WhichModel::Qwen3vl8B => "Qwen/Qwen3-VL-8B-Instruct", @@ -171,6 +173,8 @@ fn run_list() -> anyhow::Result<()> { WhichModel::Qwen2_5vl3B, WhichModel::Qwen2_5vl7B, WhichModel::Qwen3_0_6B, + WhichModel::Qwen3ASR0_6B, + WhichModel::Qwen3ASR1_7B, WhichModel::Qwen3vl2B, WhichModel::Qwen3vl4B, WhichModel::Qwen3vl8B, @@ -289,6 +293,14 @@ fn run_run(args: RunArgs) -> anyhow::Result<()> { use aha::exec::qwen3::Qwen3Exec; Qwen3Exec::run(&input, output.as_deref(), &weight_path)?; } + WhichModel::Qwen3ASR0_6B => { + use aha::exec::qwen3_asr::Qwen3ASRExec; + Qwen3ASRExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3ASR1_7B => { + use aha::exec::qwen3_asr::Qwen3ASRExec; + Qwen3ASRExec::run(&input, output.as_deref(), &weight_path)?; + } WhichModel::Qwen3vl2B => { use aha::exec::qwen3vl::Qwen3vlExec; Qwen3vlExec::run(&input, output.as_deref(), &weight_path)?; diff --git a/src/models/feature_extractor/feature_extraction_whisper.rs b/src/models/feature_extractor/feature_extraction_whisper.rs index f105fe3..507005e 100644 --- a/src/models/feature_extractor/feature_extraction_whisper.rs +++ b/src/models/feature_extractor/feature_extraction_whisper.rs @@ -7,13 +7,13 @@ use crate::utils::{ }; pub struct WhisperFeatureExtractor { - feature_size: usize, + // feature_size: usize, hop_length: usize, - chunk_length: usize, - n_samples: usize, + // chunk_length: usize, + // n_samples: usize, n_fft: usize, dither: f64, - padding_value: f32, + // padding_value: f32, sampling_rate: usize, mel_filters: Tensor, window: Tensor, @@ -23,10 +23,10 @@ impl WhisperFeatureExtractor { pub fn new( feature_size: usize, hop_length: usize, - chunk_length: usize, + // chunk_length: usize, n_fft: usize, dither: f64, - padding_value: f32, + // padding_value: f32, sampling_rate: usize, device: &Device, ) -> Result { @@ -44,15 +44,15 @@ impl WhisperFeatureExtractor { device, )? .t()?; - let n_samples = chunk_length * sampling_rate; + // let n_samples = chunk_length * sampling_rate; Ok(Self { - feature_size, + // feature_size, hop_length, - chunk_length, - n_samples, + // chunk_length, + // n_samples, n_fft, dither, - padding_value, + // padding_value, sampling_rate, mel_filters, window, diff --git a/src/models/glm_asr_nano/processor.rs b/src/models/glm_asr_nano/processor.rs index 55f0b30..e20222d 100644 --- a/src/models/glm_asr_nano/processor.rs +++ b/src/models/glm_asr_nano/processor.rs @@ -1,4 +1,3 @@ -use std::f32; use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; use anyhow::Result; @@ -9,22 +8,16 @@ use crate::{ 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, - torch_stft, - }, - tensor_utils::{log10, pad_reflect_last_dim, split_tensor}, - }, + utils::{audio_utils::extract_audios, tensor_utils::split_tensor}, }; pub struct GlmAsrNanoProcessor { sampling_rate: usize, chunk_length: usize, n_samples: usize, - n_fft: usize, - window: Tensor, - mel_filters: Tensor, + // n_fft: usize, + // window: Tensor, + // mel_filters: Tensor, hop_length: usize, audio_token: String, // audio_token_id: u32, @@ -55,29 +48,29 @@ impl GlmAsrNanoProcessor { let sampling_rate = processor_cfg.feature_extractor.sampling_rate; let chunk_length = processor_cfg.feature_extractor.chunk_length; let n_samples = processor_cfg.feature_extractor.n_samples; - let n_fft = processor_cfg.feature_extractor.n_fft; + // let n_fft = processor_cfg.feature_extractor.n_fft; let hop_length = processor_cfg.feature_extractor.hop_length; - let window = create_hann_window(n_fft, dtype, device)?; - let window = window.unsqueeze(0)?.unsqueeze(0)?; - let mel_filters = mel_filter_bank( - 1 + n_fft / 2, - processor_cfg.feature_extractor.feature_size, - 0.0, - 8000.0, - sampling_rate as f32, - Some("slaney"), - crate::utils::audio_utils::MelScale::Slaney, - false, - device, - )? - .t()?; + // let window = create_hann_window(n_fft, dtype, device)?; + // let window = window.unsqueeze(0)?.unsqueeze(0)?; + // let mel_filters = mel_filter_bank( + // 1 + n_fft / 2, + // processor_cfg.feature_extractor.feature_size, + // 0.0, + // 8000.0, + // sampling_rate as f32, + // Some("slaney"), + // crate::utils::audio_utils::MelScale::Slaney, + // false, + // 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.chunk_length, processor_cfg.feature_extractor.n_fft, processor_cfg.feature_extractor.dither, - processor_cfg.feature_extractor.padding_value, + // processor_cfg.feature_extractor.padding_value, processor_cfg.feature_extractor.sampling_rate, device, )?; @@ -85,9 +78,9 @@ impl GlmAsrNanoProcessor { sampling_rate, chunk_length, n_samples, - n_fft, - window, - mel_filters, + // n_fft, + // window, + // mel_filters, hop_length, audio_token, // audio_token_id, diff --git a/src/models/index_tts2/processor.rs b/src/models/index_tts2/processor.rs index e10e9e8..88c8e28 100644 --- a/src/models/index_tts2/processor.rs +++ b/src/models/index_tts2/processor.rs @@ -10,7 +10,7 @@ use crate::{ utils::{ audio_utils::{ create_hann_window, extract_audio_url, get_waveform_and_window_properties, kaldi_fbank, - kaldi_get_mel_banks, load_audio, mel_filter_bank, resample_simple, spectrogram, + kaldi_get_mel_banks, load_audio, mel_filter_bank, resample_simple, torch_stft, }, get_vb_model_path, diff --git a/src/models/mask_gct/model.rs b/src/models/mask_gct/model.rs index c583e58..7dd9c8f 100644 --- a/src/models/mask_gct/model.rs +++ b/src/models/mask_gct/model.rs @@ -1,5 +1,5 @@ use anyhow::Result; -use candle_core::{D, IndexOp, Tensor}; +use candle_core::{D, Tensor}; use candle_nn::{ Conv1d, Embedding, Init, LayerNorm, Linear, Module, VarBuilder, embedding, linear, }; diff --git a/src/models/mod.rs b/src/models/mod.rs index 8505a39..e0fa3d2 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,5 +1,5 @@ -pub mod common; pub mod campplus; +pub mod common; pub mod deepseek_ocr; pub mod feature_extractor; pub mod fun_asr_nano; @@ -29,8 +29,9 @@ use crate::models::{ glm_asr_nano::generate::GlmAsrNanoGenerateModel, hunyuan_ocr::generate::HunyuanOCRGenerateModel, minicpm4::generate::MiniCPMGenerateModel, paddleocr_vl::generate::PaddleOCRVLGenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel, - qwen3::generate::Qwen3GenerateModel, qwen3vl::generate::Qwen3VLGenerateModel, - rmbg2_0::generate::RMBG2_0Model, voxcpm::generate::VoxCPMGenerate, + qwen3::generate::Qwen3GenerateModel, qwen3_asr::generate::Qwen3AsrGenerateModel, + qwen3vl::generate::Qwen3VLGenerateModel, rmbg2_0::generate::RMBG2_0Model, + voxcpm::generate::VoxCPMGenerate, }; #[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] @@ -43,7 +44,11 @@ pub enum WhichModel { Qwen2_5vl7B, #[value(name = "qwen3-0.6b", hide = true)] Qwen3_0_6B, - #[value(name = "qwen3vl-2b", hide = true)] + #[value(name = "qwen3asr-0.6b", hide = true)] + Qwen3ASR0_6B, + #[value(name = "qwen3asr-1.7b", hide = true)] + Qwen3ASR1_7B, + #[value(name = "qwen3vl-4b", hide = true)] Qwen3vl2B, #[value(name = "qwen3vl-4b", hide = true)] Qwen3vl4B, @@ -88,6 +93,7 @@ pub enum ModelInstance<'a> { MiniCPM4(MiniCPMGenerateModel<'a>), Qwen2_5VL(Qwen2_5VLGenerateModel<'a>), Qwen3(Qwen3GenerateModel<'a>), + Qwen3ASR(Qwen3AsrGenerateModel<'a>), Qwen3VL(Qwen3VLGenerateModel<'a>), DeepSeekOCR(DeepseekOCRGenerateModel), HunyuanOCR(HunyuanOCRGenerateModel<'a>), @@ -104,6 +110,7 @@ impl<'a> GenerateModel for ModelInstance<'a> { ModelInstance::MiniCPM4(model) => model.generate(mes), ModelInstance::Qwen2_5VL(model) => model.generate(mes), ModelInstance::Qwen3(model) => model.generate(mes), + ModelInstance::Qwen3ASR(model) => model.generate(mes), ModelInstance::Qwen3VL(model) => model.generate(mes), ModelInstance::DeepSeekOCR(model) => model.generate(mes), ModelInstance::HunyuanOCR(model) => model.generate(mes), @@ -131,6 +138,7 @@ impl<'a> GenerateModel for ModelInstance<'a> { ModelInstance::Qwen2_5VL(model) => model.generate_stream(mes), ModelInstance::Qwen3(model) => model.generate_stream(mes), ModelInstance::Qwen3VL(model) => model.generate_stream(mes), + ModelInstance::Qwen3ASR(model) => model.generate_stream(mes), ModelInstance::DeepSeekOCR(model) => model.generate_stream(mes), ModelInstance::HunyuanOCR(model) => model.generate_stream(mes), ModelInstance::PaddleOCRVL(model) => model.generate_stream(mes), @@ -160,6 +168,14 @@ pub fn load_model(model_type: WhichModel, path: &str) -> Result { + let model = Qwen3AsrGenerateModel::init(path, None, None)?; + ModelInstance::Qwen3ASR(model) + } + WhichModel::Qwen3ASR1_7B => { + let model = Qwen3AsrGenerateModel::init(path, None, None)?; + ModelInstance::Qwen3ASR(model) + } WhichModel::Qwen3vl2B => { let model = Qwen3VLGenerateModel::init(path, None, None)?; ModelInstance::Qwen3VL(model) diff --git a/src/models/qwen3_asr/config.rs b/src/models/qwen3_asr/config.rs index b023260..7373fde 100644 --- a/src/models/qwen3_asr/config.rs +++ b/src/models/qwen3_asr/config.rs @@ -1,86 +1,202 @@ -use serde::Deserialize; +use candle_nn::Activation; -#[derive(Debug, Clone, PartialEq, Deserialize)] -pub struct FunASRNanoConfig { - pub audio_encoder_conf: AudioEncoderConf, - pub llm_conf: LlmConf, - pub audio_adaptor_conf: AudioAdaptorConf, - pub detach_ctc_decoder: bool, - pub ctc_decoder_conf: CtcDecoderConf, - pub ctc_weight: f64, - pub ctc_conf: CtcConf, - pub frontend_conf: FrontendConf, +use crate::models::qwen3::config::Qwen3Config; + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Qwen3ASRConfig { + pub model_type: String, + pub support_languages: Vec, + pub thinker_config: ThinkerConfig, } -#[derive(Debug, Clone, PartialEq, Deserialize)] -pub struct AudioEncoderConf { - pub output_size: usize, - pub attention_heads: usize, - pub linear_units: usize, - pub num_blocks: usize, - pub tp_blocks: usize, - pub dropout_rate: f64, - pub positional_dropout_rate: f64, - pub attention_dropout_rate: f64, - pub input_layer: String, - pub pos_enc_class: String, - pub normalize_before: bool, - pub kernel_size: usize, - pub sanm_shfit: usize, - pub selfattention_layer_type: String, - pub freeze: bool, - pub freeze_layer_num: i32, - pub feat_permute: bool, +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct ThinkerConfig { + pub model_type: String, + pub audio_config: Qwen3ASRAudioConfig, + pub audio_end_token_id: u32, + pub audio_start_token_id: u32, + pub audio_token_id: u32, + pub dtype: String, + pub initializer_range: f64, + pub text_config: Qwen3ASRTextConfig, } -#[derive(Debug, Clone, PartialEq, Deserialize)] -pub struct LlmConf { - pub hub: String, - pub freeze: bool, - pub llm_dtype: String, - pub init_param_path: String, +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Qwen3ASRAudioConfig { + pub activation_dropout: f32, + pub activation_function: Activation, + pub add_cross_attention: bool, + pub attention_dropout: f32, + pub bad_words_ids: Option>, + pub begin_suppress_tokens: Option>, + pub bos_token_id: Option, + pub chunk_size_feed_forward: usize, + pub conv_chunksize: usize, + pub cross_attention_hidden_size: Option, + pub d_model: usize, + pub decoder_start_token_id: Option, + pub diversity_penalty: f64, + pub do_sample: bool, + pub downsample_hidden_size: usize, + pub dropout: f32, + pub early_stopping: bool, + pub encoder_attention_heads: usize, + pub encoder_ffn_dim: usize, + pub encoder_layers: usize, + pub encoder_no_repeat_ngram_size: usize, + pub eos_token_id: Option, + pub exponential_decay_length_penalty: Option>, + pub forced_bos_token_id: Option, + pub forced_eos_token_id: Option, + pub id2label: std::collections::HashMap, + pub initializer_range: f64, + pub is_decoder: bool, + pub is_encoder_decoder: bool, + pub label2id: std::collections::HashMap, + pub length_penalty: f64, + pub max_length: usize, + pub max_source_positions: usize, + pub min_length: usize, + pub model_type: String, + pub n_window: usize, + pub n_window_infer: usize, + pub no_repeat_ngram_size: usize, + pub num_beam_groups: usize, + pub num_beams: usize, + pub num_hidden_layers: usize, + pub num_mel_bins: usize, + pub num_return_sequences: usize, + pub output_attentions: bool, + pub output_dim: usize, + pub output_hidden_states: bool, + pub output_scores: bool, + pub pad_token_id: Option, + pub prefix: Option, + pub problem_type: Option, + pub pruned_heads: std::collections::HashMap>, + pub remove_invalid_values: bool, + pub repetition_penalty: f64, + pub return_dict: bool, + pub return_dict_in_generate: bool, + pub scale_embedding: bool, + pub sep_token_id: Option, + pub suppress_tokens: Option>, + pub task_specific_params: Option>, + pub temperature: f64, + pub tf_legacy_loss: bool, + pub tie_encoder_decoder: bool, + pub tie_word_embeddings: bool, + pub tokenizer_class: Option, + pub top_k: usize, + pub top_p: f64, + pub torchscript: bool, + pub typical_p: f64, + pub use_bfloat16: bool, } -#[derive(Debug, Clone, PartialEq, Deserialize)] -pub struct AudioAdaptorConf { - pub downsample_rate: usize, - pub use_low_frame_rate: bool, - pub ffn_dim: usize, - pub llm_dim: usize, - pub encoder_dim: usize, - pub n_layer: usize, - pub freeze: bool, +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Qwen3ASRTextConfig { + pub add_cross_attention: bool, + pub attention_bias: bool, + pub attention_dropout: f64, + pub bad_words_ids: Option>, + pub begin_suppress_tokens: Option>, + pub bos_token_id: Option, + pub chunk_size_feed_forward: usize, + pub cross_attention_hidden_size: Option, + pub decoder_start_token_id: Option, + pub diversity_penalty: f64, + pub do_sample: bool, + pub dtype: Option, + pub early_stopping: bool, + pub encoder_no_repeat_ngram_size: usize, + pub eos_token_id: Option, + pub exponential_decay_length_penalty: Option>, + pub forced_bos_token_id: Option, + pub forced_eos_token_id: Option, + pub head_dim: usize, + pub hidden_act: Activation, + pub hidden_size: usize, + pub id2label: std::collections::HashMap, + pub initializer_range: f64, + pub intermediate_size: usize, + pub is_decoder: bool, + pub is_encoder_decoder: bool, + pub label2id: std::collections::HashMap, + pub length_penalty: f64, + pub max_length: usize, + pub max_position_embeddings: usize, + pub min_length: usize, + pub model_type: String, + pub no_repeat_ngram_size: usize, + pub num_attention_heads: usize, + pub num_beam_groups: usize, + pub num_beams: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub num_return_sequences: usize, + pub output_attentions: bool, + pub output_hidden_states: bool, + pub output_scores: bool, + pub pad_token_id: Option, + pub prefix: Option, + pub problem_type: Option, + pub remove_invalid_values: bool, + pub repetition_penalty: f64, + pub return_dict: bool, + pub return_dict_in_generate: bool, + pub rms_norm_eps: f64, + pub rope_scaling: Qwen3ASRRopeScaling, + pub rope_theta: f32, + pub sep_token_id: Option, + pub suppress_tokens: Option>, + pub temperature: f64, + pub tf_legacy_loss: bool, + pub tie_encoder_decoder: bool, + pub tie_word_embeddings: bool, + pub tokenizer_class: Option, + pub top_k: usize, + pub top_p: f64, + pub torchscript: bool, + pub typical_p: f64, + pub use_bfloat16: bool, + pub use_cache: bool, + pub vocab_size: usize, } -#[derive(Debug, Clone, PartialEq, Deserialize)] -pub struct CtcDecoderConf { - pub downsample_rate: u32, - pub ffn_dim: u32, - pub llm_dim: u32, - pub encoder_dim: u32, - pub n_layer: u32, - pub freeze: bool, +pub fn qwen3asr_text_config2qwen3_config(cfg: &Qwen3ASRTextConfig) -> Qwen3Config { + Qwen3Config { + attention_bias: cfg.attention_bias, + attention_dropout: cfg.attention_dropout as f64, + bos_token_id: cfg.bos_token_id.unwrap_or(151643) as u32, + eos_token_id: cfg.eos_token_id.unwrap_or(151645) as u32, + head_dim: cfg.head_dim, + hidden_act: cfg.hidden_act, + hidden_size: cfg.hidden_size, + initializer_range: cfg.initializer_range as f64, + intermediate_size: cfg.intermediate_size, + max_position_embeddings: cfg.max_position_embeddings, + max_window_layers: 0, + num_attention_heads: cfg.num_attention_heads, + num_hidden_layers: cfg.num_hidden_layers, + num_key_value_heads: cfg.num_key_value_heads, + rms_norm_eps: cfg.rms_norm_eps, + rope_theta: cfg.rope_theta, + tie_word_embeddings: true, + torch_dtype: "bfloat16".to_string(), + use_cache: cfg.use_cache, + use_sliding_window: false, + vocab_size: cfg.vocab_size, + } } -#[derive(Debug, Clone, PartialEq, Deserialize)] -pub struct CtcConf { - pub dropout_rate: f64, - pub ctc_type: String, - pub reduce: bool, - pub ignore_nan_grad: bool, -} - -#[derive(Debug, Clone, PartialEq, Deserialize)] -pub struct FrontendConf { - pub fs: usize, - pub window: String, - pub n_mels: usize, - pub frame_length: f32, - pub frame_shift: f32, - pub lfr_m: usize, - pub lfr_n: usize, - #[serde(skip_serializing_if = "Option::is_none")] - pub cmvn_file: Option, +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Qwen3ASRRopeScaling { + pub interleaved: bool, + pub mrope_interleaved: bool, + pub mrope_section: Vec, + pub rope_type: String, + pub r#type: String, } #[derive(Debug, Clone, PartialEq, serde::Deserialize)] @@ -89,4 +205,4 @@ pub struct Qwen3ASRGenerationConfig { pub eos_token_id: Vec, pub pad_token_id: usize, pub temperature: f32, -} \ No newline at end of file +} diff --git a/src/models/qwen3_asr/generate.rs b/src/models/qwen3_asr/generate.rs index f1e12ab..22b70cd 100644 --- a/src/models/qwen3_asr/generate.rs +++ b/src/models/qwen3_asr/generate.rs @@ -1,34 +1,38 @@ -use std::collections::HashMap; use aha_openai_dive::v1::resources::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, }; use anyhow::{Result, anyhow}; -use candle_core::{DType, Device, Tensor, pickle::read_all_with_key}; +use candle_core::{DType, Device, Tensor}; use candle_nn::VarBuilder; use rocket::async_stream::stream; use rocket::futures::Stream; use crate::{ - chat_template::ChatTemplate, models::{ + chat_template::ChatTemplate, + models::{ GenerateModel, - fun_asr_nano::{ - config::FunASRNanoConfig, model::FunAsrNanoModel, + feature_extractor::config::FeatureExtractor, + qwen3_asr::{ + config::{Qwen3ASRConfig, Qwen3ASRGenerationConfig}, + model::Qwen3ASRModel, + processor::Qwen3AsrProcessor, }, - qwen3::config::{Qwen3Config, Qwen3GenerationConfig}, qwen3_asr::{config::Qwen3ASRGenerationConfig, processor::Qwen3AsrProcessor}, - }, tokenizer::TokenizerModel, utils::{ + }, + tokenizer::TokenizerModel, + utils::{ build_completion_chunk_response, build_completion_response, find_type_files, get_device, get_dtype, get_logit_processor, - } + }, }; pub struct Qwen3AsrGenerateModel<'a> { chat_template: ChatTemplate<'a>, - // tokenizer: TokenizerModel, + tokenizer: TokenizerModel, processor: Qwen3AsrProcessor, - // fun_asr_nano: FunAsrNanoModel, + qwen3_asr: Qwen3ASRModel, device: Device, - // dtype: DType, + dtype: DType, eos_token_id1: u32, eos_token_id2: u32, generation_config: Qwen3ASRGenerationConfig, @@ -38,28 +42,40 @@ pub struct Qwen3AsrGenerateModel<'a> { impl<'a> Qwen3AsrGenerateModel<'a> { pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { let chat_template = ChatTemplate::init(path)?; + let tokenizer = TokenizerModel::init(path)?; let generation_config_path = path.to_string() + "/generation_config.json"; let generation_config: Qwen3ASRGenerationConfig = serde_json::from_slice(&std::fs::read(generation_config_path)?)?; let device = get_device(device); - let processor = Qwen3AsrProcessor::new(&device)?; + let preprocess_config_path = path.to_string() + "/preprocessor_config.json"; + let preprocess_config: FeatureExtractor = + serde_json::from_slice(&std::fs::read(preprocess_config_path)?)?; + let processor = Qwen3AsrProcessor::new(&device, &preprocess_config)?; + let config_path = path.to_string() + "/config.json"; + let cfg: Qwen3ASRConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; + let cfg_dtype = cfg.thinker_config.dtype.as_str(); + let dtype = get_dtype(dtype, cfg_dtype); + let model_list = find_type_files(path, "safetensors")?; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, &device)? }; + let qwen3_asr = Qwen3ASRModel::new(vb, &cfg)?; - Ok(Self { chat_template, - // tokenizer, + tokenizer, processor, - // fun_asr_nano, + qwen3_asr, device, - // dtype, + dtype, eos_token_id1: generation_config.eos_token_id[0] as u32, eos_token_id2: generation_config.eos_token_id[1] as u32, generation_config, model_name: "qwen3-asr".to_string(), }) } +} - pub fn generate(&mut self, mes: ChatCompletionParameters) -> Result<()> { +impl<'a> GenerateModel for Qwen3AsrGenerateModel<'a> { + fn generate(&mut self, mes: ChatCompletionParameters) -> Result { let temperature = match mes.temperature { None => self.generation_config.temperature, Some(tem) => tem, @@ -68,145 +84,110 @@ impl<'a> Qwen3AsrGenerateModel<'a> { None => 34562u64, Some(s) => s as u64, }; - let mut logit_processor = - get_logit_processor(Some(temperature), mes.top_p, None, seed); + let mut logit_processor = get_logit_processor(Some(temperature), mes.top_p, None, seed); let render_text = self.chat_template.apply_chat_template(&mes)?; - let audio_data = - self.processor.process_info(&mes, &render_text)?; - // for audio in audio_data { - // let text = - // } - - Ok(()) + let audio_datas = self + .processor + .process_info(&mes, &render_text, &self.tokenizer)?; + let sample_len = mes.max_tokens.unwrap_or(1024); + let mut generate = Vec::new(); + for data in audio_datas.iter() { + let mut input_ids = data.input_ids.clone(); + let mut input_features = Some(data.input_features.clone().to_dtype(self.dtype)?); + let mut seq_len = input_ids.dim(1)?; + let mut seqlen_offset = 0; + for _ in 0..sample_len { + let logits = + self.qwen3_asr + .forward(&input_ids, seqlen_offset, input_features.as_ref())?; + let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; + let next_token = logit_processor.sample(&logits)?; + generate.push(next_token); + if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 { + break; + } + seqlen_offset += seq_len; + seq_len = 1; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; + input_features = None; + } + self.qwen3_asr.clear_kv_cache(); + } + let num_token = generate.len() as u32; + let res = self.tokenizer.token_decode(generate)?; + let response = build_completion_response(res, &self.model_name, Some(num_token)); + Ok(response) } + fn generate_stream( + &mut self, + mes: ChatCompletionParameters, + ) -> Result< + Box< + dyn Stream> + + Send + + Unpin + + '_, + >, + > { + let temperature = match mes.temperature { + None => self.generation_config.temperature, + Some(tem) => tem, + }; + let seed = match mes.seed { + None => 34562u64, + Some(s) => s as u64, + }; + let mut logit_processor = get_logit_processor(Some(temperature), mes.top_p, None, seed); + let render_text = self.chat_template.apply_chat_template(&mes)?; + let audio_datas = self + .processor + .process_info(&mes, &render_text, &self.tokenizer)?; + let sample_len = mes.max_tokens.unwrap_or(1024); + let stream = stream! { + let mut error_tokens = Vec::new(); + for data in audio_datas.iter() { + let mut input_ids = data.input_ids.clone(); + let mut input_features = Some(data.input_features.clone().to_dtype(self.dtype)?); + let mut seq_len = input_ids.dim(1)?; + let mut seqlen_offset = 0; + for _ in 0..sample_len { + let logits = + self.qwen3_asr + .forward(&input_ids, seqlen_offset, input_features.as_ref())?; + let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; + let next_token = logit_processor.sample(&logits)?; + let mut decode_ids = Vec::new(); + if !error_tokens.is_empty() { + decode_ids.extend_from_slice(&error_tokens); + } + decode_ids.push(next_token); + let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{e}")))?; + if decoded_token.contains("�") { + error_tokens.push(next_token); + if error_tokens.len() > 3 { + error_tokens.clear(); + } + seqlen_offset += seq_len; + seq_len = 1; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; + input_features = None; + continue; + } + error_tokens.clear(); + let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None); + yield Ok(chunk); + if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 { + break; + } + seqlen_offset += seq_len; + seq_len = 1; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; + input_features = None; + } + self.qwen3_asr.clear_kv_cache(); + } + }; + Ok(Box::new(Box::pin(stream))) + } } - -// impl<'a> GenerateModel for Qwen3AsrGenerateModel<'a> { -// fn generate(&mut self, mes: ChatCompletionParameters) -> Result { -// let temperature = match mes.temperature { -// None => self.generation_config.temperature, -// Some(tem) => tem, -// }; -// let seed = match mes.seed { -// None => 34562u64, -// Some(s) => s as u64, -// }; -// let mut logit_processor = -// get_logit_processor(Some(temperature), mes.top_p, None, seed); -// let audio_data = -// self.processor.process_info(&mes)?; -// for audio in audio_data { -// let text = -// } -// let mut speech = Some(speech.to_dtype(self.dtype)?); -// let mut fbank_mask = Some(&fbank_mask); -// let mut seq_len = input_ids.dim(1)?; -// let mut seqlen_offset = 0; -// let mut generate = Vec::new(); -// let sample_len = mes.max_tokens.unwrap_or(1024); -// for _ in 0..sample_len { -// let logits = self.fun_asr_nano.forward( -// &input_ids, -// speech.as_ref(), -// fbank_mask, -// seqlen_offset, -// )?; -// let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; -// let next_token = logit_processor.sample(&logits)?; -// generate.push(next_token); -// if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 { -// break; -// } -// seqlen_offset += seq_len; -// seq_len = 1; -// input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; -// speech = None; -// fbank_mask = None; -// } -// let num_token = generate.len() as u32; -// let res = self.tokenizer.token_decode(generate)?; -// self.fun_asr_nano.clear_kv_cache(); -// let response = build_completion_response(res, &self.model_name, Some(num_token)); -// Ok(response) -// } - -// fn generate_stream( -// &mut self, -// mes: ChatCompletionParameters, -// ) -> Result< -// Box< -// dyn Stream> -// + Send -// + Unpin -// + '_, -// >, -// > { -// let temperature = match mes.temperature { -// None => self.generation_config.temperature, -// Some(tem) => tem, -// }; -// let top_p = match mes.top_p { -// None => self.generation_config.top_p, -// Some(top_p) => top_p, -// }; -// let top_k = self.generation_config.top_k; -// let seed = match mes.seed { -// None => 34562u64, -// Some(s) => s as u64, -// }; -// let mut logit_processor = -// get_logit_processor(Some(temperature), Some(top_p), Some(top_k), seed); -// let (speech, fbank_mask, input_ids) = self.processor.process_info(&mes, &self.tokenizer)?; -// let mut seq_len = input_ids.dim(1)?; -// let mut seqlen_offset = 0; -// let sample_len = mes.max_tokens.unwrap_or(1024); -// let stream = stream! { -// let mut error_tokens = Vec::new(); -// let mut speech = Some(speech.to_dtype(self.dtype)?); -// let mut fbank_mask = Some(&fbank_mask); -// let mut input_ids = input_ids; -// for _ in 0..sample_len { -// let logits = self.fun_asr_nano.forward( -// &input_ids, -// speech.as_ref(), -// fbank_mask, -// seqlen_offset, -// )?; -// let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?; -// let next_token = logit_processor.sample(&logits)?; -// let mut decode_ids = Vec::new(); -// if !error_tokens.is_empty() { -// decode_ids.extend_from_slice(&error_tokens); -// } -// decode_ids.push(next_token); -// let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{e}")))?; -// if decoded_token.contains("�") { -// error_tokens.push(next_token); -// if error_tokens.len() > 3 { -// error_tokens.clear(); -// } -// seqlen_offset += seq_len; -// seq_len = 1; -// input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; -// speech = None; -// fbank_mask = None; -// continue; -// } -// error_tokens.clear(); -// let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None); -// yield Ok(chunk); -// if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 { -// break; -// } -// seqlen_offset += seq_len; -// seq_len = 1; -// input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; -// speech = None; -// fbank_mask = None; -// } -// self.fun_asr_nano.clear_kv_cache(); -// }; -// Ok(Box::new(Box::pin(stream))) -// } -// } diff --git a/src/models/qwen3_asr/model.rs b/src/models/qwen3_asr/model.rs index 13f4ba6..7aae11a 100644 --- a/src/models/qwen3_asr/model.rs +++ b/src/models/qwen3_asr/model.rs @@ -1,646 +1,388 @@ use anyhow::Result; -use candle_core::{D, IndexOp, Tensor}; -use candle_nn::{Conv1d, LayerNorm, Linear, Module, VarBuilder, linear, ops::softmax_last_dim}; +use candle_core::Tensor; +use candle_nn::{ + Activation, Conv2d, Embedding, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, + linear, linear_no_bias, rms_norm, +}; use crate::{ models::{ - common::{ - NaiveAttention, TwoLinearMLP, eager_attention_forward, get_conv1d, get_layer_norm, + common::{NaiveAttention, get_conv2d, get_layer_norm}, + qwen3::model::Qwen3DecoderLayer, + qwen3_asr::{ + config::{ + Qwen3ASRAudioConfig, Qwen3ASRConfig, Qwen3ASRTextConfig, ThinkerConfig, + qwen3asr_text_config2qwen3_config, + }, + processor::get_feat_extract_output_lengths, }, - fun_asr_nano::config::FunASRNanoConfig, - qwen3::{config::Qwen3Config, model::Qwen3Model}, }, - position_embed::sinusoidal_pe::SinusoidalPositionEncoderCat, - utils::tensor_utils::{get_equal_mask, attn_masked_fill, masked_scatter_dim0}, + position_embed::{ + rope::Qwen3VLTextRotaryEmbedding, sinusoidal_pe::SinusoidalPositionEncoderCat, + }, + utils::tensor_utils::{ + get_equal_mask, masked_scatter_dim0, prepare_causal_attention_mask, split_tensor, + split_tensor_with_size, + }, }; -pub struct MultiHeadedAttentionSANM { - head_dim: usize, - n_head: usize, - linear_out: Linear, - linear_q_k_v: Linear, - fsmn_block: Conv1d, - left_padding: usize, - right_padding: usize, - scaling: f64, -} - -impl MultiHeadedAttentionSANM { - pub fn new( - vb: VarBuilder, - n_head: usize, - in_dim: usize, - hidden_dim: usize, - kernel_size: usize, - sanm_shfit: usize, - ) -> Result { - let head_dim = hidden_dim / n_head; - let linear_out = linear(hidden_dim, hidden_dim, vb.pp("linear_out"))?; - let linear_q_k_v = linear(in_dim, hidden_dim * 3, vb.pp("linear_q_k_v"))?; - let fsmn_block = get_conv1d( - vb.pp("fsmn_block"), - hidden_dim, - hidden_dim, - kernel_size, - 0, - 1, - 1, - hidden_dim, - false, - )?; - let mut left_padding = (kernel_size - 1) / 2; - if sanm_shfit > 0 { - left_padding += sanm_shfit; - } - let right_padding = kernel_size - 1 - left_padding; - let scaling = (head_dim as f64).powf(-0.5); - Ok(Self { - head_dim, - n_head, - linear_out, - linear_q_k_v, - fsmn_block, - left_padding, - right_padding, - scaling, - }) - } - - pub fn forward_fsmn( - &self, - inputs: &Tensor, - mask: Option<&Tensor>, - mask_shfit_chunk: Option<&Tensor>, - ) -> Result { - let mut inputs = inputs.clone(); - let mask = if let Some(mask) = mask { - let mut mask = mask.unsqueeze(D::Minus1)?.unsqueeze(0)?; - if let Some(mask_shfit_chunk) = mask_shfit_chunk { - mask = mask.broadcast_mul(mask_shfit_chunk)?; - } - inputs = inputs.broadcast_mul(&mask)?; - Some(mask) - } else { - None - }; - let xs = inputs.transpose(1, 2)?; - let xs = xs.pad_with_zeros(D::Minus1, self.left_padding, self.right_padding)?; - let xs = self.fsmn_block.forward(&xs)?; - let xs = xs.transpose(1, 2)?; - let mut xs = xs.add(&inputs)?; - if let Some(mask) = mask { - xs = xs.broadcast_mul(&mask)?; - } - Ok(xs) - } - pub fn forward_qkv(&self, xs: &Tensor) -> Result<(Tensor, Tensor, Tensor, Tensor)> { - let (b, t, _) = xs.dims3()?; - let q_k_v = self - .linear_q_k_v - .forward(xs)? - .reshape((b, t, 3, self.n_head, ()))? - .permute((2, 0, 3, 1, 4))? - .contiguous()?; - let q_h = q_k_v.i(0)?.contiguous()?; - let k_h = q_k_v.i(1)?.contiguous()?; - let v_h = q_k_v.i(2)?.contiguous()?; - let v = v_h.transpose(1, 2)?.reshape((b, t, ()))?; - Ok((q_h, k_h, v_h, v)) - } - - pub fn forward_attention( - &self, - values: &Tensor, - scores: &Tensor, - mask: Option<&Tensor>, - mask_att_chunk_encoder: Option<&Tensor>, - ) -> Result { - let bs = scores.dim(0)?; - let attn = if let Some(mask) = mask { - let mask = if let Some(mask_att_chunk_encoder) = mask_att_chunk_encoder { - mask.mul(mask_att_chunk_encoder)? - } else { - mask.clone() - }; - // mask: rank = 2 - let mask = get_equal_mask(&mask, 0)?; - let scores = attn_masked_fill(scores, &mask, f32::NEG_INFINITY)?; - let attn = softmax_last_dim(&scores)?; - attn_masked_fill(&attn, &mask, 0.0)? - } else { - softmax_last_dim(scores)? - }; - let xs = attn.matmul(values)?; - let xs = - xs.transpose(1, 2)? - .contiguous()? - .reshape((bs, (), self.n_head * self.head_dim))?; - let xs = self.linear_out.forward(&xs)?; - Ok(xs) - } - - pub fn forward_simple(&self, xs: &Tensor) -> Result { - let (b, t, _) = xs.dims3()?; - let q_k_v = self.linear_q_k_v.forward(xs)?; - let dim = self.head_dim * self.n_head; - let q_h = q_k_v - .narrow(D::Minus1, 0, dim)? - .reshape((b, t, self.n_head, ()))? - .permute((0, 2, 1, 3))?; - let k_h = q_k_v - .narrow(D::Minus1, dim, dim)? - .reshape((b, t, self.n_head, ()))? - .permute((0, 2, 1, 3))?; - let v = q_k_v.narrow(D::Minus1, dim * 2, dim)?; - let v_h = v.reshape((b, t, self.n_head, ()))?.permute((0, 2, 1, 3))?; - let fsmn_memory = v.transpose(1, 2)?; - let fsmn_memory = fsmn_memory - .pad_with_zeros(D::Minus1, self.left_padding, self.right_padding)? - .contiguous()?; - let fsmn_memory = self.fsmn_block.forward(&fsmn_memory)?; - // let fsmn_memory = conv1d_group_parallel(&fsmn_memory, &self.fsmn_block)?; - - let fsmn_memory = fsmn_memory.transpose(1, 2)?; - let fsmn_memory = fsmn_memory.add(&v)?; - let att_outs = eager_attention_forward(&q_h, &k_h, &v_h, None, None, self.scaling)?; - let att_outs = att_outs.reshape((b, t, ()))?; - let att_outs = self.linear_out.forward(&att_outs)?; - let att_outs = att_outs.add(&fsmn_memory)?; - Ok(att_outs) - } - - pub fn forward( - &self, - xs: &Tensor, - mask: Option<&Tensor>, - mask_shfit_chunk: Option<&Tensor>, - mask_att_chunk_encoder: Option<&Tensor>, - ) -> Result { - let (q_h, k_h, v_h, v) = self.forward_qkv(xs)?; - let fsmn_memory = self.forward_fsmn(&v, mask, mask_shfit_chunk)?; - let q_h = q_h.affine(self.scaling, 0.0)?; - let scores = q_h.matmul(&k_h.transpose(D::Minus2, D::Minus1)?)?; - let attn_outs = self.forward_attention(&v_h, &scores, mask, mask_att_chunk_encoder)?; - let att_outs = attn_outs.add(&fsmn_memory)?; - Ok(att_outs) - } -} - -pub struct EncoderLayerSANM { - self_attn: MultiHeadedAttentionSANM, - feed_forward: TwoLinearMLP, - norm1: LayerNorm, - norm2: LayerNorm, - concat_linear: Option, - normalize_before: bool, - in_dim: usize, - hidden_dim: usize, -} - -impl EncoderLayerSANM { - pub fn new( - vb: VarBuilder, - in_dim: usize, - hidden_dim: usize, - n_head: usize, - kernel_size: usize, - sanm_shfit: usize, - hidden_units: usize, - normalize_before: bool, - concat_after: bool, - ) -> Result { - let self_attn = MultiHeadedAttentionSANM::new( - vb.pp("self_attn"), - n_head, - in_dim, - hidden_dim, - kernel_size, - sanm_shfit, - )?; - let feed_forward = TwoLinearMLP::new( - vb.pp("feed_forward"), - hidden_dim, - hidden_units, - hidden_dim, - candle_nn::Activation::Relu, - true, - "w_1", - "w_2", - )?; - let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, in_dim)?; - let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, hidden_dim)?; - let concat_linear = if concat_after { - let lin = linear(hidden_dim * 2, hidden_dim, vb.pp("concat_linear"))?; - Some(lin) - } else { - None - }; - Ok(Self { - self_attn, - feed_forward, - norm1, - norm2, - concat_linear, - normalize_before, - in_dim, - hidden_dim, - }) - } - - pub fn forward( - &self, - xs: &Tensor, - mask: Option<&Tensor>, - mask_shfit_chunk: Option<&Tensor>, - mask_att_chunk_encoder: Option<&Tensor>, - ) -> Result { - let stoch_layer_coeff = 1.0f64; - let residual = xs.clone(); - let mut xs = if self.normalize_before { - self.norm1.forward(xs)? - } else { - xs.clone() - }; - if self.concat_linear.is_some() { - let attn = - self.self_attn - .forward(&xs, mask, mask_shfit_chunk, mask_att_chunk_encoder)?; - let x_concat = Tensor::cat(&[&xs, &attn], D::Minus1)?; - if self.in_dim == self.hidden_dim { - let x_concat = self - .concat_linear - .as_ref() - .unwrap() - .forward(&x_concat)? - .affine(stoch_layer_coeff, 0.0)?; - xs = residual.add(&x_concat)?; - } else { - xs = self - .concat_linear - .as_ref() - .unwrap() - .forward(&x_concat)? - .affine(stoch_layer_coeff, 0.0)?; - } - } else if self.in_dim == self.hidden_dim { - let attn = self - .self_attn - .forward(&xs, mask, mask_shfit_chunk, mask_att_chunk_encoder)? - .affine(stoch_layer_coeff, 0.0)?; - xs = residual.add(&attn)?; - } else { - xs = self - .self_attn - .forward(&xs, mask, mask_shfit_chunk, mask_att_chunk_encoder)? - .affine(stoch_layer_coeff, 0.0)?; - } - - if !self.normalize_before { - xs = self.norm1.forward(&xs)?; - } - let residual = xs.clone(); - if self.normalize_before { - xs = self.norm2.forward(&xs)?; - } - xs = self - .feed_forward - .forward(&xs)? - .affine(stoch_layer_coeff, 0.0)?; - xs = residual.add(&xs)?; - if !self.normalize_before { - xs = self.norm2.forward(&xs)?; - } - Ok(xs) - } - - pub fn forward_simple(&self, xs: &Tensor) -> Result { - let residual = xs.clone(); - let mut xs = self.norm1.forward(xs)?; - if self.in_dim == self.hidden_dim { - let attn = self.self_attn.forward_simple(&xs)?; - xs = residual.add(&attn)?; - } else { - xs = self.self_attn.forward_simple(&xs)?; - } - - let residual = xs.clone(); - let xs = self.norm2.forward(&xs)?; - - let xs = self.feed_forward.forward(&xs)?; - let xs = residual.add(&xs)?; - Ok(xs) - } -} - -pub struct SenseVoiceEncoderSmall { - embed: SinusoidalPositionEncoderCat, - encoders0: EncoderLayerSANM, - encoders: Vec, - tp_encoders: Vec, - after_norm: LayerNorm, - tp_norm: LayerNorm, - scaling: f64, -} - -impl SenseVoiceEncoderSmall { - pub fn new( - vb: VarBuilder, - input_size: usize, - output_size: usize, - attention_heads: usize, - linear_units: usize, - num_blocks: usize, - tp_blocks: usize, - normalize_before: bool, - kernel_size: usize, - sanm_shfit: usize, - ) -> Result { - let embed = SinusoidalPositionEncoderCat::new(Some(input_size), true, vb.device())?; - - let encoders0 = EncoderLayerSANM::new( - vb.pp("encoders0.0"), - input_size, - output_size, - attention_heads, - kernel_size, - sanm_shfit, - linear_units, - normalize_before, - false, - )?; - let mut encoders = vec![]; - let vb_encoders = vb.pp("encoders"); - for i in 0..(num_blocks - 1) { - let encoder_i = EncoderLayerSANM::new( - vb_encoders.pp(i), - output_size, - output_size, - attention_heads, - kernel_size, - sanm_shfit, - linear_units, - normalize_before, - false, - )?; - encoders.push(encoder_i); - } - let vb_tp_encoders = vb.pp("tp_encoders"); - let mut tp_encoders = vec![]; - for i in 0..tp_blocks { - let tp_blocks_i = EncoderLayerSANM::new( - vb_tp_encoders.pp(i), - output_size, - output_size, - attention_heads, - kernel_size, - sanm_shfit, - linear_units, - normalize_before, - false, - )?; - tp_encoders.push(tp_blocks_i); - } - let after_norm = get_layer_norm(vb.pp("after_norm"), 1e-5, output_size)?; - let tp_norm = get_layer_norm(vb.pp("tp_norm"), 1e-5, output_size)?; - let scaling = (output_size as f64).powf(0.5); - Ok(Self { - embed, - encoders0, - encoders, - tp_encoders, - after_norm, - tp_norm, - scaling, - }) - } - pub fn forward(&self, xs: &Tensor) -> Result { - let xs = xs.affine(self.scaling, 0.0)?; - let xs = self.embed.forward(&xs, 0)?; - let mut xs = self.encoders0.forward_simple(&xs)?; - for encoder_layer in &self.encoders { - xs = encoder_layer.forward_simple(&xs)?; - } - xs = self.after_norm.forward(&xs)?; - for tp_layer in &self.tp_encoders { - xs = tp_layer.forward_simple(&xs)?; - } - xs = self.tp_norm.forward(&xs)?; - Ok(xs) - } -} - -pub struct AdaptorEncoderLayer { +pub struct Qwen3ASRAudioEncoderLayer { self_attn: NaiveAttention, - feed_forward: TwoLinearMLP, - norm1: LayerNorm, - norm2: LayerNorm, - concat_linear: Option, - normalize_before: bool, + self_attn_layer_norm: LayerNorm, + activation_fn: Activation, + fc1: Linear, + fc2: Linear, + final_layer_norm: LayerNorm, } -impl AdaptorEncoderLayer { - pub fn new( - vb: VarBuilder, - llm_dim: usize, - n_head: usize, - normalize_before: bool, - concat_after: bool, - ) -> Result { +impl Qwen3ASRAudioEncoderLayer { + pub fn new(vb: VarBuilder, config: &Qwen3ASRAudioConfig) -> Result { let self_attn = NaiveAttention::new( vb.pp("self_attn"), - llm_dim, - n_head, - n_head, + config.d_model, + config.encoder_attention_heads, + config.encoder_attention_heads, None, true, - Some("linear_q"), - Some("linear_k"), - Some("linear_v"), - Some("linear_out"), + Some("q_proj"), + Some("k_proj"), + Some("v_proj"), + Some("out_proj"), )?; - let feed_forward = TwoLinearMLP::new( - vb.pp("feed_forward"), - llm_dim, - llm_dim / 4, - llm_dim, - candle_nn::Activation::Relu, - true, - "w_1", - "w_2", - )?; - let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, llm_dim)?; - let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, llm_dim)?; - let concat_linear = if concat_after { - let lin = linear(llm_dim * 2, llm_dim, vb.pp("concat_linear"))?; - Some(lin) - } else { - None - }; + let self_attn_layer_norm = + get_layer_norm(vb.pp("self_attn_layer_norm"), 1e-5, config.d_model)?; + let activation_fn = config.activation_function; + let fc1 = linear(config.d_model, config.encoder_ffn_dim, vb.pp("fc1"))?; + let fc2 = linear(config.encoder_ffn_dim, config.d_model, vb.pp("fc2"))?; + let final_layer_norm = get_layer_norm(vb.pp("final_layer_norm"), 1e-5, config.d_model)?; Ok(Self { self_attn, - feed_forward, - norm1, - norm2, - concat_linear, - normalize_before, + self_attn_layer_norm, + activation_fn, + fc1, + fc2, + final_layer_norm, }) } pub fn forward(&self, xs: &Tensor, mask: Option<&Tensor>) -> Result { - let stoch_layer_coeff = 1.0f64; let residual = xs.clone(); - let mut xs = if self.normalize_before { - self.norm1.forward(xs)? - } else { - xs.clone() - }; - if self.concat_linear.is_some() { - let attn = self.self_attn.forward(&xs, None, None, mask, false)?; - let x_concat = Tensor::cat(&[&xs, &attn], D::Minus1)?; - let x_concat = self - .concat_linear - .as_ref() - .unwrap() - .forward(&x_concat)? - .affine(stoch_layer_coeff, 0.0)?; - xs = residual.add(&x_concat)?; - } else { - let attn = self - .self_attn - .forward(&xs, None, None, mask, false)? - .affine(stoch_layer_coeff, 0.0)?; - xs = residual.add(&attn)?; - } - if !self.normalize_before { - xs = self.norm1.forward(&xs)?; - } - let residual = xs.clone(); - if self.normalize_before { - xs = self.norm2.forward(&xs)?; - } - xs = self - .feed_forward - .forward(&xs)? - .affine(stoch_layer_coeff, 0.0)?; - xs = residual.add(&xs)?; - if !self.normalize_before { - xs = self.norm2.forward(&xs)?; - } + let xs = self.self_attn_layer_norm.forward(xs)?; + let xs = self.self_attn.forward(&xs, None, None, mask, false)?; + let residual = xs.add(&residual)?; + let xs = self.final_layer_norm.forward(&residual)?; + let xs = self.fc1.forward(&xs)?.apply(&self.activation_fn)?; + let xs = self.fc2.forward(&xs)?; + let xs = xs.add(&residual)?; Ok(xs) } } -pub struct AudioAdaptor { - k: usize, - linear1: Linear, - linear2: Linear, - blocks: Vec, +pub struct Qwen3ASRAudioEncoder { + n_window: usize, + positional_embedding: SinusoidalPositionEncoderCat, + layers: Vec, + ln_post: LayerNorm, + conv2d1: Conv2d, + conv2d2: Conv2d, + conv2d3: Conv2d, + conv_out: Linear, + proj1: Linear, + act: Activation, + proj2: Linear, + // n_window_infer: usize, + conv_chunksize: usize, } -impl AudioAdaptor { - pub fn new( - vb: VarBuilder, - downsample_rate: usize, - encoder_dim: usize, - llm_dim: usize, - ffn_dim: usize, - n_layer: usize, - attention_heads: usize, - ) -> Result { - let linear1 = linear(encoder_dim * downsample_rate, ffn_dim, vb.pp("linear1"))?; - let linear2 = linear(ffn_dim, llm_dim, vb.pp("linear2"))?; - let mut blocks = vec![]; - let vb_blocks = vb.pp("blocks"); - for i in 0..n_layer { - let layer = - AdaptorEncoderLayer::new(vb_blocks.pp(i), llm_dim, attention_heads, true, false)?; - blocks.push(layer); +impl Qwen3ASRAudioEncoder { + pub fn new(vb: VarBuilder, config: &Qwen3ASRAudioConfig) -> Result { + let n_window = config.n_window; + let positional_embedding = + SinusoidalPositionEncoderCat::new(Some(config.d_model), true, vb.device())?; + let mut layers = vec![]; + let vb_layers = vb.pp("layers"); + for i in 0..config.encoder_layers { + let layer = Qwen3ASRAudioEncoderLayer::new(vb_layers.pp(i), config)?; + layers.push(layer); } + let ln_post = get_layer_norm(vb.pp("ln_post"), 1e-5, config.d_model)?; + let conv2d1 = get_conv2d( + vb.pp("conv2d1"), + 1, + config.downsample_hidden_size, + 3, + 1, + 2, + 1, + 1, + true, + )?; + let conv2d2 = get_conv2d( + vb.pp("conv2d2"), + config.downsample_hidden_size, + config.downsample_hidden_size, + 3, + 1, + 2, + 1, + 1, + true, + )?; + let conv2d3 = get_conv2d( + vb.pp("conv2d3"), + config.downsample_hidden_size, + config.downsample_hidden_size, + 3, + 1, + 2, + 1, + 1, + true, + )?; + let in_dim = + config.downsample_hidden_size * ((((config.num_mel_bins + 1) / 2 + 1) / 2 + 1) / 2); + let conv_out = linear_no_bias(in_dim, config.d_model, vb.pp("conv_out"))?; + let proj1 = linear(config.d_model, config.d_model, vb.pp("proj1"))?; + let act = config.activation_function; + let proj2 = linear(config.d_model, config.output_dim, vb.pp("proj2"))?; + // let n_window_infer = config.n_window_infer; + let conv_chunksize = config.conv_chunksize; Ok(Self { - k: downsample_rate, - linear1, - linear2, - blocks, + n_window, + positional_embedding, + layers, + ln_post, + conv2d1, + conv2d2, + conv2d3, + conv_out, + proj1, + act, + proj2, + // n_window_infer, + conv_chunksize, }) } + pub fn forward(&self, xs: &Tensor) -> Result { - let (bs, seq_len, dim) = xs.dims3()?; - let chunk_num = (seq_len - 1) / self.k + 1; - let pad_num = chunk_num * self.k - seq_len; - let xs = xs.pad_with_zeros(1, 0, pad_num)?; - let xs = xs.contiguous()?.reshape((bs, chunk_num, dim * self.k))?; - let xs = self.linear1.forward(&xs)?.relu()?; - let mut xs = self.linear2.forward(&xs)?; - for block in &self.blocks { - xs = block.forward(&xs, None)?; + // xs: (feature_dim, feature_len) + let feature_lens = xs.dim(1)?; + // let aftercnn_lens = get_feat_extract_output_lengths(feature_lens); + let chunk_num = feature_lens / (self.n_window * 2); + let mut chunk_lengths = vec![self.n_window * 2; chunk_num]; + let last = feature_lens % (self.n_window * 2); + if last > 0 { + chunk_lengths.push(last); } - Ok(xs) + let mut chunk_list = split_tensor(&xs.t()?, &chunk_lengths, 0)?; + if last > 0 { + let chunk_last = chunk_list + .pop() + .ok_or(anyhow::anyhow!(format!("chunk_list is empty")))?; + let pad_size = self.n_window * 2 - last; + let chunk_last = chunk_last.pad_with_zeros(0, 0, pad_size)?; + chunk_list.push(chunk_last); + } + let padded_feature = Tensor::stack(&chunk_list, 0)?.transpose(1, 2)?; + let feature_lens_after_cnn: Vec = chunk_lengths + .iter() + .map(|&i| get_feat_extract_output_lengths(i)) + .collect(); + let feature_len_after_cnn = feature_lens_after_cnn.iter().sum(); + let padded_feature = padded_feature.unsqueeze(1)?; + let mut padded_embeds = vec![]; + let feature_splits = split_tensor_with_size(&padded_feature, self.conv_chunksize, 0)?; + for chunk in feature_splits.iter() { + let padded_embed = self.conv2d1.forward(chunk)?.gelu()?; + let padded_embed = self.conv2d2.forward(&padded_embed)?.gelu()?; + let padded_embed = self.conv2d3.forward(&padded_embed)?.gelu()?; + padded_embeds.push(padded_embed); + } + let padded_embed = Tensor::cat(&padded_embeds, 0)?; + let (b, c, f, t) = padded_embed.dims4()?; + let padded_embed = + padded_embed + .permute((0, 3, 1, 2))? + .contiguous()? + .reshape((b, t, c * f))?; + let padded_embed = self.conv_out.forward(&padded_embed)?; + let padded_embed = self.positional_embedding.forward(&padded_embed, 0)?; + let padded_embed = padded_embed.flatten(0, 1)?; + let mut hidden_states = padded_embed + .narrow(0, 0, feature_len_after_cnn)? + .unsqueeze(0)?; + for layer in &self.layers { + hidden_states = layer.forward(&hidden_states, None)?; + } + let hidden_states = hidden_states.squeeze(0)?; + let hidden_states = self.ln_post.forward(&hidden_states)?; + let hidden_states = self.proj1.forward(&hidden_states)?.apply(&self.act)?; + let hidden_states = self.proj2.forward(&hidden_states)?; + Ok(hidden_states) } } -pub struct FunAsrNanoModel { - audio_encoder: SenseVoiceEncoderSmall, - audio_adaptor: AudioAdaptor, - llm: Qwen3Model, +pub struct Qwen3ASRThinkerTextModel { + embed_tokens: Embedding, + layers: Vec, + norm: RmsNorm, + rotary_emb: Qwen3VLTextRotaryEmbedding, + mrope_section: Vec, } -impl FunAsrNanoModel { - pub fn new(vb: VarBuilder, config: &FunASRNanoConfig, llm_cfg: &Qwen3Config) -> Result { - let input_size = config.frontend_conf.lfr_m * config.frontend_conf.n_mels; - let audio_encoder = SenseVoiceEncoderSmall::new( - vb.pp("audio_encoder"), - input_size, - config.audio_encoder_conf.output_size, - config.audio_encoder_conf.attention_heads, - config.audio_encoder_conf.linear_units, - config.audio_encoder_conf.num_blocks, - config.audio_encoder_conf.tp_blocks, - config.audio_encoder_conf.normalize_before, - config.audio_encoder_conf.kernel_size, - config.audio_encoder_conf.sanm_shfit, - )?; - let audio_adaptor = AudioAdaptor::new( - vb.pp("audio_adaptor"), - config.audio_adaptor_conf.downsample_rate, - config.audio_adaptor_conf.encoder_dim, - config.audio_adaptor_conf.llm_dim, - config.audio_adaptor_conf.ffn_dim, - config.audio_adaptor_conf.n_layer, - 8, - )?; - let llm = Qwen3Model::new(llm_cfg, vb.pp("llm"))?; + +impl Qwen3ASRThinkerTextModel { + pub fn new(vb: VarBuilder, cfg: &Qwen3ASRTextConfig) -> Result { + let embed_tokens = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?; + let mut layers = vec![]; + let vb_layers = vb.pp("layers"); + let qwen3cfg = qwen3asr_text_config2qwen3_config(cfg); + for i in 0..cfg.num_hidden_layers { + let layer = Qwen3DecoderLayer::new(&qwen3cfg, vb_layers.pp(i))?; + layers.push(layer); + } + let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm"))?; + let rotary_emb = Qwen3VLTextRotaryEmbedding::new(cfg.head_dim, cfg.rope_theta); Ok(Self { - audio_encoder, - audio_adaptor, - llm, + embed_tokens, + layers, + norm, + rotary_emb, + mrope_section: cfg.rope_scaling.mrope_section.clone(), + }) + } + + pub fn forward( + &mut self, + input_embeds: &Tensor, + seqlen_offset: usize, + position_ids: Option<&Tensor>, + ) -> Result { + let (b_size, seq_len, _) = input_embeds.dims3()?; + let position_ids = match position_ids { + Some(ids) => ids.clone(), + None => Tensor::arange( + seqlen_offset as u32, + (seq_len + seqlen_offset) as u32, + input_embeds.device(), + )? + .unsqueeze(0)? + .unsqueeze(0)? + .broadcast_as((3, b_size, seq_len))?, + }; + let (cos, sin) = self.rotary_emb.forward_asr( + &position_ids, + input_embeds.dtype(), + self.mrope_section.clone(), + )?; + let mut xs = input_embeds.clone(); + let attention_mask: Option = { + if seq_len <= 1 { + None + } else { + Some(prepare_causal_attention_mask( + b_size, + seq_len, + 0, + input_embeds.device(), + )?) + } + }; + for layer in self.layers.iter_mut() { + xs = layer.forward(&xs, &cos, &sin, attention_mask.as_ref())?; + } + let xs = self.norm.forward(&xs)?; + Ok(xs) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} + +pub struct Qwen3ASRThinker { + audio_tower: Qwen3ASRAudioEncoder, + model: Qwen3ASRThinkerTextModel, + audio_token_id: u32, + lm_head: Linear, +} + +impl Qwen3ASRThinker { + pub fn new(vb: VarBuilder, config: &ThinkerConfig) -> Result { + let audio_tower = Qwen3ASRAudioEncoder::new(vb.pp("audio_tower"), &config.audio_config)?; + let model = Qwen3ASRThinkerTextModel::new(vb.pp("model"), &config.text_config)?; + let lm_head = if config.text_config.tie_word_embeddings { + Linear::new(model.embed_tokens.embeddings().clone(), None) + } else { + linear_no_bias( + config.text_config.hidden_size, + config.text_config.vocab_size, + vb.pp("lm_head"), + )? + }; + Ok(Self { + audio_tower, + model, + audio_token_id: config.audio_token_id, + lm_head, }) } pub fn forward( &mut self, input_ids: &Tensor, - speech: Option<&Tensor>, - fbank_mask: Option<&Tensor>, seqlen_offset: usize, + input_features: Option<&Tensor>, ) -> Result { - let mut inputs_embeds = self.llm.embedding_token_id(input_ids)?; - if let Some(speech) = speech - && let Some(fbank_mask) = fbank_mask - { - let speech = self.audio_encoder.forward(speech)?; - let encoder_out = self.audio_adaptor.forward(&speech)?; - let speech_token_len = fbank_mask.sum_all()?.to_scalar::()?; - let audio_embed = encoder_out - .squeeze(0)? - .narrow(0, 0, speech_token_len as usize)?; - inputs_embeds = masked_scatter_dim0(&inputs_embeds, &audio_embed, fbank_mask)?; + let mut input_embeds = self.model.embed_tokens.forward(input_ids)?; + if let Some(input_features) = input_features { + let audio_feature = self.audio_tower.forward(&input_features)?; + // println!("audio_feature: {}", audio_feature); + let audio_mask = get_equal_mask(input_ids, self.audio_token_id)?; + let n_audio_tokens = audio_mask.sum_all()?.to_scalar::()?; + if n_audio_tokens as usize != audio_feature.dim(0)? { + return Err(anyhow::anyhow!(format!( + "n_audio_tokens num: {} not equal to audio_feature len: {}", + n_audio_tokens, + audio_feature.dim(0)? + ))); + } + input_embeds = masked_scatter_dim0(&input_embeds, &audio_feature, &audio_mask)?; } - let logits = self - .llm - .forward(None, Some(&inputs_embeds), seqlen_offset)?; + let outputs = self.model.forward(&input_embeds, seqlen_offset, None)?; + let seq_len = outputs.dim(1)?; + let hidden_state = outputs.narrow(1, seq_len - 1, 1)?; + let logits = self.lm_head.forward(&hidden_state)?; Ok(logits) } - pub fn clear_kv_cache(&mut self) { - self.llm.clear_kv_cache(); + self.model.clear_kv_cache(); + } +} + +pub struct Qwen3ASRModel { + thinker: Qwen3ASRThinker, +} + +impl Qwen3ASRModel { + pub fn new(vb: VarBuilder, config: &Qwen3ASRConfig) -> Result { + let thinker = Qwen3ASRThinker::new(vb.pp("thinker"), &config.thinker_config)?; + Ok(Self { thinker }) + } + + pub fn forward( + &mut self, + input_ids: &Tensor, + seqlen_offset: usize, + input_features: Option<&Tensor>, + ) -> Result { + let logits = self + .thinker + .forward(input_ids, seqlen_offset, input_features)?; + Ok(logits) + } + pub fn clear_kv_cache(&mut self) { + self.thinker.clear_kv_cache(); } } diff --git a/src/models/qwen3_asr/processor.rs b/src/models/qwen3_asr/processor.rs index 508ac6b..e50e9b9 100644 --- a/src/models/qwen3_asr/processor.rs +++ b/src/models/qwen3_asr/processor.rs @@ -1,12 +1,17 @@ use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; use anyhow::Result; use candle_core::{Device, Tensor}; -use serde_json::{Value, json}; -use crate::utils::{ - audio_utils::{extract_audios, split_audio_into_chunks}, - capitalize_first_letter, extract_user_text_vec, - tensor_utils::float_range_normalize, +use crate::{ + models::feature_extractor::{ + config::FeatureExtractor, feature_extraction_whisper::WhisperFeatureExtractor, + }, + tokenizer::TokenizerModel, + utils::{ + audio_utils::{extract_audios, split_audio_into_chunks}, + capitalize_first_letter, + tensor_utils::float_range_normalize, + }, }; pub struct Qwen3AsrProcessor { @@ -14,10 +19,12 @@ pub struct Qwen3AsrProcessor { sample_rate: usize, support_language: Vec, max_asr_input_seconds: f32, + whisper_feature_extracor: WhisperFeatureExtractor, + audio_token: String, } impl Qwen3AsrProcessor { - pub fn new(device: &Device) -> Result { + pub fn new(device: &Device, config: &FeatureExtractor) -> Result { let support_language: Vec = vec![ "Chinese", "English", @@ -53,11 +60,23 @@ impl Qwen3AsrProcessor { .iter() .map(|s| s.to_string()) .collect(); + let whisper_feature_extracor = WhisperFeatureExtractor::new( + config.feature_size, + config.hop_length, + // config.chunk_length, + config.n_fft, + config.dither, + // config.padding_value, + config.sampling_rate, + device, + )?; Ok(Self { device: device.clone(), sample_rate: 16000, support_language, max_asr_input_seconds: 1200.0, + whisper_feature_extracor, + audio_token: "<|audio_pad|>".to_string(), }) } @@ -73,7 +92,19 @@ impl Qwen3AsrProcessor { self.support_language.contains(lang) } - pub fn process_info(&self, mes: &ChatCompletionParameters, render: &str) -> Result<()> { + fn replace_special_tokens(&self, text: &str, token_len: usize) -> String { + let replace = "<|audio_placeholder|>".repeat(token_len as usize); + let text = text.replacen(&self.audio_token, &replace, 1); + let text = text.replace("<|audio_placeholder|>", &self.audio_token); + text + } + + pub fn process_info( + &self, + mes: &ChatCompletionParameters, + render: &str, + tokenizer: &TokenizerModel, + ) -> Result> { let audio_count = render .matches("<|audio_start|><|audio_pad|><|audio_end|>") .count(); @@ -104,24 +135,38 @@ impl Qwen3AsrProcessor { let wavs = split_audio_into_chunks(wav, self.sample_rate, self.max_asr_input_seconds)?; split_wavs.extend_from_slice(&wavs); } - - - // let mut audio_datas = vec![]; - // for (i, wav) in audio_tensors.iter().enumerate() { - // let wavs = split_audio_into_chunks(wav, self.sample_rate, self.max_asr_input_seconds)?; - // for i_w in wavs { - // let audio_data = AudioData { - // wav: i_w, - // language: langs[i].clone(), - // }; - // audio_datas.push(audio_data); - // } - // } - Ok(()) + let mut audio_datas = vec![]; + for wav in split_wavs.iter() { + let (input_features, _) = + self.whisper_feature_extracor + .call(wav, self.sample_rate, false)?; + let audio_len = input_features.dim(2)?; + let output_len = get_feat_extract_output_lengths(audio_len); + let text = self.replace_special_tokens(&render, output_len); + let input_ids = tokenizer.text_encode(text, &self.device)?; + let input_features = input_features.squeeze(0)?; + let audio = AudioData { + input_features, + input_ids, + }; + audio_datas.push(audio); + } + Ok(audio_datas) } } pub struct AudioData { - pub wav: Tensor, - pub language: Option, + pub input_features: Tensor, + pub input_ids: Tensor, +} + +pub fn get_feat_extract_output_lengths(audio_len: usize) -> usize { + let input_len_leave = audio_len % 100; + let output_len = if input_len_leave > 0 { + let feat_lengths = (input_len_leave - 1) / 2 + 1; + ((feat_lengths - 1) / 2 + 1 - 1) / 2 + 1 + (audio_len / 100) * 13 + } else { + (audio_len / 100) * 13 + }; + output_len } diff --git a/src/models/w2v_bert_2_0/model.rs b/src/models/w2v_bert_2_0/model.rs index 6b5b6e5..fd4ddfd 100644 --- a/src/models/w2v_bert_2_0/model.rs +++ b/src/models/w2v_bert_2_0/model.rs @@ -8,7 +8,7 @@ use candle_nn::{ use crate::{ models::{ common::{ - GLU, NaiveAttention, TwoLinearMLP, eager_attention_forward, get_conv1d, get_layer_norm, + GLU, TwoLinearMLP, eager_attention_forward, get_conv1d, get_layer_norm, }, w2v_bert_2_0::config::W2VBert2_0Config, }, diff --git a/src/position_embed/rope.rs b/src/position_embed/rope.rs index 7979c6c..82ed91d 100644 --- a/src/position_embed/rope.rs +++ b/src/position_embed/rope.rs @@ -340,6 +340,70 @@ impl Qwen3VLTextRotaryEmbedding { } Ok(freqs_t) } + + pub fn apply_interleaved_mrope_asr( + &self, + freqs: &Tensor, + mrope_section: Vec, + ) -> Result { + let mut freqs_t = freqs.i(0)?.contiguous()?; //(3, bs, seq_len, head_dim //2) -> (bs, seq_len, head_dim //2) + + // for dim in 1..3 { + for (dim, offset) in (1..3).enumerate() { + let dim = dim +1; + let length = mrope_section[dim]; + let idx = Tensor::arange_step(offset as u32, length as u32, 3, freqs.device())?; + let src = freqs.i(dim)?.contiguous()?; // (bs, seq_len, head_dim //2) + let src = src.index_select(&idx, D::Minus1)?.contiguous()?; + let idx = idx + .unsqueeze(0)? + .unsqueeze(0)? + .broadcast_as(src.shape())? + .contiguous()?; + freqs_t = freqs_t.scatter(&idx, &src, D::Minus1)?; + } + Ok(freqs_t) + } + + pub fn forward_asr( + &self, + position_ids: &Tensor, + dtype: DType, + mrope_section: Vec, + ) -> Result<(Tensor, Tensor)> { + // position_ids shape: (3, bs, position) -> (3, bs, 1, position) + let position_ids = if position_ids.rank() == 2 { + let (bs, len) = position_ids.dims2()?; + position_ids.unsqueeze(0)?.expand((3, bs, len))? + } else { + position_ids.clone() + }; + let position_ids_expanded = position_ids + .unsqueeze(D::Minus2)? + .to_dtype(DType::F32)? + .contiguous()?; + // inv_freq Vec -> Tensor(1, 1, head_dim / 2, 1) -> (3, bs, head_dim / 2, 1) + let inv_freq_expanded = Tensor::from_vec( + self.inv_freq.clone(), + (1, 1, self.inv_freq.len(), 1), + position_ids.device(), + )? + .broadcast_as((3, position_ids.dim(1)?, self.inv_freq.len(), 1))? + .to_dtype(DType::F32)? + .contiguous()?; + + // (3, bs, head_dim / 2, 1) matmul (3, bs, 1, position) + // -> (3, bs, head_dim / 2, seq_len) -> (3, bs, seq_len, head_dim / 2) + let freqs = inv_freq_expanded + .matmul(&position_ids_expanded)? + .transpose(2, 3)?; + let freqs = self.apply_interleaved_mrope_asr(&freqs, mrope_section)?; + let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?.contiguous()?; + let cos = emb.cos()?; + let sin = emb.sin()?; + Ok((cos.to_dtype(dtype)?, sin.to_dtype(dtype)?)) + } + pub fn forward( &self, position_ids: &Tensor, diff --git a/src/tokenizer/mod.rs b/src/tokenizer/mod.rs index 4687455..35f08bc 100644 --- a/src/tokenizer/mod.rs +++ b/src/tokenizer/mod.rs @@ -1,6 +1,10 @@ use anyhow::{Ok, Result, anyhow}; use candle_core::{Device, Tensor}; -use tokenizers::Tokenizer; +use serde_json::Value; +use tokenizers::{ + AddedToken, Tokenizer, decoders::byte_level::ByteLevel as ByteLevelDecoder, models::bpe::BPE, + pre_tokenizers::byte_level::ByteLevel, +}; pub struct TokenizerModel { pub tokenizer: Tokenizer, @@ -14,12 +18,70 @@ impl TokenizerModel { "model path file not exists" ); let tokenizer_file = path.clone() + "/tokenizer.json"; - assert!( - std::path::Path::new(&tokenizer_file).exists(), - "tokenizer.json not exists in model path" - ); - let tokenizer = Tokenizer::from_file(tokenizer_file) - .map_err(|e| anyhow!(format!("tokenizer from file error{}", e)))?; + let tokenizer = if std::path::Path::new(&tokenizer_file).exists() { + Tokenizer::from_file(tokenizer_file) + .map_err(|e| anyhow!(format!("tokenizer from file error{}", e)))? + } else { + // 如果不存在 tokenizer.json，尝试使用 vocab.json 和 merges.txt + let vocab_file = path.clone() + "/vocab.json"; + let merges_file = path.clone() + "/merges.txt"; + let config_file = path.clone() + "/tokenizer_config.json"; + + if !std::path::Path::new(&vocab_file).exists() { + return Err(anyhow!( + "Neither tokenizer.json nor vocab.json found in model path" + )); + } + + if !std::path::Path::new(&merges_file).exists() { + return Err(anyhow!( + "Neither tokenizer.json nor merges.txt found in model path" + )); + } + // 创建 BPE 模型 + let bpe = BPE::from_file(&vocab_file, &merges_file) + .build() + .map_err(|e| anyhow!(format!("failed to build BPE tokenizer: {}", e)))?; + + // 创建分词器 + let mut tokenizer = Tokenizer::new(bpe); + // 添加字节级预分词器，这会处理换行符等特殊字符 + let byte_level_pre_tokenizer = ByteLevel::new(false, true, false); + tokenizer.with_pre_tokenizer(Some(byte_level_pre_tokenizer)); + tokenizer.with_decoder(Some(ByteLevelDecoder::default())); + if std::path::Path::new(&config_file).exists() { + let config_content = std::fs::read_to_string(&config_file)?; + let config: Value = serde_json::from_str(&config_content)?; + if let Some(added_tokens_decoder) = config.get("added_tokens_decoder") { + let mut special_tokens = Vec::new(); + + if let Value::Object(tokens_map) = added_tokens_decoder { + for (_, token_info) in tokens_map { + if let Value::Object(token_obj) = token_info { + if let Some(content_val) = token_obj.get("content") { + if let Some(content) = content_val.as_str() { + let special = token_obj + .get("special") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + let added_token = + AddedToken::from(content.to_string(), special); + special_tokens.push(added_token); + } + } + } + } + } + + // 添加所有特殊标记 + if !special_tokens.is_empty() { + tokenizer.add_special_tokens(&special_tokens); + } + } + } + tokenizer + }; Ok(Self { tokenizer }) } diff --git a/src/utils/mod.rs b/src/utils/mod.rs index fad6422..ab64cee 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -19,7 +19,7 @@ use anyhow::{Result, anyhow}; use byteorder::{LittleEndian, ReadBytesExt}; use candle_core::{ Context, DType, Device, Shape, Tensor, - pickle::{Object, PthTensors, Stack, TensorInfo, read_all_with_key}, + pickle::{Object, Stack, TensorInfo, read_all_with_key}, }; use candle_nn::VarBuilder; use candle_transformers::generation::{LogitsProcessor, Sampling}; diff --git a/src/utils/tensor_utils.rs b/src/utils/tensor_utils.rs index fbb0deb..025d82a 100644 --- a/src/utils/tensor_utils.rs +++ b/src/utils/tensor_utils.rs @@ -105,13 +105,14 @@ pub fn split_tensor_with_size( let dim = dim.to_index(t.shape(), "split")?; let mut split_res = Vec::new(); let dim_size = t.dim(dim)?; - assert_eq!( - dim_size % splits_size, - 0, - "input tensor dim size % splits_size must be equal to 0" - ); - for split in (0..dim_size).step_by(splits_size) { - split_res.push(t.narrow(dim, split, splits_size)?); + // assert_eq!( + // dim_size % splits_size, + // 0, + // "input tensor dim size % splits_size must be equal to 0" + // ); + for (i, split) in (0..dim_size).step_by(splits_size).enumerate() { + let size = splits_size.min(dim_size - i*splits_size); + split_res.push(t.narrow(dim, split, size)?); } Ok(split_res) } diff --git a/tests/messy_test.rs b/tests/messy_test.rs index b653509..d7cfe61 100644 --- a/tests/messy_test.rs +++ b/tests/messy_test.rs @@ -2,9 +2,9 @@ use std::time::Instant; -use aha::utils::{audio_utils::create_hann_window, tensor_utils::interpolate_nearest_1d}; +use aha::utils::{tensor_utils::interpolate_nearest_1d}; use anyhow::Result; -use candle_core::{DType, Tensor}; +use candle_core::{Tensor}; // use symphonia::core::io::MediaSourceStream; #[test] diff --git a/tests/test_qwen3_asr.rs b/tests/test_qwen3_asr.rs index acb7937..ad237c3 100644 --- a/tests/test_qwen3_asr.rs +++ b/tests/test_qwen3_asr.rs @@ -1,12 +1,57 @@ use std::{pin::pin, time::Instant}; -use aha::models::{GenerateModel, fun_asr_nano::generate::FunAsrNanoGenerateModel, qwen3_asr::generate::Qwen3AsrGenerateModel}; +use aha::models::{GenerateModel, qwen3_asr::generate::Qwen3AsrGenerateModel}; use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; use anyhow::Result; use rocket::futures::StreamExt; #[test] fn qwen3_asr_generate() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda qwen3_asr_generate -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-ASR-0.6B/", save_dir); //Qwen/Qwen3-ASR-1.7B + let message = r#" + { + "model": "qwen3-asr", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "audio", + "audio_url": + { + "url": "https://package-release.coderbox.cn/aiway/test/other/%E5%93%AA%E5%90%92.wav" + } + } + ] + } + ] + } + "#; + // "metadata": {"language": "Chinese"} + let mes: ChatCompletionParameters = serde_json::from_str(message)?; + let i_start = Instant::now(); + let mut model = Qwen3AsrGenerateModel::init(&model_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + let i_start = Instant::now(); + let res = model.generate(mes)?; + let i_duration = i_start.elapsed(); + println!("generate: \n {:?}", res); + if res.usage.is_some() { + let num_token = res.usage.as_ref().unwrap().total_tokens; + let duration_secs = i_duration.as_secs_f64(); + let tps = num_token as f64 / duration_secs; + println!("Tokens per second (TPS): {:.2}", tps); + } + println!("Time elapsed in generate is: {:?}", i_duration); + Ok(()) +} + +#[tokio::test] +async fn qwen3_asr_stream() -> Result<()> { + // RUST_BACKTRACE=1 cargo test -F cuda qwen3_asr_stream -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-ASR-0.6B/", save_dir); @@ -26,69 +71,20 @@ fn qwen3_asr_generate() -> Result<()> { } ] } - ], - "metadata": {"language": "Chinese"} + ] } "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = Qwen3AsrGenerateModel::init(&model_path, None, None)?; + let mut fun_asr_model = Qwen3AsrGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); let i_start = Instant::now(); - let res = model.generate(mes)?; + let mut stream = pin!(fun_asr_model.generate_stream(mes)?); + while let Some(item) = stream.next().await { + println!("generate: \n {:?}", item); + } let i_duration = i_start.elapsed(); - println!("generate: \n {:?}", res); - // if res.usage.is_some() { - // let num_token = res.usage.as_ref().unwrap().total_tokens; - // let duration_secs = i_duration.as_secs_f64(); - // let tps = num_token as f64 / duration_secs; - // println!("Tokens per second (TPS): {:.2}", tps); - // } println!("Time elapsed in generate is: {:?}", i_duration); Ok(()) } - -// #[tokio::test] -// async fn qwen3_asr_stream() -> Result<()> { -// // RUST_BACKTRACE=1 cargo test -F cuda fun_asr_nano_stream -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-ASR-0.6B/", save_dir); -// let message = r#" -// { -// "model": "qwen3-asr", -// "messages": [ -// { -// "role": "user", -// "content": [ -// { -// "type": "audio", -// "audio_url": -// { -// "url": "https://package-release.coderbox.cn/aiway/test/other/%E5%93%AA%E5%90%92.wav" -// } -// }, -// { -// "type": "text", -// "text": "语音转写：" -// } -// ] -// } -// ] -// } -// "#; -// let mes: ChatCompletionParameters = serde_json::from_str(message)?; -// let i_start = Instant::now(); -// let mut fun_asr_model = FunAsrNanoGenerateModel::init(&model_path, None, None)?; -// let i_duration = i_start.elapsed(); -// println!("Time elapsed in load model is: {:?}", i_duration); -// let i_start = Instant::now(); -// let mut stream = pin!(fun_asr_model.generate_stream(mes)?); -// while let Some(item) = stream.next().await { -// println!("generate: \n {:?}", item); -// } -// let i_duration = i_start.elapsed(); -// println!("Time elapsed in generate is: {:?}", i_duration); -// Ok(()) -// }