style: cargo fmt --all (18 files)

Auto-merged by ci-doctor.
This commit is contained in:
redclawsystems
2026-05-07 16:30:04 +00:00
parent 789448f196
commit ae53983c03
18 changed files with 250 additions and 168 deletions
+4 -4
View File
@@ -2,19 +2,19 @@
//! //!
//! Provides the [`CudaDevice`] type for managing CUDA contexts and streams. //! Provides the [`CudaDevice`] type for managing CUDA contexts and streams.
#[cfg(feature = "cuda")]
use crate::CudaBackend;
#[cfg(feature = "cuda")] #[cfg(feature = "cuda")]
use cudarc::cublas::CudaBlas; use cudarc::cublas::CudaBlas;
#[cfg(feature = "cuda")] #[cfg(feature = "cuda")]
use cudarc::driver::{CudaContext, CudaStream}; use cudarc::driver::{CudaContext, CudaStream};
#[cfg(feature = "cuda")] #[cfg(feature = "cuda")]
use crate::CudaBackend;
#[cfg(feature = "cuda")]
use rtx_backend::{DeviceId, DeviceOps};
#[cfg(feature = "cuda")]
use once_cell::sync::{Lazy, OnceCell}; use once_cell::sync::{Lazy, OnceCell};
#[cfg(feature = "cuda")] #[cfg(feature = "cuda")]
use parking_lot::RwLock; use parking_lot::RwLock;
#[cfg(feature = "cuda")] #[cfg(feature = "cuda")]
use rtx_backend::{DeviceId, DeviceOps};
#[cfg(feature = "cuda")]
use std::collections::HashMap; use std::collections::HashMap;
#[cfg(feature = "cuda")] #[cfg(feature = "cuda")]
use std::sync::Arc; use std::sync::Arc;
@@ -2,13 +2,13 @@
//! //!
//! Extracted from `compute` to keep module size within limits. //! Extracted from `compute` to keep module size within limits.
use super::{create_shader_module, create_uniform_buffer};
use crate::{WebGpuDevice, WebGpuTensorPrimitive, shaders}; use crate::{WebGpuDevice, WebGpuTensorPrimitive, shaders};
use wgpu::{ use wgpu::{
BindGroupDescriptor, BindGroupEntry, BindGroupLayout, BindGroupLayoutDescriptor, BindGroupDescriptor, BindGroupEntry, BindGroupLayout, BindGroupLayoutDescriptor,
BindGroupLayoutEntry, BindingType, BufferBindingType, BufferUsages, ComputePipeline, BindGroupLayoutEntry, BindingType, BufferBindingType, BufferUsages, ComputePipeline,
ComputePipelineDescriptor, PipelineLayoutDescriptor, ShaderStages, ComputePipelineDescriptor, PipelineLayoutDescriptor, ShaderStages,
}; };
use super::{create_shader_module, create_uniform_buffer};
/// Parameters for 2D convolution. /// Parameters for 2D convolution.
#[repr(C)] #[repr(C)]
@@ -132,7 +132,11 @@ pub struct ReductionParams {
} }
/// Create a shader module from WGSL source. /// Create a shader module from WGSL source.
pub(super) fn create_shader_module(device: &WebGpuDevice, source: &str, label: &str) -> ShaderModule { pub(super) fn create_shader_module(
device: &WebGpuDevice,
source: &str,
label: &str,
) -> ShaderModule {
device device
.wgpu_device() .wgpu_device()
.create_shader_module(ShaderModuleDescriptor { .create_shader_module(ShaderModuleDescriptor {
@@ -1037,4 +1041,7 @@ pub struct BmmParams {
} }
pub mod conv; pub mod conv;
pub use conv::{Conv2dParams, Pool2dParams, AvgPool2dParams, dispatch_conv2d, dispatch_max_pool2d, dispatch_avg_pool2d}; pub use conv::{
AvgPool2dParams, Conv2dParams, Pool2dParams, dispatch_avg_pool2d, dispatch_conv2d,
dispatch_max_pool2d,
};
+28 -16
View File
@@ -178,10 +178,7 @@ fn main() -> Result<()> {
tracing::info!("loaded Moonshine-tiny"); tracing::info!("loaded Moonshine-tiny");
// Optional emotion2vec for emotion-classification metrics. // Optional emotion2vec for emotion-classification metrics.
let target_idx = cli let target_idx = cli.target_emotion.as_deref().and_then(emotion_index);
.target_emotion
.as_deref()
.and_then(emotion_index);
if cli.emotion2vec.is_some() && cli.target_emotion.is_some() && target_idx.is_none() { if cli.emotion2vec.is_some() && cli.target_emotion.is_some() && target_idx.is_none() {
anyhow::bail!( anyhow::bail!(
"--target-emotion '{}' is unrecognized — must be angry, disgusted, fearful, happy, neutral, excited, sad, or surprised", "--target-emotion '{}' is unrecognized — must be angry, disgusted, fearful, happy, neutral, excited, sad, or surprised",
@@ -222,8 +219,8 @@ fn main() -> Result<()> {
if line.trim().is_empty() { if line.trim().is_empty() {
continue; continue;
} }
let row: PairRow = serde_json::from_str(&line) let row: PairRow =
.with_context(|| format!("parse row {idx}: {line}"))?; serde_json::from_str(&line).with_context(|| format!("parse row {idx}: {line}"))?;
// Speaker similarity at 16 kHz. WavLM-SV's TDNN front-end has // Speaker similarity at 16 kHz. WavLM-SV's TDNN front-end has
// kernel size 5 with stride 2 chains, so it requires at least a // kernel size 5 with stride 2 chains, so it requires at least a
@@ -254,15 +251,19 @@ fn main() -> Result<()> {
.trim() .trim()
.to_string(); .to_string();
let wer = row.ref_text.as_ref().map(|t| word_error_rate(t, &transcript)); let wer = row
.ref_text
.as_ref()
.map(|t| word_error_rate(t, &transcript));
// Amplitude of gen_wav at native rate. // Amplitude of gen_wav at native rate.
let gen_native = audio_io::load_mono_24k(&row.gen_wav)?; let gen_native = audio_io::load_mono_24k(&row.gen_wav)?;
let (peak_db, rms_db) = peak_rms_db(&gen_native); let (peak_db, rms_db) = peak_rms_db(&gen_native);
// Optional emotion classification on gen_wav at 16 kHz. // Optional emotion classification on gen_wav at 16 kHz.
let (target_emotion_prob, top_emotion, top_emotion_prob) = let (target_emotion_prob, top_emotion, top_emotion_prob) = if let Some(emo) =
if let Some(emo) = emo_model.as_ref() { emo_model.as_ref()
{
if gen_16k.len() < MIN_SV_SAMPLES { if gen_16k.len() < MIN_SV_SAMPLES {
(None, None, None) (None, None, None)
} else { } else {
@@ -287,7 +288,9 @@ fn main() -> Result<()> {
.map(|(i, &p)| (i, p)) .map(|(i, &p)| (i, p))
.unwrap_or((0, 0.0)); .unwrap_or((0, 0.0));
// Use the trait's tag for human-readable label. // Use the trait's tag for human-readable label.
let top_label = emo.classify(&gen_16k).map(|l| format!("{l:?}")) let top_label = emo
.classify(&gen_16k)
.map(|l| format!("{l:?}"))
.unwrap_or_else(|_| format!("class_{top_i}")); .unwrap_or_else(|_| format!("class_{top_i}"));
let _ = top_i; let _ = top_i;
(target_p, Some(top_label), Some(top_p)) (target_p, Some(top_label), Some(top_p))
@@ -334,7 +337,10 @@ fn main() -> Result<()> {
eprintln!("rows {n_rows}"); eprintln!("rows {n_rows}");
eprintln!("speaker_cosine {:.3} (mean)", sum_cos / n_rows as f64); eprintln!("speaker_cosine {:.3} (mean)", sum_cos / n_rows as f64);
if n_wer > 0 { if n_wer > 0 {
eprintln!("wer {:.3} (mean over {n_wer})", sum_wer / n_wer as f64); eprintln!(
"wer {:.3} (mean over {n_wer})",
sum_wer / n_wer as f64
);
} else { } else {
eprintln!("wer (no ref_text in any row)"); eprintln!("wer (no ref_text in any row)");
} }
@@ -358,9 +364,7 @@ fn word_error_rate(reference: &str, hypothesis: &str) -> f32 {
curr[0] = i; curr[0] = i;
for j in 1..=m { for j in 1..=m {
let cost = if r[i - 1] == h[j - 1] { 0 } else { 1 }; let cost = if r[i - 1] == h[j - 1] { 0 } else { 1 };
curr[j] = (prev[j] + 1) curr[j] = (prev[j] + 1).min(curr[j - 1] + 1).min(prev[j - 1] + cost);
.min(curr[j - 1] + 1)
.min(prev[j - 1] + cost);
} }
std::mem::swap(&mut prev, &mut curr); std::mem::swap(&mut prev, &mut curr);
} }
@@ -395,7 +399,15 @@ fn peak_rms_db(pcm: &[f32]) -> (f32, f32) {
sum_sq += (s as f64) * (s as f64); sum_sq += (s as f64) * (s as f64);
} }
let rms = (sum_sq / pcm.len() as f64).sqrt() as f32; let rms = (sum_sq / pcm.len() as f64).sqrt() as f32;
let peak_db = if peak < 1e-6 { -100.0 } else { 20.0 * peak.log10() }; let peak_db = if peak < 1e-6 {
let rms_db = if rms < 1e-6 { -100.0 } else { 20.0 * rms.log10() }; -100.0
} else {
20.0 * peak.log10()
};
let rms_db = if rms < 1e-6 {
-100.0
} else {
20.0 * rms.log10()
};
(peak_db, rms_db) (peak_db, rms_db)
} }
@@ -210,10 +210,11 @@ fn main() -> Result<()> {
let mid = num_frames / 2; let mid = num_frames / 2;
let c0 = codes.narrow(2, mid, 1)?.narrow(1, 0, 1)?; let c0 = codes.narrow(2, mid, 1)?.narrow(1, 0, 1)?;
let target_c0: u32 = c0.to_dtype(DType::U32)?.flatten_all()?.to_vec1::<u32>()?[0]; let target_c0: u32 = c0.to_dtype(DType::U32)?.flatten_all()?.to_vec1::<u32>()?[0];
generator generator.model.inner.capture_decoder_activations(
.model &prompt.tokens,
.inner &prompt.mask,
.capture_decoder_activations(&prompt.tokens, &prompt.mask, target_c0)? target_c0,
)?
} else { } else {
let segment = Segment::new(speaker, row.transcript.clone(), audio); let segment = Segment::new(speaker, row.transcript.clone(), audio);
let prompt = build_prompt( let prompt = build_prompt(
+1 -2
View File
@@ -125,8 +125,7 @@ fn should_flush(buf: &str, policy: FlushPolicy) -> bool {
/// synthesis call. Use this to apply per-sentence steering (e.g. emotion /// synthesis call. Use this to apply per-sentence steering (e.g. emotion
/// shifts mid-reply). Sync — apply_steering is non-async — so the hook /// shifts mid-reply). Sync — apply_steering is non-async — so the hook
/// fits cleanly between LLM token consumption and the synthesize call. /// fits cleanly between LLM token consumption and the synthesize call.
pub type PreSentenceHook = pub type PreSentenceHook = Box<dyn FnMut(&mut Generator, &str) -> Result<()> + Send + 'static>;
Box<dyn FnMut(&mut Generator, &str) -> Result<()> + Send + 'static>;
pub struct Converse<'a, L: LlmClient> { pub struct Converse<'a, L: LlmClient> {
llm: &'a L, llm: &'a L,
+2 -1
View File
@@ -419,7 +419,8 @@ impl Layer {
// Optional layer-by-layer dump for qmatmul bisection. // Optional layer-by-layer dump for qmatmul bisection.
if std::env::var("CSM_DUMP_LAYERS").is_ok() if std::env::var("CSM_DUMP_LAYERS").is_ok()
&& let Some(idx) = dbg_idx && let Some(idx) = dbg_idx
&& seqlen_offset == 0 { && seqlen_offset == 0
{
let f = out.flatten_all()?.to_vec1::<f32>().unwrap_or_default(); let f = out.flatten_all()?.to_vec1::<f32>().unwrap_or_default();
let n: f32 = f.iter().map(|x| x * x).sum::<f32>().sqrt(); let n: f32 = f.iter().map(|x| x * x).sum::<f32>().sqrt();
eprintln!( eprintln!(
+4 -2
View File
@@ -232,7 +232,8 @@ fn vad_intervals(probs: &[f32], cfg: &DiarizationConfig) -> Vec<(usize, usize)>
for i in 0..active.len() { for i in 0..active.len() {
if active[i] { if active[i] {
if let Some(prev) = last_speech if let Some(prev) = last_speech
&& i - prev <= max_silence_chunks + 1 { && i - prev <= max_silence_chunks + 1
{
active[(prev + 1)..i].fill(true); active[(prev + 1)..i].fill(true);
} }
last_speech = Some(i); last_speech = Some(i);
@@ -302,7 +303,8 @@ fn agglomerative_cluster(embs: &[&[f32]], threshold: f32, n_speakers: Option<usi
break; break;
} }
if let Some(k) = n_speakers if let Some(k) = n_speakers
&& clusters.len() <= k { && clusters.len() <= k
{
break; break;
} }
// Find the closest pair using average linkage. // Find the closest pair using average linkage.
+12 -7
View File
@@ -344,9 +344,10 @@ impl Generator {
// CFG path: dual backbone if all preconditions are met. The // CFG path: dual backbone if all preconditions are met. The
// schedule field takes precedence over the legacy `cfg_scale`. // schedule field takes precedence over the legacy `cfg_scale`.
let cfg_schedule = opts let cfg_schedule = opts.cfg_schedule.or_else(|| {
.cfg_schedule opts.cfg_scale
.or_else(|| opts.cfg_scale.map(crate::cfg_schedule::CfgSchedule::Constant)); .map(crate::cfg_schedule::CfgSchedule::Constant)
});
let cfg_active = cfg_schedule.map(|s| s.is_active()).unwrap_or(false) let cfg_active = cfg_schedule.map(|s| s.is_active()).unwrap_or(false)
&& !context.is_empty() && !context.is_empty()
&& self.model.inner.cfg_enabled(); && self.model.inner.cfg_enabled();
@@ -416,7 +417,8 @@ impl Generator {
} }
if let Some(g) = rep_guard.as_mut() if let Some(g) = rep_guard.as_mut()
&& g.observe(&sampled) { && g.observe(&sampled)
{
tracing::warn!( tracing::warn!(
"loop-escape: repetition guard tripped at frame {frame_idx}; ending generation" "loop-escape: repetition guard tripped at frame {frame_idx}; ending generation"
); );
@@ -525,7 +527,8 @@ impl Generator {
break; break;
} }
if let Some(g) = rep_guard.as_mut() if let Some(g) = rep_guard.as_mut()
&& g.observe(&sampled) { && g.observe(&sampled)
{
tracing::warn!( tracing::warn!(
"loop-escape: repetition guard tripped at frame {frame_idx}; ending" "loop-escape: repetition guard tripped at frame {frame_idx}; ending"
); );
@@ -563,7 +566,8 @@ impl Generator {
} }
// Drain any internal Mimi buffering by feeding a `None` step. // Drain any internal Mimi buffering by feeding a `None` step.
if let Some(tail) = self.mimi.decode_step(None)? if let Some(tail) = self.mimi.decode_step(None)?
&& !tail.is_empty() { && !tail.is_empty()
{
on_chunk(&tail)?; on_chunk(&tail)?;
full_pcm.extend_from_slice(&tail); full_pcm.extend_from_slice(&tail);
} }
@@ -651,7 +655,8 @@ where
let codes = Tensor::from_vec(flat, (1, num_codebooks, k), device)?.to_dtype(DType::U32)?; let codes = Tensor::from_vec(flat, (1, num_codebooks, k), device)?.to_dtype(DType::U32)?;
pending.clear(); pending.clear();
if let Some(samples) = mimi.decode_step(Some(&codes))? if let Some(samples) = mimi.decode_step(Some(&codes))?
&& !samples.is_empty() { && !samples.is_empty()
{
on_chunk(&samples)?; on_chunk(&samples)?;
full_pcm.extend_from_slice(&samples); full_pcm.extend_from_slice(&samples);
} }
+2 -1
View File
@@ -236,7 +236,8 @@ impl LlmClient for OpenAiCompatibleClient {
let mut body = serde_json::to_value(&req) let mut body = serde_json::to_value(&req)
.map_err(|e| CsmError::Config(format!("LLM serialize: {e}")))?; .map_err(|e| CsmError::Config(format!("LLM serialize: {e}")))?;
if !config.extra_body.is_empty() if !config.extra_body.is_empty()
&& let Some(obj) = body.as_object_mut() { && let Some(obj) = body.as_object_mut()
{
for (k, v) in config.extra_body.iter() { for (k, v) in config.extra_body.iter() {
obj.insert(k.clone(), v.clone()); obj.insert(k.clone(), v.clone());
} }
+5 -4
View File
@@ -489,7 +489,10 @@ pub fn load_lora_set_from_safetensors<P: AsRef<Path>>(
); );
} }
if !b_map.is_empty() { if !b_map.is_empty() {
tracing::warn!("LoRA adapter has dangling .lora_b without .lora_a: {:?}", b_map.keys()); tracing::warn!(
"LoRA adapter has dangling .lora_b without .lora_a: {:?}",
b_map.keys()
);
} }
Ok(set) Ok(set)
} }
@@ -526,9 +529,7 @@ pub fn merge_into_safetensors<P: AsRef<Path>>(
.to_dtype(candle_core::DType::F32) .to_dtype(candle_core::DType::F32)
.map_err(|e| CsmError::Other(anyhow::anyhow!("upcast {name}: {e}")))?; .map_err(|e| CsmError::Other(anyhow::anyhow!("upcast {name}: {e}")))?;
let dims = base_f32.dims(); let dims = base_f32.dims();
if dims.len() != 2 if dims.len() != 2 || dims[0] != adapter.out_features || dims[1] != adapter.in_features
|| dims[0] != adapter.out_features
|| dims[1] != adapter.in_features
{ {
tracing::warn!( tracing::warn!(
"LoRA shape mismatch for `{name}`: base={dims:?} adapter=({},{}) — passing through unchanged", "LoRA shape mismatch for `{name}`: base={dims:?} adapter=({},{}) — passing through unchanged",
+73 -23
View File
@@ -142,7 +142,11 @@ impl ConvStem {
}, },
vb.pp("conv3"), vb.pp("conv3"),
)?; )?;
Ok(Self { conv1, conv2, conv3 }) Ok(Self {
conv1,
conv2,
conv3,
})
} }
/// Forward: `(B, 1, T_audio)` raw waveform → `(B, T_seq, 288)` where /// Forward: `(B, 1, T_audio)` raw waveform → `(B, T_seq, 288)` where
@@ -167,9 +171,7 @@ impl ConvStem {
/// Loader: open the HF safetensors and construct a `ConvStem`. Useful /// Loader: open the HF safetensors and construct a `ConvStem`. Useful
/// for the standalone Phase 8.5 smoke test. /// for the standalone Phase 8.5 smoke test.
pub fn load_conv_stem(weights_path: &std::path::Path, device: &Device) -> Result<ConvStem> { pub fn load_conv_stem(weights_path: &std::path::Path, device: &Device) -> Result<ConvStem> {
let vb = unsafe { let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device) }?;
VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device)
}?;
ConvStem::new(vb.pp("model").pp("encoder")) ConvStem::new(vb.pp("model").pp("encoder"))
} }
@@ -190,7 +192,13 @@ struct RotaryCache {
} }
impl RotaryCache { impl RotaryCache {
fn new(rotary_dim: usize, max_seq: usize, theta: f64, dtype: DType, dev: &Device) -> Result<Self> { fn new(
rotary_dim: usize,
max_seq: usize,
theta: f64,
dtype: DType,
dev: &Device,
) -> Result<Self> {
assert!(rotary_dim.is_multiple_of(2), "rotary_dim must be even"); assert!(rotary_dim.is_multiple_of(2), "rotary_dim must be even");
let inv_freq: Vec<f32> = (0..rotary_dim) let inv_freq: Vec<f32> = (0..rotary_dim)
.step_by(2) .step_by(2)
@@ -203,7 +211,11 @@ impl RotaryCache {
let freqs = positions.matmul(&inv_freq.reshape((1, rotary_dim / 2))?)?; let freqs = positions.matmul(&inv_freq.reshape((1, rotary_dim / 2))?)?;
let cos = freqs.cos()?.to_dtype(dtype)?; let cos = freqs.cos()?.to_dtype(dtype)?;
let sin = freqs.sin()?.to_dtype(dtype)?; let sin = freqs.sin()?.to_dtype(dtype)?;
Ok(Self { cos, sin, rotary_dim }) Ok(Self {
cos,
sin,
rotary_dim,
})
} }
/// Apply partial RoPE to `q` of shape `(B, H, T, head_dim)`. /// Apply partial RoPE to `q` of shape `(B, H, T, head_dim)`.
@@ -251,7 +263,14 @@ impl EncoderAttention {
let k_proj = candle_nn::linear_no_bias(h, h, vb.pp("k_proj"))?; let k_proj = candle_nn::linear_no_bias(h, h, vb.pp("k_proj"))?;
let v_proj = candle_nn::linear_no_bias(h, h, vb.pp("v_proj"))?; let v_proj = candle_nn::linear_no_bias(h, h, vb.pp("v_proj"))?;
let o_proj = candle_nn::linear_no_bias(h, h, vb.pp("o_proj"))?; let o_proj = candle_nn::linear_no_bias(h, h, vb.pp("o_proj"))?;
Ok(Self { q_proj, k_proj, v_proj, o_proj, n_heads, head_dim }) Ok(Self {
q_proj,
k_proj,
v_proj,
o_proj,
n_heads,
head_dim,
})
} }
fn forward(&self, xs: &Tensor, rope: &RotaryCache) -> Result<Tensor> { fn forward(&self, xs: &Tensor, rope: &RotaryCache) -> Result<Tensor> {
@@ -260,9 +279,18 @@ impl EncoderAttention {
let k = self.k_proj.forward(xs)?; let k = self.k_proj.forward(xs)?;
let v = self.v_proj.forward(xs)?; let v = self.v_proj.forward(xs)?;
// (B, T, H) -> (B, H_heads, T, head_dim) // (B, T, H) -> (B, H_heads, T, head_dim)
let q = q.reshape((b, t, self.n_heads, self.head_dim))?.transpose(1, 2)?.contiguous()?; let q = q
let k = k.reshape((b, t, self.n_heads, self.head_dim))?.transpose(1, 2)?.contiguous()?; .reshape((b, t, self.n_heads, self.head_dim))?
let v = v.reshape((b, t, self.n_heads, self.head_dim))?.transpose(1, 2)?.contiguous()?; .transpose(1, 2)?
.contiguous()?;
let k = k
.reshape((b, t, self.n_heads, self.head_dim))?
.transpose(1, 2)?
.contiguous()?;
let v = v
.reshape((b, t, self.n_heads, self.head_dim))?
.transpose(1, 2)?
.contiguous()?;
let q = rope.apply(&q)?; let q = rope.apply(&q)?;
let k = rope.apply(&k)?; let k = rope.apply(&k)?;
// Collapse (B, H, T, D) -> (B*H, T, D) for the matmul. candle's // Collapse (B, H, T, D) -> (B*H, T, D) for the matmul. candle's
@@ -324,7 +352,12 @@ impl EncoderLayer {
let self_attn = EncoderAttention::new(cfg, vb.pp("self_attn"))?; let self_attn = EncoderAttention::new(cfg, vb.pp("self_attn"))?;
let post_attn_ln = layer_norm_weight_only(h, 1e-5, vb.pp("post_attention_layernorm"))?; let post_attn_ln = layer_norm_weight_only(h, 1e-5, vb.pp("post_attention_layernorm"))?;
let mlp = EncoderMlp::new(cfg, vb.pp("mlp"))?; let mlp = EncoderMlp::new(cfg, vb.pp("mlp"))?;
Ok(Self { input_ln, self_attn, post_attn_ln, mlp }) Ok(Self {
input_ln,
self_attn,
post_attn_ln,
mlp,
})
} }
fn forward(&self, xs: &Tensor, rope: &RotaryCache) -> Result<Tensor> { fn forward(&self, xs: &Tensor, rope: &RotaryCache) -> Result<Tensor> {
@@ -369,8 +402,14 @@ impl Encoder {
for i in 0..cfg.encoder_num_hidden_layers { for i in 0..cfg.encoder_num_hidden_layers {
layers.push(EncoderLayer::new(cfg, layer_vb.pp(i))?); layers.push(EncoderLayer::new(cfg, layer_vb.pp(i))?);
} }
let final_ln = layer_norm_weight_only(cfg.hidden_size, 1e-5, vb.pp("encoder").pp("layer_norm"))?; let final_ln =
Ok(Self { stem, layers, final_ln, rope }) layer_norm_weight_only(cfg.hidden_size, 1e-5, vb.pp("encoder").pp("layer_norm"))?;
Ok(Self {
stem,
layers,
final_ln,
rope,
})
} }
pub fn forward(&self, pcm: &Tensor) -> Result<Tensor> { pub fn forward(&self, pcm: &Tensor) -> Result<Tensor> {
@@ -389,9 +428,7 @@ pub fn load_encoder(
device: &Device, device: &Device,
cfg: &MoonshineConfig, cfg: &MoonshineConfig,
) -> Result<Encoder> { ) -> Result<Encoder> {
let vb = unsafe { let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device) }?;
VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device)
}?;
Encoder::new(cfg, vb.pp("model")) Encoder::new(cfg, vb.pp("model"))
} }
@@ -642,7 +679,11 @@ pub struct Decoder {
impl Decoder { impl Decoder {
pub fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> { pub fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
let embed = candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("decoder").pp("embed_tokens"))?; let embed = candle_nn::embedding(
cfg.vocab_size,
cfg.hidden_size,
vb.pp("decoder").pp("embed_tokens"),
)?;
let head_dim = cfg.hidden_size / cfg.decoder_num_attention_heads; let head_dim = cfg.hidden_size / cfg.decoder_num_attention_heads;
let rotary_dim = ((head_dim as f64 * cfg.partial_rotary_factor) as usize / 2) * 2; let rotary_dim = ((head_dim as f64 * cfg.partial_rotary_factor) as usize / 2) * 2;
let rope = RotaryCache::new( let rope = RotaryCache::new(
@@ -661,7 +702,13 @@ impl Decoder {
// `encoder.layer_norm.weight`). Inspector dump confirmed. // `encoder.layer_norm.weight`). Inspector dump confirmed.
let final_ln = layer_norm_weight_only(cfg.hidden_size, 1e-5, vb.pp("decoder").pp("norm"))?; let final_ln = layer_norm_weight_only(cfg.hidden_size, 1e-5, vb.pp("decoder").pp("norm"))?;
let embed_weight = embed.embeddings().clone(); let embed_weight = embed.embeddings().clone();
Ok(Self { embed, layers, final_ln, rope, embed_weight }) Ok(Self {
embed,
layers,
final_ln,
rope,
embed_weight,
})
} }
/// Forward over `tokens` shape `(B, T)` with encoder output `enc` /// Forward over `tokens` shape `(B, T)` with encoder output `enc`
@@ -688,9 +735,7 @@ pub fn load_full(
device: &Device, device: &Device,
cfg: &MoonshineConfig, cfg: &MoonshineConfig,
) -> Result<(Encoder, Decoder)> { ) -> Result<(Encoder, Decoder)> {
let vb = unsafe { let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device) }?;
VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device)
}?;
let encoder = Encoder::new(cfg, vb.pp("model"))?; let encoder = Encoder::new(cfg, vb.pp("model"))?;
let decoder = Decoder::new(cfg, vb.pp("model"))?; let decoder = Decoder::new(cfg, vb.pp("model"))?;
Ok((encoder, decoder)) Ok((encoder, decoder))
@@ -741,7 +786,9 @@ impl Decoder {
/// Convenience: load the HF tokenizer.json for Moonshine. Caller passes /// Convenience: load the HF tokenizer.json for Moonshine. Caller passes
/// the path returned by hf_hub. /// the path returned by hf_hub.
pub fn load_tokenizer(path: &std::path::Path) -> std::result::Result<tokenizers::Tokenizer, Box<dyn std::error::Error + Send + Sync>> { pub fn load_tokenizer(
path: &std::path::Path,
) -> std::result::Result<tokenizers::Tokenizer, Box<dyn std::error::Error + Send + Sync>> {
tokenizers::Tokenizer::from_file(path) tokenizers::Tokenizer::from_file(path)
} }
@@ -914,7 +961,10 @@ impl Decoder {
)?; )?;
h = (h + attn_out)?; h = (h + attn_out)?;
let normed = layer.post_attn_ln.forward(&h)?; let normed = layer.post_attn_ln.forward(&h)?;
let cross_out = layer.cross_attn.forward_step(&normed, &cache.cross_k[i], &cache.cross_v[i])?; let cross_out =
layer
.cross_attn
.forward_step(&normed, &cache.cross_k[i], &cache.cross_v[i])?;
h = (h + cross_out)?; h = (h + cross_out)?;
let normed = layer.final_ln.forward(&h)?; let normed = layer.final_ln.forward(&h)?;
let mlp_out = layer.mlp.forward(&normed)?; let mlp_out = layer.mlp.forward(&normed)?;
+2 -2
View File
@@ -79,8 +79,8 @@ impl StreamingHpfState {
pub fn process(&mut self, samples: &mut [f32]) { pub fn process(&mut self, samples: &mut [f32]) {
for s in samples.iter_mut() { for s in samples.iter_mut() {
let x0 = *s; let x0 = *s;
let y0 = let y0 = self.b0 * x0 + self.b1 * self.x1 + self.b2 * self.x2
self.b0 * x0 + self.b1 * self.x1 + self.b2 * self.x2 - self.a1 * self.y1 - self.a1 * self.y1
- self.a2 * self.y2; - self.a2 * self.y2;
self.x2 = self.x1; self.x2 = self.x1;
self.x1 = x0; self.x1 = x0;
+4 -2
View File
@@ -178,8 +178,10 @@ impl ProsodyDetector {
min_lag, min_lag,
max_lag, max_lag,
self.voiced_threshold, self.voiced_threshold,
) ) && voiced
&& voiced && f0 >= self.min_f0_hz && f0 <= self.max_f0_hz { && f0 >= self.min_f0_hz
&& f0 <= self.max_f0_hz
{
voiced_f0s.push(f0); voiced_f0s.push(f0);
} }
total_frames += 1; total_frames += 1;
+5 -2
View File
@@ -71,7 +71,9 @@ impl LayerSteering {
} }
let dims = vec.dims(); let dims = vec.dims();
let reshaped = match dims.len() { let reshaped = match dims.len() {
1 => vec.reshape((1, dims[0])).map_err(|e| CsmError::Config(e.to_string()))?, 1 => vec
.reshape((1, dims[0]))
.map_err(|e| CsmError::Config(e.to_string()))?,
2 if dims[0] == 1 => vec, 2 if dims[0] == 1 => vec,
_ => { _ => {
return Err(CsmError::Config(format!( return Err(CsmError::Config(format!(
@@ -193,7 +195,8 @@ mod tests {
let dev = Device::Cpu; let dev = Device::Cpu;
let mut s = LayerSteering::empty(8); let mut s = LayerSteering::empty(8);
for i in 0..8 { for i in 0..8 {
s.set_layer(i, Tensor::ones((4,), DType::F32, &dev).unwrap()).unwrap(); s.set_layer(i, Tensor::ones((4,), DType::F32, &dev).unwrap())
.unwrap();
} }
assert_eq!(s.active_layers(), vec![0, 1, 2, 3, 4, 5, 6, 7]); assert_eq!(s.active_layers(), vec![0, 1, 2, 3, 4, 5, 6, 7]);
s.restrict_to_layers(&[2, 5]); s.restrict_to_layers(&[2, 5]);
+2 -1
View File
@@ -600,7 +600,8 @@ pub fn ctc_greedy_decode(logits: &Tensor, vocab: &[&str]) -> CsmResult<String> {
continue; continue;
} }
if (id as usize) != CTC_BLANK_ID if (id as usize) != CTC_BLANK_ID
&& let Some(tok) = vocab.get(id as usize) { && let Some(tok) = vocab.get(id as usize)
{
if *tok == "|" { if *tok == "|" {
out.push(' '); out.push(' ');
} else if !tok.starts_with('<') { } else if !tok.starts_with('<') {
+1 -4
View File
@@ -378,10 +378,7 @@ impl MetricsEngine {
"throughput".to_owned(), "throughput".to_owned(),
self.calculate_throughput(predictions.len()), self.calculate_throughput(predictions.len()),
); );
metrics.insert( metrics.insert("efficiency".to_owned(), self.calculate_efficiency(&metrics));
"efficiency".to_owned(),
self.calculate_efficiency(&metrics),
);
metrics.insert( metrics.insert(
"latency".to_owned(), "latency".to_owned(),
self.calculate_latency(predictions.len()), self.calculate_latency(predictions.len()),
@@ -332,8 +332,7 @@ impl MambaBlock {
z ^ (z >> 31) z ^ (z >> 31)
}; };
let in_proj = let in_proj = Tensor::randn_seeded(&[config.d_model, d_inner * 2], device, next_seed())?;
Tensor::randn_seeded(&[config.d_model, d_inner * 2], device, next_seed())?;
let conv1d_weight = let conv1d_weight =
Tensor::randn_seeded(&[d_inner, 1, config.d_conv], device, next_seed())?; Tensor::randn_seeded(&[d_inner, 1, config.d_conv], device, next_seed())?;
let conv1d_bias = if config.conv_bias { let conv1d_bias = if config.conv_bias {
@@ -341,11 +340,9 @@ impl MambaBlock {
} else { } else {
None None
}; };
let A_log = let A_log = Tensor::randn_seeded(&[d_inner, config.d_state], device, next_seed())?;
Tensor::randn_seeded(&[d_inner, config.d_state], device, next_seed())?;
let dt_proj = Tensor::randn_seeded(&[dt_rank, d_inner], device, next_seed())?; let dt_proj = Tensor::randn_seeded(&[dt_rank, d_inner], device, next_seed())?;
let out_proj = let out_proj = Tensor::randn_seeded(&[d_inner, config.d_model], device, next_seed())?;
Tensor::randn_seeded(&[d_inner, config.d_model], device, next_seed())?;
let selective_scan = SelectiveScan::new(d_inner, config.d_state); let selective_scan = SelectiveScan::new(d_inner, config.d_state);
@@ -409,8 +406,7 @@ impl MambaBlock {
.into() .into()
}) })
}; };
let assert_shape = let assert_shape = |t: &Tensor, key: &str, expected: &[usize]| -> Result<()> {
|t: &Tensor, key: &str, expected: &[usize]| -> Result<()> {
if t.shape().dims() != expected { if t.shape().dims() != expected {
return Err(anyhow::anyhow!( return Err(anyhow::anyhow!(
"MambaBlock::from_persistence_tensors: tensor `{key}` shape {:?} != expected {:?}", "MambaBlock::from_persistence_tensors: tensor `{key}` shape {:?} != expected {:?}",
@@ -426,7 +422,11 @@ impl MambaBlock {
assert_shape(&in_proj, "in_proj", &[config.d_model, d_inner * 2])?; assert_shape(&in_proj, "in_proj", &[config.d_model, d_inner * 2])?;
let conv1d_weight = take(&mut tensors, "conv1d_weight")?; let conv1d_weight = take(&mut tensors, "conv1d_weight")?;
assert_shape(&conv1d_weight, "conv1d_weight", &[d_inner, 1, config.d_conv])?; assert_shape(
&conv1d_weight,
"conv1d_weight",
&[d_inner, 1, config.d_conv],
)?;
let conv1d_bias = if config.conv_bias { let conv1d_bias = if config.conv_bias {
let b = take(&mut tensors, "conv1d_bias")?; let b = take(&mut tensors, "conv1d_bias")?;