osobhandClaude Opus 4.7
a5cedfb46a
8-gen bench (4 emotions × 2 corpora) at seed=42 against firdhokk Whisper-LV3: target RAVDESS CREMA-D happy happy (0.999) ✓ happy (0.999) ✓ angry neutral (0.92) sad (0.99) fearful happy (0.998) fearful (0.984) ✓ sad angry (0.99) fearful (0.99) CREMA-D 2/4 vs RAVDESS 1/4. Larger / more naturalistic corpus produces more class-pure fearful direction. Neither corpus solves angry or sad — recipe shifts into 'vague expressivity' rather than class-specific corners. Practical: prefer CREMA-D when available; A/B both per emotion if class precision matters. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
911 lines
34 KiB
Rust
911 lines
34 KiB
Rust
//! 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::{D, DType, Device, IndexOp, Module, Tensor};
|
||
use candle_nn::{
|
||
Conv1d, Conv1dConfig, GroupNorm, LayerNorm, LayerNormConfig, Linear, VarBuilder, conv1d,
|
||
group_norm, layer_norm, linear, ops,
|
||
};
|
||
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_erf()
|
||
}
|
||
}
|
||
|
||
#[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.is_multiple_of(2) && out_len > in_len {
|
||
h.narrow(D::Minus1, 0, in_len)?
|
||
} else {
|
||
h
|
||
};
|
||
let h = h.gelu_erf()?;
|
||
// Back to (B, T, C).
|
||
h.transpose(1, 2)?.contiguous()
|
||
}
|
||
}
|
||
|
||
// -- 4. Transformer encoder layer with gated relative-position attention --
|
||
|
||
/// T5-style relative-position bucket. Maps a relative offset (k - q) into
|
||
/// `[0, num_buckets)`. Half the buckets cover negative offsets, half cover
|
||
/// positive; within each half, the first `max_exact = num_buckets/4` are
|
||
/// linear, the rest log-spaced up to `max_distance`.
|
||
fn relative_position_bucket(rel_pos: i64, num_buckets: usize, max_distance: usize) -> u32 {
|
||
let half = num_buckets / 2;
|
||
let mut bucket = if rel_pos > 0 { half } else { 0 };
|
||
let abs_pos = rel_pos.unsigned_abs() as usize;
|
||
let max_exact = half / 2;
|
||
if abs_pos < max_exact {
|
||
bucket += abs_pos;
|
||
} else {
|
||
let log_ratio = (abs_pos as f64 / max_exact as f64).ln();
|
||
let log_factor = (max_distance as f64 / max_exact as f64).ln();
|
||
let log_bucket = (log_ratio / log_factor) * (half - max_exact) as f64;
|
||
let large = max_exact + log_bucket as usize;
|
||
bucket += large.min(half - 1);
|
||
}
|
||
bucket as u32
|
||
}
|
||
|
||
/// One WavLM encoder layer. Post-norm: residual+attn → LN → FFN-residual → final-LN.
|
||
#[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.get((1, NUM_HEADS, 1, 1), "gru_rel_pos_const")?;
|
||
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,
|
||
})
|
||
}
|
||
|
||
/// Compute the bucketed relative-position bias `(num_heads, T, T)` for
|
||
/// a given query length. Only callable on layer 0 (the layer that owns
|
||
/// `rel_attn_embed`). Result is reused by the rest of the stack.
|
||
pub fn compute_position_bias(&self, t: usize) -> candle_core::Result<Tensor> {
|
||
let embed = self
|
||
.rel_attn_embed
|
||
.as_ref()
|
||
.ok_or_else(|| candle_core::Error::Msg(
|
||
"compute_position_bias called on layer without rel_attn_embed (only layer 0 has one)".into(),
|
||
))?;
|
||
// Build (T, T) bucket index tensor on the host.
|
||
let mut buckets = Vec::with_capacity(t * t);
|
||
for q in 0..t {
|
||
for k in 0..t {
|
||
let rel = k as i64 - q as i64;
|
||
buckets.push(relative_position_bucket(
|
||
rel,
|
||
REL_NUM_BUCKETS,
|
||
REL_MAX_DISTANCE,
|
||
));
|
||
}
|
||
}
|
||
let device = embed.embeddings().device();
|
||
let idx = Tensor::from_vec(buckets, (t, t), device)?;
|
||
// Embedding lookup: (T, T) → (T, T, num_heads). Permute to
|
||
// (num_heads, T, T) to match the HF compute_bias output.
|
||
let values = embed.forward(&idx)?;
|
||
values.permute((2, 0, 1))?.contiguous()
|
||
}
|
||
|
||
/// Compute the gated bias to add to attention scores.
|
||
/// `position_bias`: `(num_heads, T, T)`; `xs`: `(B, T, embed_dim)`.
|
||
/// Output: `(B, num_heads, T, T)` — the per-batch, per-head gated bias.
|
||
fn gated_position_bias(
|
||
&self,
|
||
position_bias: &Tensor,
|
||
xs: &Tensor,
|
||
) -> candle_core::Result<Tensor> {
|
||
let (b, t, _) = xs.dims3()?;
|
||
// (B, T, embed_dim) → (B, T, num_heads, head_dim) → (B, num_heads, T, head_dim)
|
||
let h = xs
|
||
.reshape((b, t, NUM_HEADS, HEAD_DIM))?
|
||
.permute((0, 2, 1, 3))?
|
||
.contiguous()?;
|
||
// Linear(head_dim → 8) → (B, num_heads, T, 8)
|
||
let proj = h.apply(&self.gru_rel_pos_linear)?;
|
||
// (B, num_heads, T, 2, 4) → sum(-1) → (B, num_heads, T, 2)
|
||
let proj = proj.reshape((b, NUM_HEADS, t, 2, 4))?.sum(D::Minus1)?;
|
||
let gates = ops::sigmoid(&proj)?; // (B, num_heads, T, 2)
|
||
let gate_a = gates.narrow(D::Minus1, 0, 1)?; // (B, num_heads, T, 1)
|
||
let gate_b = gates.narrow(D::Minus1, 1, 1)?; // (B, num_heads, T, 1)
|
||
// gate_output = gate_a * (gate_b * const - 1) + 2
|
||
let const_g = self
|
||
.gru_rel_pos_const
|
||
.broadcast_as((b, NUM_HEADS, t, 1))?
|
||
.to_dtype(gate_b.dtype())?;
|
||
let inner = ((gate_b * const_g)? - 1.0f64)?;
|
||
let gate_out = ((gate_a * inner)? + 2.0f64)?; // (B, num_heads, T, 1)
|
||
// Broadcast position_bias (num_heads, T, T) → (1, num_heads, T, T) → (B, num_heads, T, T)
|
||
let bias = position_bias
|
||
.unsqueeze(0)?
|
||
.broadcast_as((b, NUM_HEADS, t, t))?
|
||
.to_dtype(gate_out.dtype())?;
|
||
gate_out.broadcast_mul(&bias)
|
||
}
|
||
|
||
/// Multi-head self-attention with the gated bias added to the score
|
||
/// matrix before softmax.
|
||
fn attention(&self, xs: &Tensor, gated_bias: &Tensor) -> candle_core::Result<Tensor> {
|
||
let (b, t, _) = xs.dims3()?;
|
||
let split_heads = |proj: Tensor| -> candle_core::Result<Tensor> {
|
||
proj.reshape((b, t, NUM_HEADS, HEAD_DIM))?
|
||
.permute((0, 2, 1, 3))?
|
||
.contiguous()
|
||
};
|
||
let q = split_heads(xs.apply(&self.q_proj)?)?;
|
||
let k = split_heads(xs.apply(&self.k_proj)?)?;
|
||
let v = split_heads(xs.apply(&self.v_proj)?)?;
|
||
let scale = (HEAD_DIM as f64).powf(-0.5);
|
||
let scores = (q.matmul(&k.transpose(2, 3)?.contiguous()?)? * scale)?;
|
||
let scores = (scores + gated_bias)?;
|
||
let attn = ops::softmax_last_dim(&scores)?;
|
||
let out = attn.matmul(&v)?; // (B, num_heads, T, head_dim)
|
||
let out = out
|
||
.permute((0, 2, 1, 3))?
|
||
.contiguous()?
|
||
.reshape((b, t, HIDDEN_DIM))?;
|
||
out.apply(&self.out_proj)
|
||
}
|
||
|
||
/// Forward pass. Returns `(hidden_states, position_bias)` so the
|
||
/// encoder can thread the bias through subsequent layers.
|
||
///
|
||
/// `in_bias` is `Some` for layers 1..N — the bias computed by layer 0.
|
||
/// `None` is acceptable on layer 0 (we'll compute it locally) and an
|
||
/// error on any other layer (caller's bug).
|
||
pub fn forward_with_bias(
|
||
&self,
|
||
xs: &Tensor,
|
||
in_bias: Option<&Tensor>,
|
||
) -> candle_core::Result<(Tensor, Tensor)> {
|
||
let owned;
|
||
let bias = match in_bias {
|
||
Some(b) => b,
|
||
None => {
|
||
let t = xs.dim(D::Minus2)?;
|
||
owned = self.compute_position_bias(t)?;
|
||
&owned
|
||
}
|
||
};
|
||
let gated = self.gated_position_bias(bias, xs)?;
|
||
let attn_out = self.attention(xs, &gated)?;
|
||
let h = (xs + attn_out)?;
|
||
let h = h.apply(&self.attn_norm)?;
|
||
// FFN: 768 → 3072 → GELU → 768
|
||
let ffn = h.apply(&self.fc1)?.gelu_erf()?.apply(&self.fc2)?;
|
||
let h = (h + ffn)?;
|
||
let h = h.apply(&self.final_norm)?;
|
||
Ok((h, bias.clone()))
|
||
}
|
||
|
||
/// Convenience used by the smoke test where the caller doesn't carry a
|
||
/// bias — equivalent to layer 0's `forward_with_bias(xs, None)`.
|
||
pub fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
let (out, _bias) = self.forward_with_bias(xs, None)?;
|
||
Ok(out)
|
||
}
|
||
}
|
||
|
||
// -- 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`.
|
||
/// The first hidden state is the post-norm input embedding (matches
|
||
/// the HF `output_hidden_states` convention used by `WavLMForXVector`).
|
||
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());
|
||
// Layer 0 computes the position_bias (shared across all 12 layers).
|
||
let mut position_bias: Option<Tensor> = None;
|
||
for layer in &self.layers {
|
||
let (next_h, bias) = layer.forward_with_bias(&h, position_bias.as_ref())?;
|
||
h = next_h;
|
||
position_bias = Some(bias);
|
||
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.get(NUM_LAYERS + 1, "layer_weights")?;
|
||
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 relative_position_bucket_basics() {
|
||
// Distance 0 → bucket 0 (small, left side).
|
||
assert_eq!(relative_position_bucket(0, 320, 800), 0);
|
||
// Tiny positive → goes to right half (bucket >= 160).
|
||
let b1 = relative_position_bucket(1, 320, 800);
|
||
assert!(b1 >= 160 && b1 < 320);
|
||
// Large positive → still on right half, capped to half - 1.
|
||
let b_large = relative_position_bucket(10_000, 320, 800);
|
||
assert_eq!(b_large, 320 - 1);
|
||
// Negative on the left half.
|
||
let b_neg = relative_position_bucket(-1, 320, 800);
|
||
assert!(b_neg < 160);
|
||
}
|
||
|
||
#[test]
|
||
fn encoder_layer_runs_with_layer0_bias() {
|
||
let device = Device::Cpu;
|
||
let (_vm, vb) = random_vb(&device);
|
||
let layer = WavLmEncoderLayer::new(0, vb).unwrap();
|
||
let xs = Tensor::randn(0f32, 1f32, (1, 24, HIDDEN_DIM), &device).unwrap();
|
||
let (out, bias) = layer.forward_with_bias(&xs, None).unwrap();
|
||
assert_eq!(out.dims(), &[1, 24, HIDDEN_DIM]);
|
||
assert_eq!(bias.dims(), &[NUM_HEADS, 24, 24]);
|
||
// Verify the layer actually transformed the input (not a no-op).
|
||
let diff = (&out - &xs).unwrap();
|
||
let l2 = diff.sqr().unwrap().sum_all().unwrap();
|
||
let l2: f32 = l2.to_dtype(DType::F32).unwrap().to_scalar().unwrap();
|
||
assert!(l2 > 1e-6, "encoder layer is a no-op (l2 = {l2})");
|
||
}
|
||
|
||
#[test]
|
||
fn encoder_threads_position_bias_through_stack() {
|
||
let device = Device::Cpu;
|
||
let (_vm, vb) = random_vb(&device);
|
||
let enc = Encoder::new(vb).unwrap();
|
||
let xs = Tensor::randn(0f32, 1f32, (1, 24, HIDDEN_DIM), &device).unwrap();
|
||
let states = enc.forward_all_layers(&xs).unwrap();
|
||
assert_eq!(states.len(), NUM_LAYERS + 1);
|
||
for s in &states {
|
||
assert_eq!(s.dims(), &[1, 24, HIDDEN_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}"
|
||
);
|
||
}
|
||
}
|