# 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.