diff --git a/src/models/glm_asr_nano/config.rs b/src/models/glm_asr_nano/config.rs new file mode 100644 index 0000000..6dacf41 --- /dev/null +++ b/src/models/glm_asr_nano/config.rs @@ -0,0 +1,24 @@ +use serde::{Deserialize}; + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct GlmAsrNanoProcessorConfig { + pub audio_token: String, + pub default_transcription_prompt: String, + pub feature_extractor: FeatureExtractor, + 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, +} \ No newline at end of file diff --git a/src/models/glm_asr_nano/generate.rs b/src/models/glm_asr_nano/generate.rs new file mode 100644 index 0000000..dc73896 --- /dev/null +++ b/src/models/glm_asr_nano/generate.rs @@ -0,0 +1,50 @@ +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::Result; +use candle_core::{DType, Device}; + +use crate::{ + chat_template::ChatTemplate, + models::glm_asr_nano::{config::GlmAsrNanoProcessorConfig, processor::GlmAsrNanoProcessor}, + tokenizer::TokenizerModel, + utils::{get_device, get_dtype}, +}; + +pub struct GlmAsrNanoGenerateModel<'a> { + chat_template: ChatTemplate<'a>, + tokenizer: TokenizerModel, + processor: GlmAsrNanoProcessor, + // glm_asr_nano: GlmAsrNanoModel, + device: Device, + // eos_token_id1: u32, + // eos_token_id2: u32, + // eos_token_id3: u32, + // generation_config: GlmAsrNanoGenerationConfig, + model_name: String, +} + +impl<'a> GlmAsrNanoGenerateModel<'a> { + pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { + let chat_template = ChatTemplate::init(path)?; + let tokenizer = TokenizerModel::init(path)?; + let device = get_device(device); + let processor = GlmAsrNanoProcessor::new(path, &device, DType::F32)?; + // let cfg_dtype = cfg.dtype.as_str(); + // let dtype = get_dtype(dtype, cfg_dtype); + Ok(Self { + chat_template, + tokenizer, + processor, + device, + // eos_token_id1, + // eos_token_id2, + // eos_token_id3, + model_name: "glm-asr-nano".to_string(), + }) + } + + pub fn generate(&self, mes: ChatCompletionParameters) -> Result<()> { + let render_text = self.chat_template.apply_chat_template(&mes)?; + let audio = self.processor.process_info(&mes)?; + Ok(()) + } +} diff --git a/src/models/glm_asr_nano/mod.rs b/src/models/glm_asr_nano/mod.rs new file mode 100644 index 0000000..ce9b2a0 --- /dev/null +++ b/src/models/glm_asr_nano/mod.rs @@ -0,0 +1,4 @@ +pub mod config; +pub mod generate; +pub mod model; +pub mod processor; \ No newline at end of file diff --git a/src/models/glm_asr_nano/model.rs b/src/models/glm_asr_nano/model.rs new file mode 100644 index 0000000..e69de29 diff --git a/src/models/glm_asr_nano/processor.rs b/src/models/glm_asr_nano/processor.rs new file mode 100644 index 0000000..b857573 --- /dev/null +++ b/src/models/glm_asr_nano/processor.rs @@ -0,0 +1,94 @@ +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::Result; +use candle_core::{DType, Device, IndexOp, Tensor}; + +use crate::{ + models::glm_asr_nano::config::GlmAsrNanoProcessorConfig, + tokenizer::TokenizerModel, + utils::{audio_utils::extract_audios, extract_user_text}, +}; + +pub struct WhisperFeatureExtractor { + feature_size: usize, + sampling_rate: usize, + padding_value: f32, + hop_length: usize, + chunk_length: usize, + n_fft: usize, + dither: f32, +} + +pub struct GlmAsrNanoProcessor { + sampling_rate: usize, + chunk_length: usize, + audio_token: String, + audio_token_id: u32, + max_audio_len: usize, + default_transcription_prompt: String, + device: Device, +} + +impl GlmAsrNanoProcessor { + pub fn new(path: &str, device: &Device, dtype: DType) -> Result { + let path = path.to_string(); + assert!( + std::path::Path::new(&path).exists(), + "model path file not exists" + ); + let processor_config_path = path.to_string() + "/processor_config.json"; + assert!( + std::path::Path::new(&processor_config_path).exists(), + "processor_config.json not exists in model path" + ); + let processor_cfg: GlmAsrNanoProcessorConfig = + serde_json::from_slice(&std::fs::read(processor_config_path)?)?; + let audio_token = processor_cfg.audio_token.clone(); + let audio_token_id = 59260u32; + let max_audio_len = processor_cfg.max_audio_len; + let default_transcription_prompt = processor_cfg.default_transcription_prompt.clone(); + let sampling_rate = processor_cfg.feature_extractor.sampling_rate; + let chunk_length = processor_cfg.feature_extractor.chunk_length; + Ok(Self { + sampling_rate, + chunk_length, + audio_token, + audio_token_id, + max_audio_len, + default_transcription_prompt, + device: device.clone(), + }) + } + + // pub fn process_audio(&self, audios: Vec) -> Result { + // let window_size = self.sampling_rate * self.chunk_length; + // let max_windows = self.max_audio_len / self.chunk_length; + // let mut per_sample_windows = vec![]; + // let mut flat_chunks = vec![]; + // for audio_el in audios { + // let n_samples = audio_el.dim(0)?; + // let n_win = ((n_samples + window_size - 1) / window_size).max(1); + // let n_win = if n_win > max_windows { + // max_windows + // } else { + // n_win + // }; + // per_sample_windows.push(n_win); + // let time_cap = (n_win * window_size).min(n_samples); + // for i in 0..n_win { + // let start = i * window_size; + // let end = ((i + 1) * window_size).min(time_cap); + // flat_chunks.push(audio_el.i(start..end)?); + // } + // } + // } + + pub fn process_info( + &self, + mes: &ChatCompletionParameters + ) -> Result { + let audio_tensors = extract_audios(mes, &self.device, Some(self.sampling_rate))?; + println!("audio: {}", audio_tensors[0]); + // let audio = self.process_audio(audio_tensors)?; + Ok(audio_tensors[0].clone()) + } +} diff --git a/src/models/mod.rs b/src/models/mod.rs index 705d69e..3423ae2 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -7,6 +7,7 @@ pub mod qwen2_5vl; pub mod qwen3vl; pub mod rmbg2_0; pub mod voxcpm; +pub mod glm_asr_nano; use aha_openai_dive::v1::resources::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, diff --git a/src/models/voxcpm/generate.rs b/src/models/voxcpm/generate.rs index 97268b2..a7dee6c 100644 --- a/src/models/voxcpm/generate.rs +++ b/src/models/voxcpm/generate.rs @@ -213,7 +213,7 @@ impl GenerateModel for VoxCPMGenerate { extract_metadata_value::(&mes.metadata, "retry_badcase_ratio_threshold") .unwrap_or(6.0); let target_text = extract_user_text(&mes)?; - let prompt_wav = extract_audio_url(&mes)?; + let prompt_wav = extract_audio_url(&mes); let prompt_wav_path = if !prompt_wav.is_empty() { Some(prompt_wav[0].clone()) } else { diff --git a/src/models/voxcpm/model.rs b/src/models/voxcpm/model.rs index 08135c9..79da2ef 100644 --- a/src/models/voxcpm/model.rs +++ b/src/models/voxcpm/model.rs @@ -513,7 +513,7 @@ impl VoxCPMModel { let text_token = Tensor::cat(&[text_token, audio_start], D::Minus1)?; let text_length = text_token.dim(0)?; let mut audio = - load_audio_with_resample(&path, self.device.clone(), Some(self.sample_rate))?; + load_audio_with_resample(&path, &self.device, Some(self.sample_rate))?; let patch_len = self.patch_size * self.chunk_size; if audio.dim(1)? % patch_len != 0 { audio = audio.pad_with_zeros( @@ -730,7 +730,7 @@ impl VoxCPMModel { let text_token = Tensor::from_slice(&text_token, text_token.len(), &self.device)?; let mut audio = load_audio_with_resample( &prompt_wav_path, - self.device.clone(), + &self.device, Some(self.sample_rate), )?; let patch_len = self.patch_size * self.chunk_size; diff --git a/src/utils/audio_utils.rs b/src/utils/audio_utils.rs index 33af770..8a17778 100644 --- a/src/utils/audio_utils.rs +++ b/src/utils/audio_utils.rs @@ -14,6 +14,7 @@ use candle_core::{D, Device, Tensor}; use candle_nn::{Conv1d, Conv1dConfig, Module}; use hound::{SampleFormat, WavReader}; use num::integer::gcd; +use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; use crate::utils::get_default_save_dir; @@ -276,7 +277,7 @@ pub fn get_audio_path(path_str: &str) -> Result { } } -pub fn load_audio(path: &str, device: Device) -> Result<(Tensor, usize)> { +pub fn load_audio(path: &str, device: &Device) -> Result<(Tensor, usize)> { let audio_path = get_audio_path(path)?; let mut reader = WavReader::open(audio_path)?; let spec = reader.spec(); @@ -317,7 +318,7 @@ pub fn load_audio(path: &str, device: Device) -> Result<(Tensor, usize)> { samples.len() / spec.channels as usize, spec.channels as usize, ), - &device, + device, )? .t()?; if spec.channels > 1 { @@ -329,7 +330,7 @@ pub fn load_audio(path: &str, device: Device) -> Result<(Tensor, usize)> { pub fn load_audio_with_resample( path: &str, - device: Device, + device: &Device, target_sample_rate: Option, ) -> Result { let (mut audio, sr) = load_audio(path, device)?; @@ -387,7 +388,7 @@ pub fn get_audio_wav_u8(audio: &Tensor, sample_rate: u32) -> Result> { Ok(wav_buffer) } -pub fn extract_audio_url(mes: &ChatCompletionParameters) -> Result> { +pub fn extract_audio_url(mes: &ChatCompletionParameters) -> Vec { let mut audio_vec = Vec::new(); for chat_mes in mes.messages.clone() { if let ChatMessage::User { content, .. } = chat_mes.clone() @@ -400,20 +401,14 @@ pub fn extract_audio_url(mes: &ChatCompletionParameters) -> Result> } } } - // 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 = ret + &text + "\n" - // } - // } - // } - // } } - Ok(audio_vec) + audio_vec +} + +pub fn extract_audios(mes: &ChatCompletionParameters, device: &Device, target_sample_rate: Option) -> Result> { + let audio_url_vec = extract_audio_url(mes); + // 并行加载音频 + audio_url_vec.par_iter().map(|url| load_audio_with_resample(url, device, target_sample_rate)).collect() } // 从 ChatCompletionResponse 中提取音频数据 diff --git a/src/utils/img_utils.rs b/src/utils/img_utils.rs index 7120ab0..78838e4 100644 --- a/src/utils/img_utils.rs +++ b/src/utils/img_utils.rs @@ -8,7 +8,7 @@ use anyhow::{Result, anyhow}; use base64::{Engine, engine::general_purpose}; use candle_core::{DType, Device, Tensor}; use image::{DynamicImage, ImageBuffer, ImageReader, Rgb, RgbImage, imageops}; -use rayon::prelude::*; +use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; use crate::utils::{ceil_by_factor, floor_by_factor, round_by_factor}; diff --git a/tests/test_glm_asr_nano.rs b/tests/test_glm_asr_nano.rs new file mode 100644 index 0000000..0df6e9e --- /dev/null +++ b/tests/test_glm_asr_nano.rs @@ -0,0 +1,36 @@ +use aha::models::glm_asr_nano::generate::GlmAsrNanoGenerateModel; +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::{Result}; + +#[test] +fn glm_asr_nano_generate() -> Result<()> { + // RUST_BACKTRACE=1 cargo test -F cuda glm_asr_nano_generate -r -- --nocapture + let model_path = "/home/jhq/huggingface_model/zai-org/GLM-ASR-Nano-2512/"; + let message = r#" + { + "model": "glm-asr-nano", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "audio", + "audio_url": + { + "url": "file://./assets/audio/voice_01.wav" + } + }, + { + "type": "text", + "text": "Please transcribe this audio into text" + } + ] + } + ] + } + "#; + let mes: ChatCompletionParameters = serde_json::from_str(message)?; + let glm_asr_model = GlmAsrNanoGenerateModel::init(model_path, None, None)?; + let _ = glm_asr_model.generate(mes)?; + Ok(()) +} \ No newline at end of file