Files

T

osobhandClaude Opus 4.7 209279c13e rtx-csm: Phase 13.9 — wav2vec2 candle port (slices 1+2+3, real ASR working)

Full port of facebook/wav2vec2-base-960h (94.4 M params, MIT) closing
the WhisperX-class word-alignment gap from the audio-ML survey. Same
staged-scaffolding pattern that worked for emotion2vec — but landed
slices 1+2+3 in one session.

src/wav2vec2.rs ships:
  - Wav2Vec2Config::base_960h
  - FeatureExtractor — 7 Conv1d (1→512, total stride 320). Layer 0
    uses GroupNorm with num_groups=num_channels=512 (HF's wav2vec2
    feat_extract_norm: "group"). Critical: state-dict key is
    layer_norm.* but the OP is GroupNorm — loading as LayerNorm
    produces empty CTC output.
  - FeatureProjection — LayerNorm(512) + Linear(512→768)
  - ConvPosEmbedding — kernel 128 grouped Conv1d, materialized at
    load time from upstream weight_g + weight_v (fairseq's weight_norm
    on dim=2; eps-guarded division for numerical stability)
  - Block — POST-norm transformer with separate Q/K/V (vs emotion2vec's
    fused QKV), uses (B*H, T, D) Metal 3D-matmul workaround from
    Phase 8.8 Moonshine
  - Encoder — pos_conv + initial LayerNorm + 12 Blocks
  - Wav2Vec2 top-level — load_from_safetensors via mmap'd VarBuilder
  - ctc_greedy_decode + VOCAB_960H constant for the 32-char alphabet

examples/wav2vec2_inspect.rs (slice 1): dumps tensor layout + config
examples/wav2vec2_smoke.rs (slice 3): real-weight load + ASR forward

Verified on Metal:
  loaded model in 0.28 s
  forward in 9 ms for 10.42 s audio (~1150× realtime)
  transcript: "HE HOPED THERE WOULD BE STEW FOR DINNER TURNIPS AND
              CARROTS AND BRUISED POTATOES AND FAT MUTTON PIECES TO
              BE LADLED OUT IN THICK PEPPERED FLOWER FAT AND SAUCE"

Numerical parity with upstream Python — the FLOWER-for-FLOUR typo is
the known wav2vec2-base-960h failure mode, matches HF reference exactly.

7 new unit tests; lib suite 127/127 (was 120).

Slice 4 remaining: Viterbi forced alignment given known transcript,
to emit (token, frame_start_ms, frame_end_ms) for word-boundary cuts.
The ASR path itself is now production-ready.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

2026-04-28 03:33:35 -07:00

3.6 KiB

Raw Blame History

wav2vec2 candle port — design notes

Inspector: cargo run -p rtx-csm --release --example wav2vec2_inspect

Repo: https://huggingface.co/facebook/wav2vec2-base-960h (model.safetensors, 378 MB, 94.4 M params, F32). MIT license.

Goal: word-level forced alignment for the data-prep stack — given known transcript T and audio A, run wav2vec2 + CTC + Viterbi to get a (token, frame_start, frame_end) table. Closes the WhisperX-class gap identified in the audio-ML Rust survey.

Architecture (from `config.json` + 212 inspected tensors)

audio 16 kHz (1, T)
 → feature_extractor          7 × Conv1d, total stride 320
                              kernels [10, 3, 3, 3, 3, 2, 2]
                              strides [5, 2, 2, 2, 2, 2, 2]
                              layer 0 = Conv1d → GroupNorm(512) → GELU
                              layers 1-6 = Conv1d → GELU (no norm)
 → feature_projection         LayerNorm(512) + Linear(512 → 768)
 → conv pos embedding         Conv1d(768→768, kernel 128, groups 16)
                              + GELU; output added to input as bias
 → encoder.layers.0..11       12 × POST-norm transformer block
                              (separate Q/K/V, NOT fused like emotion2vec)
 → lm_head                    Linear(768 → 32)  ← CTC head
 → CTC argmax (greedy decode) or Viterbi (forced alignment)

Per-block (encoder.layers.{i}.*, 16 tensors each):

attention.q_proj.weight/bias    Linear(768→768)
attention.k_proj.weight/bias    Linear(768→768)
attention.v_proj.weight/bias    Linear(768→768)
attention.out_proj.weight/bias  Linear(768→768)
layer_norm.weight/bias          LayerNorm(768)  — POST-norm (after attn+residual)
feed_forward.intermediate_dense.weight/bias  Linear(768→3072)
feed_forward.output_dense.weight/bias        Linear(3072→768)
final_layer_norm.weight/bias    LayerNorm(768)  — POST-norm (after FFN+residual)

POST-norm forward: x = layer_norm(x + attn(x)); x = final_layer_norm(x + ffn(x)).

CTC vocab (32 chars, from `vocab.json`)

<pad>=0  <s>=1  </s>=2  <unk>=3  |=4 (word separator)
E=5 T=6 A=7 O=8 N=9 I=10 H=11 S=12 R=13 D=14 L=15 U=16 M=17 W=18 C=19
F=20 G=21 Y=22 P=23 B=24 V=25 K=26 '=27 X=28 J=29 Q=30 Z=31

| = word separator (used between words during alignment).

Differences vs the Phase 13.8 emotion2vec port

Aspect	emotion2vec_plus_base	wav2vec2-base-960h
Norm order	PRE-norm	POST-norm
QKV	Fused (qkv 768→2304)	Separate q/k/v Linear
Pos encoding	5-stack Conv1d, kernel 19	1 Conv1d, kernel 128
Feature norm	LayerNorm every layer	GroupNorm only on layer 0
Output	9-class (softmax)	32-char (CTC log-softmax)
Pickle	`.pt` + descend `model`	clean safetensors mmap

Slicing plan

✅ Slice 1 (this commit): inspector + design notes
⏳ Slice 2a (~1 h): Wav2Vec2Config + FeatureExtractor (7 Conv1d
- GroupNorm on layer 0)
⏳ Slice 2b (~30 min): FeatureProjection (LN + Linear)
⏳ Slice 2c (~30 min): ConvPosEmbedding (single Conv1d kernel 128, with same-padding handling for even kernel)
⏳ Slice 2d (~1 h): Wav2Vec2Block POST-norm + Wav2Vec2Encoder 12 blocks
⏳ Slice 2e (~1 h): top-level Wav2Vec2 + safetensors loader + lm_head + a wav2vec2_smoke example
⏳ Slice 3 (~1 h): greedy CTC decode → ASR transcript on real audio
⏳ Slice 4 (~1-2 h): Viterbi forced alignment given a known transcript; emit (token, frame_start_ms, frame_end_ms) JSON

Total: ~5-6 hours of focused work. Each slice is independently shippable + testable.

3.6 KiB Raw Blame History Unescape Escape