add qwen3 and fun-asr-nano

This commit is contained in:
jhqxxx
2026-01-15 21:57:12 +08:00
parent 2c1d5e3a14
commit d9b803d27e
39 changed files with 2577 additions and 307 deletions
+402 -12
View File
@@ -1,6 +1,7 @@
use std::fs::File;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::thread;
use std::{f64::consts::PI, io::Cursor};
use aha_openai_dive::v1::resources::chat::{
@@ -31,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;
use crate::utils::tensor_utils::{linspace, pad_replicate_last_dim};
// 重采样方法枚举
#[derive(Debug, Clone, Copy)]
@@ -279,19 +280,39 @@ pub fn load_audio_from_url(url: &str) -> Result<PathBuf> {
})
}
// pub fn _load_audio_bytes_from_url(url: &str) -> Result<Vec<u8>> {
// tokio::task::block_in_place(|| {
// let client = reqwest::blocking::Client::new();
// let response = client.get(url).send()?;
// if !response.status().is_success() {
// return Err(anyhow::anyhow!(
// "Failed to download file: {}",
// response.status()
// ));
// }
// let bytes = response.bytes()?.to_vec();
// Ok(bytes)
// })
// }
pub fn load_audio_bytes_from_url(url: &str) -> Result<Vec<u8>> {
tokio::task::block_in_place(|| {
let client = reqwest::blocking::Client::new();
let response = client.get(url).send()?;
if !response.status().is_success() {
return Err(anyhow::anyhow!(
"Failed to download file: {}",
response.status()
));
}
let bytes = response.bytes()?.to_vec();
Ok(bytes)
let url = url.to_string();
thread::spawn(move || {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let response = reqwest::get(&url).await?;
if !response.status().is_success() {
return Err(anyhow::anyhow!(
"Failed to download file: {}",
response.status()
));
}
let bytes = response.bytes().await?.to_vec();
Ok(bytes)
})
})
.join()
.unwrap()
}
pub fn get_audio_path(path_str: &str) -> Result<PathBuf> {
@@ -936,6 +957,31 @@ pub fn create_hann_window(window_size: usize, dtype: DType, device: &Device) ->
Ok(Tensor::from_vec(window, window_size, device)?.to_dtype(dtype)?)
}
pub fn crate_hamming_window(
window_size: usize,
periodic: bool,
alpha: f64,
beta: f64,
dtype: DType,
device: &Device,
) -> Result<Tensor> {
let denominator = if periodic {
window_size as f64
} else {
(window_size - 1) as f64
};
let window: Vec<f32> = (0..window_size)
.map(|i| {
let i_f64 = i as f64;
let val = alpha - beta * (2.0 * std::f64::consts::PI * i_f64 / denominator).cos();
val as f32
})
.collect();
Ok(Tensor::from_vec(window, window_size, device)?.to_dtype(dtype)?)
}
/// 梅尔频率刻度类型
#[derive(Debug, Clone, Copy)]
pub enum MelScale {
@@ -1099,3 +1145,347 @@ pub fn stft_audio(n_fft: usize, frame_wave: &[f32]) -> Result<Vec<f32>> {
let output: Vec<f32> = spectrum.iter().map(|complex| complex.norm_sqr()).collect();
Ok(output)
}
pub fn apply_stft(waveform: &Tensor) -> Result<Tensor> {
// waveform: (bs, n_frames, window_size)
let mut wave_fft = vec![];
let (batch_size, _, window_size) = waveform.dims3()?;
for bs in 0..batch_size {
let wave_i = waveform.i(bs)?;
let wave_i_vec = wave_i.to_vec2::<f32>()?;
let wave_i_fft_vec: Result<Vec<Vec<f32>>> = wave_i_vec
.par_iter()
.map(|frame_wave| stft_audio(window_size, frame_wave))
.collect();
let wave_i_fft_vec = wave_i_fft_vec?;
let wave_i_fft = Tensor::new(wave_i_fft_vec, waveform.device())?.unsqueeze(0)?;
wave_fft.push(wave_i_fft);
}
let magnitudes = Tensor::cat(&wave_fft, 0)?;
Ok(magnitudes)
}
pub fn kaldi_fbank(
waveform: &Tensor,
mel_energies: &Tensor,
window_shift: usize,
window_size: usize,
padded_window_size: usize,
dither: f32,
// energy_floor: f32,
// window_type: &str,
// sample_frequency: usize,
// snip_edges: bool,
) -> Result<Tensor> {
let (strided_input, _) = get_window(
waveform,
padded_window_size,
window_size,
window_shift,
dither,
true,
true,
0.97,
)?;
let spectrum = apply_stft(&strided_input)?;
let mel_energies = spectrum.broadcast_matmul(mel_energies)?;
let epsilon =
Tensor::new(1.192_092_9e-7_f32, waveform.device())?.broadcast_as(mel_energies.shape())?;
let mel_energies = mel_energies.maximum(&epsilon)?.log()?;
Ok(mel_energies)
}
pub fn apply_lfr(inputs: &Tensor, lfr_m: usize, lfr_n: usize) -> Result<Tensor> {
let (t, feat_dim) = inputs.dims2()?;
let t_lfr = (t as f32 / lfr_n as f32).ceil() as usize;
let left_padding_size = (lfr_m - 1) / 2;
let left_padding = inputs.narrow(0, 0, 1)?.repeat((left_padding_size, 1))?;
let mut inputs = Tensor::cat(&[&left_padding, inputs], 0)?;
let t = t + left_padding_size;
let last_idx = (t - lfr_m) / lfr_n + 1;
let num_padding = lfr_m - (t - last_idx * lfr_n);
if num_padding > 0 {
let num_padding =
(2 * lfr_m - 2 * t + (t_lfr - 1 + last_idx) * lfr_n) / 2 * (t_lfr - last_idx);
let right_padding = inputs.narrow(0, t - 1, 1)?.repeat((num_padding, 1))?;
inputs = Tensor::cat(&[&inputs, &right_padding], 0)?;
}
let mut outputs = vec![];
for i in 0..t_lfr {
let start = i * lfr_n;
let frame = inputs
.narrow(0, start, lfr_m)?
.reshape((1, lfr_m * feat_dim))?;
outputs.push(frame);
}
let lfr_outputs = Tensor::cat(&outputs, 0)?;
Ok(lfr_outputs)
}
pub fn get_waveform_and_window_properties(
sample_frequency: usize,
frame_shift: f32,
frame_length: f32,
round_to_power_of_two: bool,
) -> Result<(usize, usize, usize)> {
let window_shift = (sample_frequency as f32 * frame_shift * 0.001) as usize;
let window_size = (sample_frequency as f32 * frame_length * 0.001) as usize;
let padded_window_size = if round_to_power_of_two {
(window_size - 1).next_power_of_two()
} else {
window_size
};
Ok((window_shift, window_size, padded_window_size))
}
pub fn get_window(
waveform: &Tensor,
padded_window_size: usize,
window_size: usize,
window_shift: usize,
dither: f32,
remove_dc_offset: bool,
raw_energy: bool,
preemphasis_coefficient: f32,
) -> Result<(Tensor, Tensor)> {
let mut strided_input = extract_frames(waveform, window_size, window_shift)?;
// (ba, m, window_size)
if dither != 0.0 {
let rand_gauss = strided_input
.randn_like(0.0, 1.0)?
.affine(dither as f64, 0.0)?;
strided_input = strided_input.add(&rand_gauss)?;
}
if remove_dc_offset {
let row_means = strided_input.mean_keepdim(D::Minus1)?;
strided_input = strided_input.broadcast_sub(&row_means)?;
}
let signal_log_energy = if raw_energy {
let energy = strided_input.powf(2.0)?.sum(1)?.log()?;
Some(energy)
} else {
None
};
if preemphasis_coefficient != 0.0 {
let offset_strided_input = pad_replicate_last_dim(&strided_input, (1, 0))?
.affine(preemphasis_coefficient as f64, 0.0)?;
strided_input =
strided_input.sub(&offset_strided_input.narrow(D::Minus1, 0, window_size)?)?;
}
let windows = crate_hamming_window(
window_size,
false,
0.54,
0.46,
waveform.dtype(),
waveform.device(),
)?
.unsqueeze(0)?
.unsqueeze(0)?;
strided_input = strided_input.broadcast_mul(&windows)?;
if padded_window_size != window_size {
let padding_right = padded_window_size - window_size;
strided_input = strided_input.pad_with_zeros(D::Minus1, 0, padding_right)?;
}
let signal_log_energy = signal_log_energy.unwrap_or(strided_input.powf(2.0)?.sum(1)?.log()?);
Ok((strided_input, signal_log_energy))
}
/// 提取音频帧
pub fn extract_frames(
waveform: &Tensor,
window_size: usize,
window_shift: usize,
) -> Result<Tensor> {
// waveform ->(1, audio_len)
let waveform_len = waveform.dim(1)?;
let n_frames = 1 + (waveform_len - window_size) / window_shift;
let mut frames = Vec::with_capacity(n_frames);
for i in 0..n_frames {
let start = i * window_shift;
let frame = waveform.narrow(D::Minus1, start, window_size)?;
frames.push(frame);
}
let result = Tensor::cat(&frames, D::Minus1)?;
let bs = result.dim(0)?;
let reshaped = result.reshape((bs, n_frames, window_size))?;
Ok(reshaped)
}
pub fn inverse_mel_scale(mel_freq: &Tensor) -> Result<Tensor> {
Ok(mel_freq
.affine(1.0 / 1127.0, 0.0)?
.exp()?
.affine(1.0, -1.0)?
.affine(700.0, 0.0)?)
}
pub fn mel_scale(freq: &Tensor) -> Result<Tensor> {
Ok(freq.affine(1.0 / 700.0, 1.0)?.log()?.affine(1127.0, 0.0)?)
}
pub fn kaldi_get_mel_banks(
num_bins: usize,
window_length_padded: usize,
sample_freq: f32,
low_freq: f32,
high_freq: f32,
// vtln_low: f32,
// vtln_high: f32,
// vtln_warp_factor: f32,
device: &Device,
) -> Result<(Tensor, Tensor)> {
assert!(num_bins > 3, "Must have at least 3 mel bins");
assert!(
window_length_padded.is_multiple_of(2),
"window_length_padded must be even"
);
let num_fft_bins = window_length_padded as f32 / 2.0;
let nyquist = 0.5 * sample_freq;
let mut high_freq = high_freq;
if high_freq <= 0.0 {
high_freq += nyquist;
}
assert!(
(0.0 <= low_freq && low_freq < nyquist)
&& (0.0 < high_freq && high_freq <= nyquist)
&& (low_freq < high_freq),
"Bad values in options: low-freq {} and high-freq {} vs. nyquist {}",
low_freq,
high_freq,
nyquist
);
// FFT bin 宽度
let fft_bin_width = sample_freq / (window_length_padded as f32);
let mel_low_freq = hertz_to_mel(low_freq, MelScale::Kaldi);
let mel_high_freq = hertz_to_mel(high_freq, MelScale::Kaldi);
// 分频点之间的间隔
let mel_freq_delta = (mel_high_freq - mel_low_freq) / ((num_bins + 1) as f32);
// let mut vtln_high = vtln_high;
// if vtln_high < 0.0 {
// vtln_high += nyquist;
// }
// if vtln_warp_factor != 1.0 {
// assert!(
// low_freq < vtln_low
// && vtln_low < high_freq
// && 0.0 < vtln_high
// && vtln_high < high_freq
// && vtln_low < vtln_high,
// "Bad values in options: vtln-low {} and vtln-high {}, versus low-freq {} and high-freq {}",
// vtln_low,
// vtln_high,
// low_freq,
// high_freq
// );
// }
// 创建 bin 索引张量
let bins = Tensor::arange(0u32, num_bins as u32, device)?
.to_dtype(candle_core::DType::F32)?
.unsqueeze(1)?; // size(num_bins, 1)
// 计算梅尔刻度下的边界频率
let left_mel = bins.affine(mel_freq_delta as f64, mel_low_freq as f64)?;
let center_mel = bins
.affine(1.0, 1.0)?
.affine(mel_freq_delta as f64, mel_low_freq as f64)?;
let right_mel = bins
.affine(1.0, 2.0)?
.affine(mel_freq_delta as f64, mel_low_freq as f64)?;
// 如果使用 VTLN,则对频率进行扭曲
// let (left_mel, center_mel, right_mel) = if vtln_warp_factor != 1.0 {
// (
// vtln_warp_mel_freq(
// vtln_low,
// vtln_high,
// low_freq,
// high_freq,
// vtln_warp_factor,
// &left_mel,
// )?,
// vtln_warp_mel_freq(
// vtln_low,
// vtln_high,
// low_freq,
// high_freq,
// vtln_warp_factor,
// &center_mel,
// )?,
// vtln_warp_mel_freq(
// vtln_low,
// vtln_high,
// low_freq,
// high_freq,
// vtln_warp_factor,
// &right_mel,
// )?,
// )
// } else {
// (left_mel, center_mel, right_mel)
// };
// 转换中心频率回赫兹单位
let center_freqs = inverse_mel_scale(&center_mel)?;
// 创建 FFT bin 频率
let fft_bins = Tensor::arange(0u32, num_fft_bins as u32, device)?
.to_dtype(candle_core::DType::F32)?
.affine(fft_bin_width as f64, 0.0)?;
let mel = mel_scale(&fft_bins)?.unsqueeze(0)?; // size(1, num_fft_bins)
// 计算斜率
let up_slope = mel
.broadcast_sub(&left_mel)?
.broadcast_div(&center_mel.broadcast_sub(&left_mel)?)?;
let down_slope = right_mel
.broadcast_sub(&mel)?
.broadcast_div(&right_mel.broadcast_sub(&center_mel)?)?;
// left_mel < center_mel < right_mel 所以我们可以取两个斜率的最小值并限制负值
let min_slopes = up_slope.minimum(&down_slope)?;
let zeros = Tensor::zeros(min_slopes.dims(), candle_core::DType::F32, device)?;
let bins_tensor = min_slopes.maximum(&zeros)?;
// let bins_tensor = if vtln_warp_factor == 1.0 {
// // left_mel < center_mel < right_mel 所以我们可以取两个斜率的最小值并限制负值
// let min_slopes = up_slope.minimum(&down_slope)?;
// let zeros = Tensor::zeros(min_slopes.dims(), candle_core::DType::F32, device)?;
// min_slopes.maximum(&zeros)?
// } else {
// // 扭曲可能会改变 left_mel, center_mel, right_mel 的顺序
// let zeros = Tensor::zeros(up_slope.dims(), candle_core::DType::F32, device)?;
// let mut bins_tensor = zeros.clone();
// // 创建索引掩码
// let up_idx = mel
// .gt_tensor(&left_mel)?
// .and(&mel.le_tensor(&center_mel)?)?; // left_mel < mel <= center_mel
// let down_idx = mel
// .gt_tensor(&center_mel)?
// .and(&mel.lt_tensor(&right_mel)?)?; // center_mel < mel < right_mel
// bins_tensor = bins_tensor.where_cond(&up_idx, &up_slope)?;
// bins_tensor = bins_tensor.where_cond(&down_idx, &down_slope)?;
// bins_tensor
// };
Ok((bins_tensor, center_freqs))
}