rtx-csm: WavLM-Base+ SV scaffold for speaker similarity
Phase 5a — architectural skeleton for the microsoft/wavlm-base-plus-sv reference, drop-in replacement for the SpectralCentroidSimilarity weak baseline in speaker_sim.rs. Modules in src/wavlm_sv.rs (~600 LOC): - FeatureExtractor: 7-layer Conv1d, 320× downsample, GroupNorm at layer 0 (num_groups=num_channels=512), GELU activations. - FeatureProjection: LayerNorm + Linear 512→768. - PosConv: Conv1d(768, 768, k=128, groups=16, pad=64) + GELU; SamePad strips trailing frame for even kernel. - WavLmEncoderLayer: struct shape complete (Q/K/V/out projections, pre- attention LN, FFN intermediate/output, final LN, gru_rel_pos_const + gru_rel_pos_linear, optional rel_attn_embed at layer 0). forward() is a STUB; Phase 5b implements gated rel-pos attention. - Encoder: 12 stacked layers, returns Vec<Tensor> of 13 hidden states. - Tdnn: dilated unfold + Linear(in*kernel, out) — matches HF impl. - XVectorHead: softmax-weighted layer sum + projector 768→512 + 5 TDNN layers (kernels [5,3,3,1,1] dilations [1,2,3,1,1]) + statistics pool + 3000→512 embedding projection. - WavLmSv top-level + zero-mean unit-variance normalize + cosine similarity helper for verification scoring. 9 shape-correctness tests; 72 lib tests total green. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -0,0 +1,712 @@
|
||||
//! WavLM-Base+ Speaker Verification — Rust port (scaffold).
|
||||
//!
|
||||
//! Architectural port of `microsoft/wavlm-base-plus-sv` against candle 0.9.
|
||||
//! WavLM-Base+ is 12-layer / 768-dim transformer over 16 kHz waveforms;
|
||||
//! the `-sv` variant adds an X-vector head producing 512-d speaker
|
||||
//! embeddings. Cosine similarity between embeddings → speaker similarity
|
||||
//! score (drop-in replacement for the weak `SpectralCentroidSimilarity`
|
||||
//! in `speaker_sim.rs`).
|
||||
//!
|
||||
//! ## Pipeline (input → embedding)
|
||||
//!
|
||||
//! ```text
|
||||
//! waveform [B, T] (16 kHz)
|
||||
//! │
|
||||
//! ▼
|
||||
//! FeatureExtractor (7 conv layers, 320× downsample) → [B, 512, T/320]
|
||||
//! │
|
||||
//! transpose ▼
|
||||
//! FeatureProjection: LN → Linear 512→768 → [B, T/320, 768]
|
||||
//! │
|
||||
//! ▼
|
||||
//! PosConv (Conv1d 768→768, k=128, groups=16, weight_norm) + GELU
|
||||
//! │
|
||||
//! add ▼ residual
|
||||
//! LayerNorm → 12 × WavLMEncoderLayer (with rel_pos_bias) → 13 hidden states
|
||||
//! │
|
||||
//! ▼
|
||||
//! Weighted sum (softmax over 13 layer_weights) → [B, T/320, 768]
|
||||
//! │
|
||||
//! ▼
|
||||
//! X-Vector head:
|
||||
//! - projector Linear 768→512
|
||||
//! - 5 TDNN layers with kernels [5,3,3,1,1] dilations [1,2,3,1,1]
|
||||
//! - ReLU after each
|
||||
//! - Statistics pooling [mean | std] over time → [B, 3000]
|
||||
//! - feature_extractor Linear 3000→512 → [B, 512]
|
||||
//! ```
|
||||
//!
|
||||
//! ## Status: PARTIAL SCAFFOLD
|
||||
//!
|
||||
//! What this module ships:
|
||||
//! - All struct definitions for the full pipeline
|
||||
//! - `FeatureExtractor` (7-layer conv, GroupNorm at layer 0, GELU)
|
||||
//! - `FeatureProjection` (LayerNorm + Linear 512→768)
|
||||
//! - `XVectorHead` (5 TDNN + statistics pool + 3000→512 projection)
|
||||
//! - Stubs for `PosConv`, `WavLMEncoderLayer`, `Encoder`
|
||||
//! - Shape-correctness tests for the implemented blocks
|
||||
//!
|
||||
//! Deferred to Phase 5b:
|
||||
//! - Bucketed relative-position bias (T5-style, 320 buckets)
|
||||
//! - Gated relative-position bias (per-layer 1×12×1×1 const + 64→8 linear)
|
||||
//! - The full `WavLMEncoderLayer` forward (attention + FFN)
|
||||
//!
|
||||
//! Deferred to Phase 5c:
|
||||
//! - PyTorch `pytorch_model.bin` → safetensors conversion (handles weight_norm,
|
||||
//! `layer_weights` softmax-weights, TDNN `kernel.weight` reshape)
|
||||
//! - HF Hub asset resolution
|
||||
//!
|
||||
//! Deferred to Phase 5d:
|
||||
//! - Numerical parity vs the HF reference (cosine similarity within 1e-4 on
|
||||
//! a paired-utterance test set)
|
||||
|
||||
use crate::error::{CsmError, Result};
|
||||
use candle_core::{DType, Device, IndexOp, Module, Tensor, D};
|
||||
use candle_nn::{
|
||||
conv1d, group_norm, layer_norm, linear, ops, Activation, Conv1d, Conv1dConfig, GroupNorm,
|
||||
LayerNorm, LayerNormConfig, Linear, VarBuilder,
|
||||
};
|
||||
use std::path::Path;
|
||||
|
||||
pub const SAMPLE_RATE: u32 = 16_000;
|
||||
pub const HIDDEN_DIM: usize = 768;
|
||||
pub const NUM_HEADS: usize = 12;
|
||||
pub const HEAD_DIM: usize = HIDDEN_DIM / NUM_HEADS; // 64
|
||||
pub const FFN_DIM: usize = 3072;
|
||||
pub const NUM_LAYERS: usize = 12;
|
||||
pub const FEATURE_DIM: usize = 512;
|
||||
pub const HOP_LENGTH: usize = 320; // total CNN downsampling
|
||||
pub const REL_NUM_BUCKETS: usize = 320;
|
||||
pub const REL_MAX_DISTANCE: usize = 800;
|
||||
pub const EMBEDDING_DIM: usize = 512;
|
||||
pub const STAT_POOL_DIM: usize = 3000; // 1500 * 2 (mean+std)
|
||||
|
||||
// -- 1. Feature extractor (7-layer Conv1d, 320× downsampling) --------------
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct ConvLayer {
|
||||
conv: Conv1d,
|
||||
norm: Option<GroupNorm>,
|
||||
#[allow(dead_code)]
|
||||
kernel: usize,
|
||||
#[allow(dead_code)]
|
||||
stride: usize,
|
||||
}
|
||||
|
||||
impl ConvLayer {
|
||||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let h = xs.apply(&self.conv)?;
|
||||
let h = match self.norm.as_ref() {
|
||||
Some(gn) => h.apply(gn)?,
|
||||
None => h,
|
||||
};
|
||||
h.gelu()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FeatureExtractor {
|
||||
layers: Vec<ConvLayer>,
|
||||
}
|
||||
|
||||
impl FeatureExtractor {
|
||||
/// vb is rooted at `wavlm.feature_extractor`. Each layer's path is
|
||||
/// `conv_layers.{i}.conv.weight`. Layer 0 also has
|
||||
/// `conv_layers.0.layer_norm.{weight,bias}` (WavLM uses GroupNorm with
|
||||
/// num_groups = num_channels = 512, despite the parameter name).
|
||||
pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
|
||||
// (out_channels, kernel, stride). Input is mono so first in=1.
|
||||
let specs: [(usize, usize, usize); 7] = [
|
||||
(FEATURE_DIM, 10, 5),
|
||||
(FEATURE_DIM, 3, 2),
|
||||
(FEATURE_DIM, 3, 2),
|
||||
(FEATURE_DIM, 3, 2),
|
||||
(FEATURE_DIM, 3, 2),
|
||||
(FEATURE_DIM, 2, 2),
|
||||
(FEATURE_DIM, 2, 2),
|
||||
];
|
||||
let mut layers = Vec::with_capacity(7);
|
||||
let mut in_ch = 1usize;
|
||||
let vb_layers = vb.pp("conv_layers");
|
||||
for (i, (out_ch, k, s)) in specs.iter().enumerate() {
|
||||
let cfg = Conv1dConfig {
|
||||
stride: *s,
|
||||
..Default::default()
|
||||
};
|
||||
let conv = candle_nn::conv1d_no_bias(in_ch, *out_ch, *k, cfg, vb_layers.pp(i.to_string()).pp("conv"))?;
|
||||
let norm = if i == 0 {
|
||||
// WavLM/wav2vec2 GroupNorm layer 0: num_groups = num_channels.
|
||||
Some(group_norm(*out_ch, *out_ch, 1e-5, vb_layers.pp(i.to_string()).pp("layer_norm"))?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
layers.push(ConvLayer {
|
||||
conv,
|
||||
norm,
|
||||
kernel: *k,
|
||||
stride: *s,
|
||||
});
|
||||
in_ch = *out_ch;
|
||||
}
|
||||
Ok(Self { layers })
|
||||
}
|
||||
}
|
||||
|
||||
impl Module for FeatureExtractor {
|
||||
/// Input `(B, 1, T)`; output `(B, FEATURE_DIM=512, T/320)`.
|
||||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let mut h = xs.clone();
|
||||
for layer in &self.layers {
|
||||
h = layer.forward(&h)?;
|
||||
}
|
||||
Ok(h)
|
||||
}
|
||||
}
|
||||
|
||||
// -- 2. Feature projection (LN + Linear 512 → 768) ------------------------
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FeatureProjection {
|
||||
layer_norm: LayerNorm,
|
||||
projection: Linear,
|
||||
}
|
||||
|
||||
impl FeatureProjection {
|
||||
/// vb is rooted at `wavlm.feature_projection`.
|
||||
pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
|
||||
let layer_norm = layer_norm(
|
||||
FEATURE_DIM,
|
||||
LayerNormConfig {
|
||||
eps: 1e-5,
|
||||
..Default::default()
|
||||
},
|
||||
vb.pp("layer_norm"),
|
||||
)?;
|
||||
let projection = linear(FEATURE_DIM, HIDDEN_DIM, vb.pp("projection"))?;
|
||||
Ok(Self {
|
||||
layer_norm,
|
||||
projection,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Module for FeatureProjection {
|
||||
/// Input `(B, T/320, 512)`; output `(B, T/320, 768)`.
|
||||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
xs.apply(&self.layer_norm)?.apply(&self.projection)
|
||||
}
|
||||
}
|
||||
|
||||
// -- 3. Positional convolution (depthwise Conv1d with weight_norm) --------
|
||||
|
||||
/// `Conv1d(768, 768, kernel=128, padding=64, groups=16)` with `weight_norm`.
|
||||
/// Output is GELU-activated and added to the input as a residual.
|
||||
///
|
||||
/// Note: weight_norm storage requires a converter pass (g/v → weight),
|
||||
/// matching the AudioSeal converter pattern. Currently builds a plain
|
||||
/// Conv1d; safetensors loader must merge before reading.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PosConv {
|
||||
conv: Conv1d,
|
||||
kernel: usize,
|
||||
}
|
||||
|
||||
impl PosConv {
|
||||
pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
|
||||
let cfg = Conv1dConfig {
|
||||
padding: 64,
|
||||
groups: 16,
|
||||
..Default::default()
|
||||
};
|
||||
let conv = conv1d(HIDDEN_DIM, HIDDEN_DIM, 128, cfg, vb.pp("pos_conv_embed.conv"))?;
|
||||
Ok(Self { conv, kernel: 128 })
|
||||
}
|
||||
}
|
||||
|
||||
impl Module for PosConv {
|
||||
/// Input/output: `(B, T/320, 768)` (treats time axis -2 as conv axis).
|
||||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
// (B, T, C) → (B, C, T) → conv → (B, C, T) → drop trailing pad → (B, T, C)
|
||||
let in_len = xs.dim(D::Minus2)?;
|
||||
let xs_bct = xs.transpose(1, 2)?.contiguous()?;
|
||||
let h = xs_bct.apply(&self.conv)?;
|
||||
// SamePad strips 1 trailing frame because kernel=128 is even.
|
||||
let out_len = h.dim(D::Minus1)?;
|
||||
let h = if self.kernel % 2 == 0 && out_len > in_len {
|
||||
h.narrow(D::Minus1, 0, in_len)?
|
||||
} else {
|
||||
h
|
||||
};
|
||||
let h = h.gelu()?;
|
||||
// Back to (B, T, C).
|
||||
h.transpose(1, 2)?.contiguous()
|
||||
}
|
||||
}
|
||||
|
||||
// -- 4. Transformer encoder layer (STUB — Phase 5b) -----------------------
|
||||
|
||||
/// One WavLM encoder layer. Post-norm: residual+attn → LN → FFN-residual → final-LN.
|
||||
///
|
||||
/// **Stub**: forward returns input unchanged. Phase 5b will implement
|
||||
/// gated relative-position bias attention + 768→3072→768 FFN.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WavLmEncoderLayer {
|
||||
// Phase 5b — these will be populated:
|
||||
pub q_proj: Linear,
|
||||
pub k_proj: Linear,
|
||||
pub v_proj: Linear,
|
||||
pub out_proj: Linear,
|
||||
pub attn_norm: LayerNorm,
|
||||
pub fc1: Linear,
|
||||
pub fc2: Linear,
|
||||
pub final_norm: LayerNorm,
|
||||
/// Per-layer gated rel-pos bias parameters (1, 12, 1, 1).
|
||||
pub gru_rel_pos_const: Tensor,
|
||||
/// Linear from head_dim=64 → 8 for the 2-gate split.
|
||||
pub gru_rel_pos_linear: Linear,
|
||||
/// Only set for layer 0 — bucketed embedding (320, 12).
|
||||
pub rel_attn_embed: Option<candle_nn::Embedding>,
|
||||
}
|
||||
|
||||
impl WavLmEncoderLayer {
|
||||
pub fn new(layer_idx: usize, vb: VarBuilder) -> candle_core::Result<Self> {
|
||||
let attn = vb.pp("attention");
|
||||
let q_proj = linear(HIDDEN_DIM, HIDDEN_DIM, attn.pp("q_proj"))?;
|
||||
let k_proj = linear(HIDDEN_DIM, HIDDEN_DIM, attn.pp("k_proj"))?;
|
||||
let v_proj = linear(HIDDEN_DIM, HIDDEN_DIM, attn.pp("v_proj"))?;
|
||||
let out_proj = linear(HIDDEN_DIM, HIDDEN_DIM, attn.pp("out_proj"))?;
|
||||
let attn_norm = layer_norm(HIDDEN_DIM, 1e-5, vb.pp("layer_norm"))?;
|
||||
let ff = vb.pp("feed_forward");
|
||||
let fc1 = linear(HIDDEN_DIM, FFN_DIM, ff.pp("intermediate_dense"))?;
|
||||
let fc2 = linear(FFN_DIM, HIDDEN_DIM, ff.pp("output_dense"))?;
|
||||
let final_norm = layer_norm(HIDDEN_DIM, 1e-5, vb.pp("final_layer_norm"))?;
|
||||
let gru_rel_pos_const = attn
|
||||
.pp("gru_rel_pos_const")
|
||||
.get((1, NUM_HEADS, 1, 1), "")
|
||||
.or_else(|_| {
|
||||
Tensor::zeros((1, NUM_HEADS, 1, 1), vb.dtype(), vb.device())
|
||||
})?;
|
||||
let gru_rel_pos_linear = linear(HEAD_DIM, 8, attn.pp("gru_rel_pos_linear"))?;
|
||||
let rel_attn_embed = if layer_idx == 0 {
|
||||
Some(candle_nn::embedding(
|
||||
REL_NUM_BUCKETS,
|
||||
NUM_HEADS,
|
||||
attn.pp("rel_attn_embed"),
|
||||
)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
Ok(Self {
|
||||
q_proj,
|
||||
k_proj,
|
||||
v_proj,
|
||||
out_proj,
|
||||
attn_norm,
|
||||
fc1,
|
||||
fc2,
|
||||
final_norm,
|
||||
gru_rel_pos_const,
|
||||
gru_rel_pos_linear,
|
||||
rel_attn_embed,
|
||||
})
|
||||
}
|
||||
|
||||
/// **STUB**: returns input unchanged. Phase 5b implements full
|
||||
/// attention + FFN. Shape `(B, T, 768)` in/out.
|
||||
pub fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
Ok(xs.clone())
|
||||
}
|
||||
}
|
||||
|
||||
// -- 5. Encoder (12 stacked layers, returns 13 hidden states) -------------
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Encoder {
|
||||
pos_conv: PosConv,
|
||||
layer_norm: LayerNorm,
|
||||
layers: Vec<WavLmEncoderLayer>,
|
||||
}
|
||||
|
||||
impl Encoder {
|
||||
pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
|
||||
let pos_conv = PosConv::new(vb.clone())?;
|
||||
let layer_norm = layer_norm(HIDDEN_DIM, 1e-5, vb.pp("layer_norm"))?;
|
||||
let mut layers = Vec::with_capacity(NUM_LAYERS);
|
||||
let vb_layers = vb.pp("layers");
|
||||
for i in 0..NUM_LAYERS {
|
||||
layers.push(WavLmEncoderLayer::new(i, vb_layers.pp(i.to_string()))?);
|
||||
}
|
||||
Ok(Self {
|
||||
pos_conv,
|
||||
layer_norm,
|
||||
layers,
|
||||
})
|
||||
}
|
||||
|
||||
/// Input `(B, T, 768)`; output is `(B, T, 768)` per layer + initial,
|
||||
/// returned as a `Vec<Tensor>` of length `NUM_LAYERS + 1 = 13`.
|
||||
pub fn forward_all_layers(&self, xs: &Tensor) -> candle_core::Result<Vec<Tensor>> {
|
||||
let pos = self.pos_conv.forward(xs)?;
|
||||
let mut h = (xs + &pos)?;
|
||||
h = h.apply(&self.layer_norm)?;
|
||||
let mut hidden_states = Vec::with_capacity(NUM_LAYERS + 1);
|
||||
hidden_states.push(h.clone());
|
||||
for layer in &self.layers {
|
||||
h = layer.forward(&h)?;
|
||||
hidden_states.push(h.clone());
|
||||
}
|
||||
Ok(hidden_states)
|
||||
}
|
||||
}
|
||||
|
||||
// -- 6. X-Vector head (TDNN + stats pool + projection) --------------------
|
||||
|
||||
/// One TDNN layer: implemented per the HF reference as `Linear(in*kernel, out)`
|
||||
/// fed via a manually-strided unfold over the time axis with `dilation`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Tdnn {
|
||||
kernel_linear: Linear,
|
||||
#[allow(dead_code)]
|
||||
in_dim: usize,
|
||||
#[allow(dead_code)]
|
||||
out_dim: usize,
|
||||
kernel: usize,
|
||||
dilation: usize,
|
||||
}
|
||||
|
||||
impl Tdnn {
|
||||
pub fn new(
|
||||
in_dim: usize,
|
||||
out_dim: usize,
|
||||
kernel: usize,
|
||||
dilation: usize,
|
||||
vb: VarBuilder,
|
||||
) -> candle_core::Result<Self> {
|
||||
let kernel_linear = linear(in_dim * kernel, out_dim, vb.pp("kernel"))?;
|
||||
Ok(Self {
|
||||
kernel_linear,
|
||||
in_dim,
|
||||
out_dim,
|
||||
kernel,
|
||||
dilation,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Module for Tdnn {
|
||||
/// Input `(B, T, in_dim)`; output `(B, T_out, out_dim)` where
|
||||
/// `T_out = T - (kernel - 1) * dilation` (no padding, valid-only).
|
||||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let (b, t, c) = xs.dims3()?;
|
||||
let span = (self.kernel - 1) * self.dilation;
|
||||
if t <= span {
|
||||
return Err(candle_core::Error::Msg(format!(
|
||||
"TDNN input length {t} too short for kernel {} dilation {}",
|
||||
self.kernel, self.dilation
|
||||
)));
|
||||
}
|
||||
let t_out = t - span;
|
||||
// Build a (B, T_out, kernel*in_dim) tensor by gathering kernel
|
||||
// dilated samples per output frame.
|
||||
let mut frames: Vec<Tensor> = Vec::with_capacity(self.kernel);
|
||||
for k in 0..self.kernel {
|
||||
let offset = k * self.dilation;
|
||||
// (B, T_out, in_dim)
|
||||
let slice = xs.narrow(1, offset, t_out)?;
|
||||
frames.push(slice);
|
||||
}
|
||||
let stacked = Tensor::cat(&frames, 2)?; // (B, T_out, kernel * in_dim)
|
||||
debug_assert_eq!(stacked.dims(), &[b, t_out, self.kernel * c]);
|
||||
// Linear → (B, T_out, out_dim) → ReLU
|
||||
stacked.apply(&self.kernel_linear)?.relu()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct XVectorHead {
|
||||
/// Softmax over `layer_weights` of length `NUM_LAYERS + 1 = 13`.
|
||||
layer_weights: Tensor,
|
||||
projector: Linear,
|
||||
tdnn: Vec<Tdnn>,
|
||||
/// Linear 3000 → 512 — the "feature_extractor" key in the HF state_dict.
|
||||
/// We rename to avoid collision with the WavLM CNN feature extractor.
|
||||
embedding_proj: Linear,
|
||||
}
|
||||
|
||||
impl XVectorHead {
|
||||
/// vb is rooted at the *top* of the WavLMForXVector state_dict (so we
|
||||
/// read `layer_weights`, `projector`, `tdnn.{0..4}`, `feature_extractor`).
|
||||
pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
|
||||
let layer_weights = vb.pp("layer_weights").get(NUM_LAYERS + 1, "")?;
|
||||
let projector = linear(HIDDEN_DIM, FEATURE_DIM, vb.pp("projector"))?;
|
||||
// TDNN specs: (in, out, kernel, dilation)
|
||||
let tdnn_specs: [(usize, usize, usize, usize); 5] = [
|
||||
(FEATURE_DIM, 512, 5, 1),
|
||||
(512, 512, 3, 2),
|
||||
(512, 512, 3, 3),
|
||||
(512, 512, 1, 1),
|
||||
(512, 1500, 1, 1),
|
||||
];
|
||||
let mut tdnn = Vec::with_capacity(5);
|
||||
let vb_tdnn = vb.pp("tdnn");
|
||||
for (i, (in_dim, out_dim, k, d)) in tdnn_specs.iter().enumerate() {
|
||||
tdnn.push(Tdnn::new(*in_dim, *out_dim, *k, *d, vb_tdnn.pp(i.to_string()))?);
|
||||
}
|
||||
let embedding_proj = linear(STAT_POOL_DIM, EMBEDDING_DIM, vb.pp("feature_extractor"))?;
|
||||
Ok(Self {
|
||||
layer_weights,
|
||||
projector,
|
||||
tdnn,
|
||||
embedding_proj,
|
||||
})
|
||||
}
|
||||
|
||||
/// Input: `Vec<Tensor>` of length 13, each `(B, T, 768)`. Output: `(B, 512)`.
|
||||
pub fn forward(&self, hidden_states: &[Tensor]) -> candle_core::Result<Tensor> {
|
||||
if hidden_states.len() != NUM_LAYERS + 1 {
|
||||
return Err(candle_core::Error::Msg(format!(
|
||||
"expected {} hidden states, got {}",
|
||||
NUM_LAYERS + 1,
|
||||
hidden_states.len()
|
||||
)));
|
||||
}
|
||||
// Softmax-weighted sum over layers.
|
||||
let weights = ops::softmax(&self.layer_weights, 0)?; // (13,)
|
||||
let weights = weights.to_dtype(hidden_states[0].dtype())?;
|
||||
let mut sum: Option<Tensor> = None;
|
||||
for (i, h) in hidden_states.iter().enumerate() {
|
||||
let w = weights.i(i)?; // scalar
|
||||
let scaled = h.broadcast_mul(&w.reshape((1, 1, 1))?)?;
|
||||
sum = Some(match sum {
|
||||
Some(prev) => (prev + scaled)?,
|
||||
None => scaled,
|
||||
});
|
||||
}
|
||||
let h = sum.expect("at least one hidden state");
|
||||
|
||||
// Projector 768 → 512.
|
||||
let mut h = h.apply(&self.projector)?;
|
||||
// 5 TDNN layers (reduce time dim each time).
|
||||
for layer in &self.tdnn {
|
||||
h = layer.forward(&h)?;
|
||||
}
|
||||
// Statistics pooling: mean & std over time axis (dim=1).
|
||||
let mean = h.mean(1)?; // (B, 1500)
|
||||
let var = h.var(1)?;
|
||||
let std = (var + 1e-9)?.sqrt()?;
|
||||
let stat = Tensor::cat(&[&mean, &std], D::Minus1)?; // (B, 3000)
|
||||
// Final projection 3000 → 512.
|
||||
stat.apply(&self.embedding_proj)
|
||||
}
|
||||
}
|
||||
|
||||
// -- 7. Top-level model ----------------------------------------------------
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WavLmSv {
|
||||
feature_extractor: FeatureExtractor,
|
||||
feature_projection: FeatureProjection,
|
||||
encoder: Encoder,
|
||||
head: XVectorHead,
|
||||
}
|
||||
|
||||
impl WavLmSv {
|
||||
pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
|
||||
let backbone = vb.pp("wavlm");
|
||||
let feature_extractor = FeatureExtractor::new(backbone.pp("feature_extractor"))?;
|
||||
let feature_projection = FeatureProjection::new(backbone.pp("feature_projection"))?;
|
||||
let encoder = Encoder::new(backbone.pp("encoder"))?;
|
||||
let head = XVectorHead::new(vb.clone())?;
|
||||
Ok(Self {
|
||||
feature_extractor,
|
||||
feature_projection,
|
||||
encoder,
|
||||
head,
|
||||
})
|
||||
}
|
||||
|
||||
/// Per-utterance zero-mean/unit-variance normalization (matches
|
||||
/// HF Wav2Vec2FeatureExtractor with do_normalize=True).
|
||||
pub fn normalize_waveform(samples: &[f32]) -> Vec<f32> {
|
||||
if samples.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let n = samples.len() as f32;
|
||||
let mean = samples.iter().sum::<f32>() / n;
|
||||
let var = samples.iter().map(|s| (s - mean).powi(2)).sum::<f32>() / n;
|
||||
let std = (var + 1e-7).sqrt();
|
||||
samples.iter().map(|s| (s - mean) / std).collect()
|
||||
}
|
||||
|
||||
/// Compute a `(B, 512)` speaker embedding from a 16 kHz mono waveform
|
||||
/// tensor `(B, 1, T)`. Caller is responsible for normalization.
|
||||
pub fn embed(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let features = self.feature_extractor.forward(xs)?; // (B, 512, T/320)
|
||||
// (B, 512, T') → (B, T', 512)
|
||||
let features = features.transpose(1, 2)?.contiguous()?;
|
||||
let projected = self.feature_projection.forward(&features)?; // (B, T', 768)
|
||||
let hidden_states = self.encoder.forward_all_layers(&projected)?; // 13 × (B, T', 768)
|
||||
self.head.forward(&hidden_states)
|
||||
}
|
||||
|
||||
/// Convenience: embed a `Vec<f32>` and return the embedding as a Vec<f32>.
|
||||
pub fn embed_samples(&self, samples: &[f32], device: &Device) -> Result<Vec<f32>> {
|
||||
let normalized = Self::normalize_waveform(samples);
|
||||
let xs = Tensor::from_slice(&normalized, (1, 1, normalized.len()), device)
|
||||
.map_err(|e| CsmError::Config(format!("embed: tensor: {e}")))?;
|
||||
let emb = self
|
||||
.embed(&xs)
|
||||
.map_err(|e| CsmError::Config(format!("embed: forward: {e}")))?;
|
||||
emb.i(0)?
|
||||
.to_dtype(DType::F32)
|
||||
.and_then(|t| t.to_vec1::<f32>())
|
||||
.map_err(|e| CsmError::Config(format!("embed: to_vec: {e}")))
|
||||
}
|
||||
|
||||
/// Cosine similarity between two embeddings, both expected as `(D,)`
|
||||
/// or `(B, D)` tensors of compatible shape.
|
||||
pub fn cosine_similarity(a: &Tensor, b: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let a_norm = a.broadcast_div(
|
||||
&(a.sqr()?.sum_keepdim(D::Minus1)? + 1e-9)?
|
||||
.sqrt()?,
|
||||
)?;
|
||||
let b_norm = b.broadcast_div(
|
||||
&(b.sqr()?.sum_keepdim(D::Minus1)? + 1e-9)?
|
||||
.sqrt()?,
|
||||
)?;
|
||||
(a_norm * b_norm)?.sum(D::Minus1)
|
||||
}
|
||||
}
|
||||
|
||||
/// Stub loader. Phase 5c will implement `pytorch_model.bin` → safetensors
|
||||
/// conversion (weight_norm merge for pos_conv + TDNN kernel reshape) and
|
||||
/// HF Hub asset resolution under `microsoft/wavlm-base-plus-sv`.
|
||||
pub fn load_from_safetensors<P: AsRef<Path>>(
|
||||
safetensors: P,
|
||||
device: &Device,
|
||||
) -> Result<WavLmSv> {
|
||||
let vb = unsafe {
|
||||
candle_nn::VarBuilder::from_mmaped_safetensors(
|
||||
&[safetensors.as_ref()],
|
||||
DType::F32,
|
||||
device,
|
||||
)
|
||||
}
|
||||
.map_err(|e| CsmError::Config(format!("opening WavLM safetensors: {e}")))?;
|
||||
WavLmSv::new(vb).map_err(|e| CsmError::Config(format!("WavLmSv::new: {e}")))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_nn::VarMap;
|
||||
|
||||
fn random_vb(device: &Device) -> (VarMap, VarBuilder<'static>) {
|
||||
let vm = VarMap::new();
|
||||
let vb = VarBuilder::from_varmap(&vm, DType::F32, device);
|
||||
(vm, vb)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feature_extractor_downsamples_by_320() {
|
||||
let device = Device::Cpu;
|
||||
let (_vm, vb) = random_vb(&device);
|
||||
let fe = FeatureExtractor::new(vb).unwrap();
|
||||
// 16000 samples = 1s @ 16 kHz → ~50 frames.
|
||||
let xs = Tensor::randn(0f32, 1f32, (1, 1, 16000), &device).unwrap();
|
||||
let out = fe.forward(&xs).unwrap();
|
||||
let frames = out.dim(D::Minus1).unwrap();
|
||||
assert_eq!(out.dim(0).unwrap(), 1);
|
||||
assert_eq!(out.dim(1).unwrap(), FEATURE_DIM);
|
||||
assert!(
|
||||
(frames as i64 - 49).abs() <= 2,
|
||||
"expected ~49 frames, got {frames}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feature_projection_lifts_512_to_768() {
|
||||
let device = Device::Cpu;
|
||||
let (_vm, vb) = random_vb(&device);
|
||||
let fp = FeatureProjection::new(vb).unwrap();
|
||||
let xs = Tensor::randn(0f32, 1f32, (2, 49, FEATURE_DIM), &device).unwrap();
|
||||
let out = fp.forward(&xs).unwrap();
|
||||
assert_eq!(out.dims(), &[2, 49, HIDDEN_DIM]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pos_conv_preserves_shape() {
|
||||
let device = Device::Cpu;
|
||||
let (_vm, vb) = random_vb(&device);
|
||||
let pc = PosConv::new(vb).unwrap();
|
||||
let xs = Tensor::randn(0f32, 1f32, (1, 49, HIDDEN_DIM), &device).unwrap();
|
||||
let out = pc.forward(&xs).unwrap();
|
||||
assert_eq!(out.dims(), &[1, 49, HIDDEN_DIM]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tdnn_reduces_time_axis() {
|
||||
let device = Device::Cpu;
|
||||
let (_vm, vb) = random_vb(&device);
|
||||
// kernel=5, dilation=1 → T_out = T - 4
|
||||
let t = Tdnn::new(FEATURE_DIM, 512, 5, 1, vb.clone()).unwrap();
|
||||
let xs = Tensor::randn(0f32, 1f32, (1, 49, FEATURE_DIM), &device).unwrap();
|
||||
let out = t.forward(&xs).unwrap();
|
||||
assert_eq!(out.dims(), &[1, 45, 512]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tdnn_dilated_reduces_correctly() {
|
||||
let device = Device::Cpu;
|
||||
let (_vm, vb) = random_vb(&device);
|
||||
// kernel=3, dilation=2 → span = 4, T_out = T - 4
|
||||
let t = Tdnn::new(512, 512, 3, 2, vb).unwrap();
|
||||
let xs = Tensor::randn(0f32, 1f32, (1, 30, 512), &device).unwrap();
|
||||
let out = t.forward(&xs).unwrap();
|
||||
assert_eq!(out.dims(), &[1, 26, 512]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn xvector_head_produces_512d_embedding() {
|
||||
let device = Device::Cpu;
|
||||
let (_vm, vb) = random_vb(&device);
|
||||
let head = XVectorHead::new(vb).unwrap();
|
||||
// T must be at least 5+4+4 = 13 frames after the TDNN cascade.
|
||||
// Cascade reductions: 4, 4, 4, 0, 0 = 12 frames lost. Need T >= 13.
|
||||
let t = 49usize;
|
||||
let hidden_states: Vec<Tensor> = (0..NUM_LAYERS + 1)
|
||||
.map(|_| Tensor::randn(0f32, 1f32, (1, t, HIDDEN_DIM), &device).unwrap())
|
||||
.collect();
|
||||
let emb = head.forward(&hidden_states).unwrap();
|
||||
assert_eq!(emb.dims(), &[1, EMBEDDING_DIM]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn end_to_end_forward_smoke() {
|
||||
let device = Device::Cpu;
|
||||
let (_vm, vb) = random_vb(&device);
|
||||
let model = WavLmSv::new(vb).unwrap();
|
||||
// 1s @ 16 kHz = 16000 samples.
|
||||
let xs = Tensor::randn(0f32, 1f32, (1, 1, 16000), &device).unwrap();
|
||||
let emb = model.embed(&xs).unwrap();
|
||||
assert_eq!(emb.dims(), &[1, EMBEDDING_DIM]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cosine_similarity_self_is_one() {
|
||||
let device = Device::Cpu;
|
||||
let v = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0], (4,), &device).unwrap();
|
||||
let s = WavLmSv::cosine_similarity(&v, &v).unwrap();
|
||||
let v: f32 = s.to_dtype(DType::F32).unwrap().to_scalar().unwrap();
|
||||
assert!((v - 1.0).abs() < 1e-5, "expected 1.0, got {v}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_waveform_zero_mean() {
|
||||
let s = WavLmSv::normalize_waveform(&[1.0, 2.0, 3.0, 4.0, 5.0]);
|
||||
let mean: f32 = s.iter().sum::<f32>() / s.len() as f32;
|
||||
assert!(mean.abs() < 1e-5);
|
||||
let var = s.iter().map(|x| x * x).sum::<f32>() / s.len() as f32;
|
||||
assert!((var - 1.0).abs() < 1e-3, "expected unit variance, got {var}");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user