1071 lines
38 KiB
Rust
1071 lines
38 KiB
Rust
//! Phase 13.9 — candle port of `facebook/wav2vec2-base-960h` for
|
||
//! word-level CTC forced alignment.
|
||
//!
|
||
//! Builds on the patterns from the Phase 13.8 emotion2vec port (`emotion2vec.rs`)
|
||
//! but with the wav2vec2-specific differences: POST-norm transformer
|
||
//! blocks, separate Q/K/V linears (not fused), GroupNorm on the first
|
||
//! feature-extractor layer only, single conv positional embedding
|
||
//! (kernel 128) instead of a 5-stack, and a 32-char CTC output head.
|
||
//!
|
||
//! Architecture (full notes in `docs/wav2vec2_port_notes.md`):
|
||
//! ```
|
||
//! audio 16 kHz (B, 1, T)
|
||
//! → feature_extractor 7 × Conv1d, T → T/320
|
||
//! → feature_projection LayerNorm(512) + Linear(512 → 768)
|
||
//! → conv_pos_embedding Conv1d(768→768, kernel 128, groups 16) + GELU; bias
|
||
//! → encoder.layers.0..11 12 × POST-norm transformer block
|
||
//! → lm_head Linear(768 → 32)
|
||
//! → log_softmax → 32-char CTC log-probs (B, T', 32)
|
||
//! ```
|
||
//!
|
||
//! Slice plan:
|
||
//! - **Slice 1**: inspector + design notes ✅
|
||
//! - **Slice 2 (this module)**: full candle port with safetensors loader
|
||
//! - **Slice 3**: greedy CTC decode → ASR transcript smoke test
|
||
//! - **Slice 4**: Viterbi forced alignment given a known transcript
|
||
|
||
use crate::error::{CsmError, Result as CsmResult};
|
||
use candle_core::{Device, Module, Tensor};
|
||
use candle_nn::{
|
||
Conv1d, Conv1dConfig, GroupNorm, LayerNorm, Linear, VarBuilder, conv1d, conv1d_no_bias,
|
||
group_norm, layer_norm, linear,
|
||
};
|
||
|
||
/// Architecture hyperparameters from `config.json`. Only the fields
|
||
/// needed by the inference path are surfaced.
|
||
#[derive(Debug, Clone)]
|
||
pub struct Wav2Vec2Config {
|
||
pub embed_dim: usize,
|
||
pub num_heads: usize,
|
||
pub mlp_dim: usize,
|
||
pub num_layers: usize,
|
||
pub vocab_size: usize,
|
||
pub norm_eps: f64,
|
||
pub feature_dim: usize,
|
||
/// `(out_channels, kernel, stride)` per conv layer.
|
||
pub conv_layers: Vec<(usize, usize, usize)>,
|
||
/// Conv positional embedding kernel (128 for `_base`).
|
||
pub conv_pos_kernel: usize,
|
||
/// Conv positional embedding groups (16 for `_base`).
|
||
pub conv_pos_groups: usize,
|
||
}
|
||
|
||
impl Wav2Vec2Config {
|
||
pub fn base_960h() -> Self {
|
||
Self {
|
||
embed_dim: 768,
|
||
num_heads: 12,
|
||
mlp_dim: 3072,
|
||
num_layers: 12,
|
||
vocab_size: 32,
|
||
norm_eps: 1e-5,
|
||
feature_dim: 512,
|
||
conv_layers: vec![
|
||
(512, 10, 5),
|
||
(512, 3, 2),
|
||
(512, 3, 2),
|
||
(512, 3, 2),
|
||
(512, 3, 2),
|
||
(512, 2, 2),
|
||
(512, 2, 2),
|
||
],
|
||
conv_pos_kernel: 128,
|
||
conv_pos_groups: 16,
|
||
}
|
||
}
|
||
|
||
pub fn conv_stride_total(&self) -> usize {
|
||
self.conv_layers.iter().map(|(_, _, s)| *s).product()
|
||
}
|
||
}
|
||
|
||
/// One conv layer in the feature extractor. Layer 0 in `_base` carries a
|
||
/// `GroupNorm` with `num_groups = num_channels` (= out_conv_dim, 512) —
|
||
/// per HF's wav2vec2 with `feat_extract_norm: "group"`. The state-dict
|
||
/// param name is `layer_norm.weight/bias` purely because of the wrapping
|
||
/// `Wav2Vec2GroupNormConvLayer.layer_norm` attribute name; the *operation*
|
||
/// is GroupNorm, not LayerNorm. Affine shape `[512]` matches both ops, so
|
||
/// loaders that picked LayerNorm here produced empty/garbage output.
|
||
/// Layers 1..6 are bare `Conv1d → GELU`.
|
||
#[derive(Debug)]
|
||
struct FeatureExtractorBlock {
|
||
conv: Conv1d,
|
||
/// Only present on layer 0 for `_base`. GroupNorm with
|
||
/// num_groups = num_channels = InstanceNorm semantics.
|
||
norm: Option<GroupNorm>,
|
||
}
|
||
|
||
impl FeatureExtractorBlock {
|
||
fn new(
|
||
in_channels: usize,
|
||
out_channels: usize,
|
||
kernel_size: usize,
|
||
stride: usize,
|
||
norm_eps: f64,
|
||
with_norm: bool,
|
||
vb: VarBuilder,
|
||
) -> CsmResult<Self> {
|
||
let cfg = Conv1dConfig {
|
||
stride,
|
||
..Default::default()
|
||
};
|
||
let conv = conv1d_no_bias(in_channels, out_channels, kernel_size, cfg, vb.pp("conv"))
|
||
.map_err(|e| CsmError::Config(format!("feat ext conv: {e}")))?;
|
||
let norm = if with_norm {
|
||
// num_groups = num_channels per HF wav2vec2 group-norm config.
|
||
Some(
|
||
group_norm(out_channels, out_channels, norm_eps, vb.pp("layer_norm"))
|
||
.map_err(|e| CsmError::Config(format!("feat ext group_norm: {e}")))?,
|
||
)
|
||
} else {
|
||
None
|
||
};
|
||
Ok(Self { conv, norm })
|
||
}
|
||
}
|
||
|
||
impl Module for FeatureExtractorBlock {
|
||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
// (B, C_in, T_in) → (B, C_out, T_out)
|
||
let xs = self.conv.forward(xs)?;
|
||
let xs = if let Some(n) = &self.norm {
|
||
// GroupNorm operates on (B, C, T) directly — no transpose.
|
||
n.forward(&xs)?
|
||
} else {
|
||
xs
|
||
};
|
||
xs.gelu()
|
||
}
|
||
}
|
||
|
||
/// 7-layer feature extractor: raw 16 kHz mono → `(B, 512, T/320)`.
|
||
#[derive(Debug)]
|
||
pub struct FeatureExtractor {
|
||
blocks: Vec<FeatureExtractorBlock>,
|
||
}
|
||
|
||
impl FeatureExtractor {
|
||
pub fn new(cfg: &Wav2Vec2Config, vb: VarBuilder) -> CsmResult<Self> {
|
||
let mut blocks = Vec::with_capacity(cfg.conv_layers.len());
|
||
let mut in_ch = 1usize;
|
||
for (i, (out_ch, k, s)) in cfg.conv_layers.iter().enumerate() {
|
||
let with_norm = i == 0; // `_base` only norms layer 0.
|
||
let block = FeatureExtractorBlock::new(
|
||
in_ch,
|
||
*out_ch,
|
||
*k,
|
||
*s,
|
||
cfg.norm_eps,
|
||
with_norm,
|
||
vb.pp("conv_layers").pp(i.to_string()),
|
||
)?;
|
||
blocks.push(block);
|
||
in_ch = *out_ch;
|
||
}
|
||
Ok(Self { blocks })
|
||
}
|
||
}
|
||
|
||
impl Module for FeatureExtractor {
|
||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
let mut xs = xs.clone();
|
||
for block in self.blocks.iter() {
|
||
xs = block.forward(&xs)?;
|
||
}
|
||
Ok(xs)
|
||
}
|
||
}
|
||
|
||
/// `LayerNorm(feature_dim) + Linear(feature_dim → embed_dim)` projecting
|
||
/// the feature extractor output up to the transformer dim.
|
||
#[derive(Debug)]
|
||
pub struct FeatureProjection {
|
||
norm: LayerNorm,
|
||
proj: Linear,
|
||
}
|
||
|
||
impl FeatureProjection {
|
||
pub fn new(in_dim: usize, out_dim: usize, norm_eps: f64, vb: VarBuilder) -> CsmResult<Self> {
|
||
let norm = layer_norm(in_dim, norm_eps, vb.pp("layer_norm"))
|
||
.map_err(|e| CsmError::Config(format!("feat proj norm: {e}")))?;
|
||
let proj = linear(in_dim, out_dim, vb.pp("projection"))
|
||
.map_err(|e| CsmError::Config(format!("feat proj linear: {e}")))?;
|
||
Ok(Self { norm, proj })
|
||
}
|
||
}
|
||
|
||
impl Module for FeatureProjection {
|
||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
// xs: (B, T, in_dim) — channel-last
|
||
let xs = self.norm.forward(xs)?;
|
||
self.proj.forward(&xs)
|
||
}
|
||
}
|
||
|
||
/// Conv-based positional embedding: a single grouped Conv1d (kernel 128,
|
||
/// groups 16) with same-padding, GELU activation, output added to the
|
||
/// input as a positional bias. The pickle key prefix is
|
||
/// `encoder.pos_conv_embed.conv` (Conv1d only; the GELU has no params).
|
||
///
|
||
/// Note: for an EVEN kernel (128), exact same-padding requires asymmetric
|
||
/// padding. Standard transformers-style impl pads both sides by 64 and
|
||
/// drops the last frame to keep T unchanged.
|
||
#[derive(Debug)]
|
||
pub struct ConvPosEmbedding {
|
||
conv: Conv1d,
|
||
drop_last: bool,
|
||
}
|
||
|
||
impl ConvPosEmbedding {
|
||
/// Build with random init (used by tests). Production code path is
|
||
/// [`Self::from_weight_normed`] because the upstream pickle stores
|
||
/// `weight_g` + `weight_v` instead of a materialized `weight`.
|
||
pub fn new(embed_dim: usize, kernel: usize, groups: usize, vb: VarBuilder) -> CsmResult<Self> {
|
||
let (pad, drop_last) = pad_for(kernel);
|
||
let cfg = Conv1dConfig {
|
||
padding: pad,
|
||
stride: 1,
|
||
dilation: 1,
|
||
groups,
|
||
cudnn_fwd_algo: None,
|
||
};
|
||
let conv = conv1d(embed_dim, embed_dim, kernel, cfg, vb.pp("conv"))
|
||
.map_err(|e| CsmError::Config(format!("pos conv: {e}")))?;
|
||
Ok(Self { conv, drop_last })
|
||
}
|
||
|
||
/// Materialize a weight-normed Conv1d from upstream's `weight_g` +
|
||
/// `weight_v` storage. fairseq's wav2vec2 applies
|
||
/// `nn.utils.weight_norm(self.pos_conv, name="weight", dim=2)` so the
|
||
/// stored shapes are:
|
||
/// weight_v: (out, in/groups, kernel)
|
||
/// weight_g: (1, 1, kernel) ← per-kernel-position gain
|
||
/// And the materialized weight is `weight_g * weight_v / ||weight_v||_2`
|
||
/// where the L2 norm is along axis 2 (kernel), keepdim=true so it
|
||
/// broadcasts back to `weight_v`'s shape.
|
||
pub fn from_weight_normed(
|
||
embed_dim: usize,
|
||
kernel: usize,
|
||
groups: usize,
|
||
vb: VarBuilder,
|
||
) -> CsmResult<Self> {
|
||
let (pad, drop_last) = pad_for(kernel);
|
||
let cfg = Conv1dConfig {
|
||
padding: pad,
|
||
stride: 1,
|
||
dilation: 1,
|
||
groups,
|
||
cudnn_fwd_algo: None,
|
||
};
|
||
let v_shape = (embed_dim, embed_dim / groups, kernel);
|
||
let weight_v = vb
|
||
.pp("conv")
|
||
.get(v_shape, "weight_v")
|
||
.map_err(|e| CsmError::Config(format!("pos conv weight_v: {e}")))?;
|
||
let weight_g = vb
|
||
.pp("conv")
|
||
.get((1, 1, kernel), "weight_g")
|
||
.map_err(|e| CsmError::Config(format!("pos conv weight_g: {e}")))?;
|
||
let bias = vb
|
||
.pp("conv")
|
||
.get(embed_dim, "bias")
|
||
.map_err(|e| CsmError::Config(format!("pos conv bias: {e}")))?;
|
||
// Materialize: weight = weight_g * weight_v / (||weight_v||_dim=2 + eps)
|
||
// Eps mirrors PyTorch's weight_norm numerical guard and keeps
|
||
// random-init tests numerically stable.
|
||
let norm = weight_v
|
||
.sqr()
|
||
.and_then(|t| t.sum_keepdim(2))
|
||
.and_then(|t| t.sqrt())
|
||
.and_then(|t| (t + 1e-12)?.broadcast_as(weight_v.shape()))
|
||
.map_err(|e| CsmError::Config(format!("pos conv weight norm: {e}")))?;
|
||
let weight = (weight_v / &norm)
|
||
.and_then(|t| t.broadcast_mul(&weight_g))
|
||
.map_err(|e| CsmError::Config(format!("pos conv materialize: {e}")))?;
|
||
let conv = Conv1d::new(weight, Some(bias), cfg);
|
||
Ok(Self { conv, drop_last })
|
||
}
|
||
}
|
||
|
||
/// Compute `(padding, drop_last_frame)` for same-padding given a kernel.
|
||
/// Even kernels need asymmetric handling: pad symmetrically by k/2 and
|
||
/// trim the trailing extra frame.
|
||
fn pad_for(kernel: usize) -> (usize, bool) {
|
||
if kernel.is_multiple_of(2) {
|
||
(kernel / 2, true)
|
||
} else {
|
||
((kernel - 1) / 2, false)
|
||
}
|
||
}
|
||
|
||
impl Module for ConvPosEmbedding {
|
||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
// xs: (B, T, C). Conv1d operates on (B, C, T).
|
||
let h = xs.transpose(1, 2)?.contiguous()?;
|
||
let h = self.conv.forward(&h)?;
|
||
let h = if self.drop_last {
|
||
// Trim trailing frame from even-kernel padding.
|
||
let t = h.dim(2)?;
|
||
h.narrow(2, 0, t - 1)?
|
||
} else {
|
||
h
|
||
};
|
||
let h = h.gelu()?;
|
||
let bias = h.transpose(1, 2)?.contiguous()?;
|
||
xs + bias
|
||
}
|
||
}
|
||
|
||
/// One POST-norm transformer block. Pickle keys per block (16 tensors):
|
||
/// `attention.{q,k,v,out}_proj.{weight,bias}`,
|
||
/// `layer_norm.{weight,bias}`,
|
||
/// `feed_forward.{intermediate,output}_dense.{weight,bias}`,
|
||
/// `final_layer_norm.{weight,bias}`.
|
||
#[derive(Debug)]
|
||
pub struct Block {
|
||
q: Linear,
|
||
k: Linear,
|
||
v: Linear,
|
||
out: Linear,
|
||
layer_norm: LayerNorm,
|
||
fc1: Linear,
|
||
fc2: Linear,
|
||
final_layer_norm: LayerNorm,
|
||
num_heads: usize,
|
||
head_dim: usize,
|
||
scale: f64,
|
||
}
|
||
|
||
impl Block {
|
||
pub fn new(
|
||
embed_dim: usize,
|
||
num_heads: usize,
|
||
mlp_dim: usize,
|
||
norm_eps: f64,
|
||
vb: VarBuilder,
|
||
) -> CsmResult<Self> {
|
||
let head_dim = embed_dim / num_heads;
|
||
let scale = 1.0 / (head_dim as f64).sqrt();
|
||
let q = linear(embed_dim, embed_dim, vb.pp("attention").pp("q_proj"))
|
||
.map_err(|e| CsmError::Config(format!("block q: {e}")))?;
|
||
let k = linear(embed_dim, embed_dim, vb.pp("attention").pp("k_proj"))
|
||
.map_err(|e| CsmError::Config(format!("block k: {e}")))?;
|
||
let v = linear(embed_dim, embed_dim, vb.pp("attention").pp("v_proj"))
|
||
.map_err(|e| CsmError::Config(format!("block v: {e}")))?;
|
||
let out = linear(embed_dim, embed_dim, vb.pp("attention").pp("out_proj"))
|
||
.map_err(|e| CsmError::Config(format!("block out: {e}")))?;
|
||
let mid_norm = layer_norm(embed_dim, norm_eps, vb.pp("layer_norm"))
|
||
.map_err(|e| CsmError::Config(format!("block ln: {e}")))?;
|
||
let fc1 = linear(
|
||
embed_dim,
|
||
mlp_dim,
|
||
vb.pp("feed_forward").pp("intermediate_dense"),
|
||
)
|
||
.map_err(|e| CsmError::Config(format!("block fc1: {e}")))?;
|
||
let fc2 = linear(mlp_dim, embed_dim, vb.pp("feed_forward").pp("output_dense"))
|
||
.map_err(|e| CsmError::Config(format!("block fc2: {e}")))?;
|
||
let final_norm = layer_norm(embed_dim, norm_eps, vb.pp("final_layer_norm"))
|
||
.map_err(|e| CsmError::Config(format!("block fln: {e}")))?;
|
||
Ok(Self {
|
||
q,
|
||
k,
|
||
v,
|
||
out,
|
||
layer_norm: mid_norm,
|
||
fc1,
|
||
fc2,
|
||
final_layer_norm: final_norm,
|
||
num_heads,
|
||
head_dim,
|
||
scale,
|
||
})
|
||
}
|
||
|
||
fn attention(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
// xs: (B, T, D). Standard separate-QKV multi-head attention.
|
||
let (b, t, _d) = xs.dims3()?;
|
||
let q = self
|
||
.q
|
||
.forward(xs)?
|
||
.reshape((b, t, self.num_heads, self.head_dim))?
|
||
.transpose(1, 2)?
|
||
.contiguous()?;
|
||
let k = self
|
||
.k
|
||
.forward(xs)?
|
||
.reshape((b, t, self.num_heads, self.head_dim))?
|
||
.transpose(1, 2)?
|
||
.contiguous()?;
|
||
let v = self
|
||
.v
|
||
.forward(xs)?
|
||
.reshape((b, t, self.num_heads, self.head_dim))?
|
||
.transpose(1, 2)?
|
||
.contiguous()?;
|
||
// Collapse (B, H) for the Metal-friendly 3D matmul (Phase 8.8 Moonshine
|
||
// workaround for candle's 4D matmul shape bug).
|
||
let bh = b * self.num_heads;
|
||
let q3 = q.reshape((bh, t, self.head_dim))?;
|
||
let k3 = k.reshape((bh, t, self.head_dim))?;
|
||
let v3 = v.reshape((bh, t, self.head_dim))?;
|
||
let scores = (q3.matmul(&k3.transpose(1, 2)?)? * self.scale)?;
|
||
let attn = candle_nn::ops::softmax_last_dim(&scores)?;
|
||
let out = attn.matmul(&v3)?;
|
||
let out = out
|
||
.reshape((b, self.num_heads, t, self.head_dim))?
|
||
.transpose(1, 2)?
|
||
.reshape((b, t, self.num_heads * self.head_dim))?
|
||
.contiguous()?;
|
||
self.out.forward(&out)
|
||
}
|
||
|
||
fn ffn(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
let xs = self.fc1.forward(xs)?;
|
||
let xs = xs.gelu()?;
|
||
self.fc2.forward(&xs)
|
||
}
|
||
}
|
||
|
||
impl Module for Block {
|
||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
// POST-norm: x = layer_norm(x + attn(x)); x = final_layer_norm(x + ffn(x))
|
||
let h = self.attention(xs)?;
|
||
let xs = (xs + h)?;
|
||
let xs = self.layer_norm.forward(&xs)?;
|
||
let h = self.ffn(&xs)?;
|
||
let xs = (xs + h)?;
|
||
self.final_layer_norm.forward(&xs)
|
||
}
|
||
}
|
||
|
||
/// 12-block encoder + the ConvPosEmbedding bias added before block 0.
|
||
/// State-dict keys live under `wav2vec2.encoder.*`:
|
||
/// `encoder.pos_conv_embed.conv.{weight,bias}`
|
||
/// `encoder.layer_norm.{weight,bias}` ← initial LN before blocks
|
||
/// `encoder.layers.{0..11}.*`
|
||
#[derive(Debug)]
|
||
pub struct Encoder {
|
||
pos: ConvPosEmbedding,
|
||
init_norm: LayerNorm,
|
||
blocks: Vec<Block>,
|
||
}
|
||
|
||
impl Encoder {
|
||
pub fn new(cfg: &Wav2Vec2Config, vb: VarBuilder) -> CsmResult<Self> {
|
||
let pos = ConvPosEmbedding::from_weight_normed(
|
||
cfg.embed_dim,
|
||
cfg.conv_pos_kernel,
|
||
cfg.conv_pos_groups,
|
||
vb.pp("pos_conv_embed"),
|
||
)?;
|
||
let init_norm = layer_norm(cfg.embed_dim, cfg.norm_eps, vb.pp("layer_norm"))
|
||
.map_err(|e| CsmError::Config(format!("encoder init_norm: {e}")))?;
|
||
let mut blocks = Vec::with_capacity(cfg.num_layers);
|
||
let vb_layers = vb.pp("layers");
|
||
for i in 0..cfg.num_layers {
|
||
blocks.push(Block::new(
|
||
cfg.embed_dim,
|
||
cfg.num_heads,
|
||
cfg.mlp_dim,
|
||
cfg.norm_eps,
|
||
vb_layers.pp(i.to_string()),
|
||
)?);
|
||
}
|
||
Ok(Self {
|
||
pos,
|
||
init_norm,
|
||
blocks,
|
||
})
|
||
}
|
||
}
|
||
|
||
impl Module for Encoder {
|
||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
// xs: (B, T, D)
|
||
let xs = self.pos.forward(xs)?;
|
||
let mut xs = self.init_norm.forward(&xs)?;
|
||
for block in self.blocks.iter() {
|
||
xs = block.forward(&xs)?;
|
||
}
|
||
Ok(xs)
|
||
}
|
||
}
|
||
|
||
/// Top-level wav2vec2 + CTC head. State-dict prefix: everything below
|
||
/// `wav2vec2.*` (feature_extractor, feature_projection, encoder), plus
|
||
/// `lm_head.{weight,bias}` at the root.
|
||
#[derive(Debug)]
|
||
pub struct Wav2Vec2 {
|
||
cfg: Wav2Vec2Config,
|
||
feature_extractor: FeatureExtractor,
|
||
feature_projection: FeatureProjection,
|
||
encoder: Encoder,
|
||
lm_head: Linear,
|
||
device: Device,
|
||
}
|
||
|
||
impl Wav2Vec2 {
|
||
pub fn new(cfg: Wav2Vec2Config, vb: VarBuilder, device: Device) -> CsmResult<Self> {
|
||
let vb_w = vb.pp("wav2vec2");
|
||
let feature_extractor = FeatureExtractor::new(&cfg, vb_w.pp("feature_extractor"))?;
|
||
let feature_projection = FeatureProjection::new(
|
||
cfg.feature_dim,
|
||
cfg.embed_dim,
|
||
cfg.norm_eps,
|
||
vb_w.pp("feature_projection"),
|
||
)?;
|
||
let encoder = Encoder::new(&cfg, vb_w.pp("encoder"))?;
|
||
let lm_head = linear(cfg.embed_dim, cfg.vocab_size, vb.pp("lm_head"))
|
||
.map_err(|e| CsmError::Config(format!("lm_head: {e}")))?;
|
||
Ok(Self {
|
||
cfg,
|
||
feature_extractor,
|
||
feature_projection,
|
||
encoder,
|
||
lm_head,
|
||
device,
|
||
})
|
||
}
|
||
|
||
/// Load from `model.safetensors` via mmap.
|
||
pub fn load_from_safetensors<P: AsRef<std::path::Path>>(
|
||
path: P,
|
||
device: &Device,
|
||
) -> CsmResult<Self> {
|
||
let vb = unsafe {
|
||
candle_nn::VarBuilder::from_mmaped_safetensors(
|
||
&[path.as_ref()],
|
||
candle_core::DType::F32,
|
||
device,
|
||
)
|
||
}
|
||
.map_err(|e| CsmError::Config(format!("wav2vec2 safetensors load: {e}")))?;
|
||
Self::new(Wav2Vec2Config::base_960h(), vb, device.clone())
|
||
}
|
||
|
||
/// Forward pass: 16 kHz audio (B, 1, T) → CTC logits (B, T', 32).
|
||
/// Caller applies `log_softmax` on the last dim for proper CTC math.
|
||
pub fn forward(&self, audio_16k: &Tensor) -> candle_core::Result<Tensor> {
|
||
let feats = self.feature_extractor.forward(audio_16k)?;
|
||
// (B, 512, T') → (B, T', 512)
|
||
let feats = feats.transpose(1, 2)?.contiguous()?;
|
||
let h = self.feature_projection.forward(&feats)?;
|
||
let h = self.encoder.forward(&h)?;
|
||
self.lm_head.forward(&h)
|
||
}
|
||
|
||
pub fn config(&self) -> &Wav2Vec2Config {
|
||
&self.cfg
|
||
}
|
||
|
||
pub fn device(&self) -> &Device {
|
||
&self.device
|
||
}
|
||
}
|
||
|
||
/// Default CTC vocab for `wav2vec2-base-960h` per the upstream `vocab.json`.
|
||
/// Used by greedy + Viterbi decoders when the user doesn't supply a custom
|
||
/// tokenizer. Index = token id; value = char (`'|'` is the word separator).
|
||
pub const VOCAB_960H: &[&str] = &[
|
||
"<pad>", "<s>", "</s>", "<unk>", "|", "E", "T", "A", "O", "N", "I", "H", "S", "R", "D", "L",
|
||
"U", "M", "W", "C", "F", "G", "Y", "P", "B", "V", "K", "'", "X", "J", "Q", "Z",
|
||
];
|
||
|
||
/// CTC blank token id (`<pad>` = 0 by upstream convention).
|
||
pub const CTC_BLANK_ID: usize = 0;
|
||
|
||
/// Greedy CTC decode: per frame argmax → collapse repeats → drop blanks.
|
||
/// Returns the decoded transcript as a String. `vocab` indexes are looked
|
||
/// up from [`VOCAB_960H`] by default; supply a custom slice for non-960h
|
||
/// checkpoints.
|
||
pub fn ctc_greedy_decode(logits: &Tensor, vocab: &[&str]) -> CsmResult<String> {
|
||
// logits: (1, T', V) or (T', V); reduce to (T', V).
|
||
let logits = if logits.dims().len() == 3 {
|
||
logits
|
||
.i((0, .., ..))
|
||
.map_err(|e| CsmError::Config(format!("ctc decode squeeze: {e}")))?
|
||
} else {
|
||
logits.clone()
|
||
};
|
||
let argmax: Vec<u32> = logits
|
||
.argmax(candle_core::D::Minus1)
|
||
.and_then(|t| t.to_vec1::<u32>())
|
||
.map_err(|e| CsmError::Config(format!("ctc argmax: {e}")))?;
|
||
// CTC collapse: skip if same as previous; drop blanks.
|
||
let mut out = String::new();
|
||
let mut prev: i64 = -1;
|
||
for &id in argmax.iter() {
|
||
let id_i = id as i64;
|
||
if id_i == prev {
|
||
continue;
|
||
}
|
||
if (id as usize) != CTC_BLANK_ID
|
||
&& let Some(tok) = vocab.get(id as usize)
|
||
{
|
||
if *tok == "|" {
|
||
out.push(' ');
|
||
} else if !tok.starts_with('<') {
|
||
out.push_str(tok);
|
||
}
|
||
}
|
||
prev = id_i;
|
||
}
|
||
Ok(out)
|
||
}
|
||
|
||
use candle_core::IndexOp;
|
||
|
||
/// One aligned token — output of [`viterbi_align`]. `token` is the raw
|
||
/// CTC unit (single character for the default `_base-960h` vocab).
|
||
/// `frame_start` / `frame_end` index 50 Hz feature frames; multiply by
|
||
/// 20 ms (1000 ms / 50 Hz) to convert to milliseconds.
|
||
#[derive(Debug, Clone, serde::Serialize)]
|
||
pub struct AlignedToken {
|
||
pub token: String,
|
||
pub frame_start: usize,
|
||
pub frame_end: usize,
|
||
}
|
||
|
||
/// Group consecutive aligned tokens into words. Word-separator handling:
|
||
/// the upstream `_base-960h` vocab uses `|` as the inter-word boundary,
|
||
/// which ctc_greedy_decode renders as a space. After [`viterbi_align`]
|
||
/// the `|` shows up as its own token; this helper folds adjacent
|
||
/// non-separator tokens into a `Word { text, frame_start, frame_end }`.
|
||
#[derive(Debug, Clone, serde::Serialize)]
|
||
pub struct AlignedWord {
|
||
pub word: String,
|
||
pub frame_start: usize,
|
||
pub frame_end: usize,
|
||
}
|
||
|
||
/// Convert a transcript string into the CTC token id sequence over the
|
||
/// supplied vocab. Spaces in the transcript become the word-separator
|
||
/// token (`|` for `_base-960h`); other characters are looked up directly.
|
||
/// Unknown characters are replaced with the `<unk>` token.
|
||
pub fn transcript_to_token_ids(transcript: &str, vocab: &[&str]) -> CsmResult<Vec<u32>> {
|
||
let mut id_for: std::collections::HashMap<&str, u32> = Default::default();
|
||
for (i, t) in vocab.iter().enumerate() {
|
||
id_for.insert(*t, i as u32);
|
||
}
|
||
let unk_id = *id_for
|
||
.get("<unk>")
|
||
.ok_or_else(|| CsmError::Config("vocab missing <unk>".into()))?;
|
||
let sep_id = *id_for
|
||
.get("|")
|
||
.ok_or_else(|| CsmError::Config("vocab missing | (word sep)".into()))?;
|
||
let mut out = Vec::new();
|
||
let upper = transcript.to_uppercase();
|
||
let chars: Vec<char> = upper.chars().collect();
|
||
let mut i = 0;
|
||
while i < chars.len() {
|
||
let c = chars[i];
|
||
if c == ' ' {
|
||
// Collapse runs of spaces into a single separator.
|
||
out.push(sep_id);
|
||
while i < chars.len() && chars[i] == ' ' {
|
||
i += 1;
|
||
}
|
||
continue;
|
||
}
|
||
let s = c.to_string();
|
||
let id = id_for.get(s.as_str()).copied().unwrap_or(unk_id);
|
||
out.push(id);
|
||
i += 1;
|
||
}
|
||
Ok(out)
|
||
}
|
||
|
||
/// CTC Viterbi forced alignment.
|
||
///
|
||
/// Given per-frame log-probabilities `log_probs` of shape `(T, V)` (caller
|
||
/// supplies the log-softmax) and a target token sequence `tokens`, find
|
||
/// the maximum-likelihood alignment of `tokens` against the frames using
|
||
/// the standard CTC transition lattice (interleave each token with a
|
||
/// blank, allow self-loop or advance-to-next at every step).
|
||
///
|
||
/// Returns one [`AlignedToken`] per token in the input sequence with
|
||
/// inclusive `frame_start..=frame_end` indices into the 50 Hz feature
|
||
/// frame grid.
|
||
///
|
||
/// Error if `T < tokens.len()` (not enough frames) or if `tokens` is empty.
|
||
pub fn viterbi_align(
|
||
log_probs: &Tensor,
|
||
tokens: &[u32],
|
||
blank_id: usize,
|
||
vocab: &[&str],
|
||
) -> CsmResult<Vec<AlignedToken>> {
|
||
if tokens.is_empty() {
|
||
return Err(CsmError::Config("viterbi_align: tokens is empty".into()));
|
||
}
|
||
let lp = if log_probs.dims().len() == 3 {
|
||
log_probs
|
||
.i((0, .., ..))
|
||
.map_err(|e| CsmError::Config(format!("viterbi squeeze: {e}")))?
|
||
} else {
|
||
log_probs.clone()
|
||
};
|
||
let lp_vec = lp
|
||
.to_vec2::<f32>()
|
||
.map_err(|e| CsmError::Config(format!("viterbi to_vec2: {e}")))?;
|
||
let big_t = lp_vec.len();
|
||
let v = if big_t > 0 { lp_vec[0].len() } else { 0 };
|
||
if big_t == 0 || v == 0 {
|
||
return Err(CsmError::Config("viterbi_align: empty log_probs".into()));
|
||
}
|
||
|
||
// CTC trellis state sequence: [blank, t0, blank, t1, blank, ..., tN, blank]
|
||
// length S = 2 * tokens.len() + 1
|
||
let n = tokens.len();
|
||
let big_s = 2 * n + 1;
|
||
if big_t < n {
|
||
return Err(CsmError::Config(format!(
|
||
"viterbi_align: T={big_t} < tokens.len()={n}, can't align"
|
||
)));
|
||
}
|
||
|
||
// Build the state token-id table.
|
||
let mut state_id: Vec<usize> = Vec::with_capacity(big_s);
|
||
for (i, _) in (0..big_s).enumerate() {
|
||
if i % 2 == 0 {
|
||
state_id.push(blank_id);
|
||
} else {
|
||
state_id.push(tokens[i / 2] as usize);
|
||
}
|
||
}
|
||
|
||
let neg_inf = f32::NEG_INFINITY;
|
||
// dp[t][s] = best log-prob for reaching state s at frame t
|
||
// backptr[t][s] = state at t-1 we came from (for path recovery)
|
||
let mut dp = vec![vec![neg_inf; big_s]; big_t];
|
||
let mut backptr = vec![vec![0usize; big_s]; big_t];
|
||
|
||
// Initialization: at t=0 we can be in state 0 (initial blank) or
|
||
// state 1 (first token).
|
||
dp[0][0] = lp_vec[0][state_id[0]];
|
||
if big_s > 1 {
|
||
dp[0][1] = lp_vec[0][state_id[1]];
|
||
}
|
||
|
||
for t in 1..big_t {
|
||
for s in 0..big_s {
|
||
// Stay in s, or advance from s-1, or skip from s-2 (only if
|
||
// s is a non-blank state AND state s-2 is a different token —
|
||
// CTC's "ε-skip" rule allows skipping blank between two
|
||
// *different* tokens but NOT between two same tokens).
|
||
let mut best = dp[t - 1][s];
|
||
let mut best_prev = s;
|
||
|
||
if s >= 1 && dp[t - 1][s - 1] > best {
|
||
best = dp[t - 1][s - 1];
|
||
best_prev = s - 1;
|
||
}
|
||
|
||
// ε-skip: only when current state is a non-blank token AND
|
||
// the token at s-2 is different (so we can skip the
|
||
// intermediate blank).
|
||
if s >= 2 && s % 2 == 1 {
|
||
let cur_tok = state_id[s];
|
||
let prev_tok = state_id[s - 2];
|
||
if cur_tok != prev_tok && dp[t - 1][s - 2] > best {
|
||
best = dp[t - 1][s - 2];
|
||
best_prev = s - 2;
|
||
}
|
||
}
|
||
// Add the emission cost for being in state s at frame t.
|
||
if best > neg_inf {
|
||
dp[t][s] = best + lp_vec[t][state_id[s]];
|
||
backptr[t][s] = best_prev;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Final state: must end in either the last token (s = 2N-1) or the
|
||
// trailing blank (s = 2N). Pick whichever has the higher log-prob.
|
||
let last_token_state = big_s - 2;
|
||
let last_blank_state = big_s - 1;
|
||
let (mut s, _) = if dp[big_t - 1][last_token_state] >= dp[big_t - 1][last_blank_state] {
|
||
(last_token_state, dp[big_t - 1][last_token_state])
|
||
} else {
|
||
(last_blank_state, dp[big_t - 1][last_blank_state])
|
||
};
|
||
|
||
// Recover the per-frame state path (reverse).
|
||
let mut path = vec![0usize; big_t];
|
||
path[big_t - 1] = s;
|
||
for t in (1..big_t).rev() {
|
||
s = backptr[t][s];
|
||
path[t - 1] = s;
|
||
}
|
||
|
||
// Convert to per-token spans by finding the first and last frames
|
||
// each non-blank state index (1, 3, 5, ..., 2N-1) appears in path.
|
||
let mut out: Vec<AlignedToken> = Vec::with_capacity(n);
|
||
for (i, &tok_id) in tokens.iter().enumerate() {
|
||
let target_state = 2 * i + 1;
|
||
let (mut start, mut end) = (None, None);
|
||
for (t, &p) in path.iter().enumerate() {
|
||
if p == target_state {
|
||
if start.is_none() {
|
||
start = Some(t);
|
||
}
|
||
end = Some(t);
|
||
}
|
||
}
|
||
// Fallback: if a token is degenerate / never visited, place it
|
||
// at the path's nearest blank-bracket. This shouldn't happen
|
||
// for legitimate inputs but keeps the function total.
|
||
let (frame_start, frame_end) = match (start, end) {
|
||
(Some(s0), Some(e0)) => (s0, e0),
|
||
_ => (0, 0),
|
||
};
|
||
let token = vocab
|
||
.get(tok_id as usize)
|
||
.map(|s| s.to_string())
|
||
.unwrap_or_else(|| format!("<id_{tok_id}>"));
|
||
out.push(AlignedToken {
|
||
token,
|
||
frame_start,
|
||
frame_end,
|
||
});
|
||
}
|
||
Ok(out)
|
||
}
|
||
|
||
/// Group [`AlignedToken`]s into words at `|` boundaries.
|
||
pub fn group_into_words(tokens: &[AlignedToken]) -> Vec<AlignedWord> {
|
||
let mut out = Vec::new();
|
||
let mut buf = String::new();
|
||
let mut start: Option<usize> = None;
|
||
let mut end: usize = 0;
|
||
for tok in tokens.iter() {
|
||
if tok.token == "|" {
|
||
if !buf.is_empty() {
|
||
out.push(AlignedWord {
|
||
word: std::mem::take(&mut buf),
|
||
frame_start: start.unwrap_or(tok.frame_start),
|
||
frame_end: end,
|
||
});
|
||
start = None;
|
||
}
|
||
continue;
|
||
}
|
||
if start.is_none() {
|
||
start = Some(tok.frame_start);
|
||
}
|
||
buf.push_str(&tok.token);
|
||
end = tok.frame_end;
|
||
}
|
||
if !buf.is_empty() {
|
||
out.push(AlignedWord {
|
||
word: buf,
|
||
frame_start: start.unwrap_or(0),
|
||
frame_end: end,
|
||
});
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Convert a frame index (50 Hz, conv stride 320 at 16 kHz) to milliseconds.
|
||
pub fn frame_to_ms(frame: usize) -> f32 {
|
||
frame as f32 * 20.0
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use candle_core::DType;
|
||
use candle_nn::VarMap;
|
||
|
||
#[test]
|
||
fn config_stride_product_matches_320() {
|
||
let cfg = Wav2Vec2Config::base_960h();
|
||
assert_eq!(cfg.conv_stride_total(), 320);
|
||
assert_eq!(cfg.conv_layers.len(), 7);
|
||
assert_eq!(cfg.num_layers, 12);
|
||
assert_eq!(cfg.vocab_size, 32);
|
||
}
|
||
|
||
#[test]
|
||
fn vocab_960h_size_matches_config() {
|
||
assert_eq!(VOCAB_960H.len(), Wav2Vec2Config::base_960h().vocab_size);
|
||
assert_eq!(VOCAB_960H[CTC_BLANK_ID], "<pad>");
|
||
assert_eq!(VOCAB_960H[4], "|");
|
||
}
|
||
|
||
#[test]
|
||
fn feature_extractor_random_init_shape() {
|
||
let dev = Device::Cpu;
|
||
let cfg = Wav2Vec2Config::base_960h();
|
||
let vm = VarMap::new();
|
||
let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
|
||
let fe = FeatureExtractor::new(&cfg, vb).expect("build");
|
||
let xs = Tensor::zeros((1, 1, 16_000), DType::F32, &dev).unwrap();
|
||
let ys = <FeatureExtractor as Module>::forward(&fe, &xs).expect("forward");
|
||
let dims = ys.dims();
|
||
assert_eq!(dims.len(), 3);
|
||
assert_eq!(dims[1], 512);
|
||
// Same conv arithmetic as emotion2vec (identical conv stack).
|
||
let ideal = 16_000 / cfg.conv_stride_total();
|
||
assert!(dims[2] <= ideal && dims[2] >= ideal.saturating_sub(2));
|
||
}
|
||
|
||
#[test]
|
||
fn conv_pos_embedding_preserves_t_with_even_kernel() {
|
||
let dev = Device::Cpu;
|
||
let vm = VarMap::new();
|
||
let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
|
||
// Kernel 128, groups 16 — exact `_base` config.
|
||
let pos = ConvPosEmbedding::new(768, 128, 16, vb).expect("build");
|
||
let xs = Tensor::zeros((1, 50, 768), DType::F32, &dev).unwrap();
|
||
let ys = <ConvPosEmbedding as Module>::forward(&pos, &xs).expect("forward");
|
||
// Same-padding via even-kernel-trim must keep T=50 exactly.
|
||
assert_eq!(ys.dims(), &[1, 50, 768]);
|
||
}
|
||
|
||
#[test]
|
||
fn block_post_norm_random_init_shape() {
|
||
let dev = Device::Cpu;
|
||
let cfg = Wav2Vec2Config::base_960h();
|
||
let vm = VarMap::new();
|
||
let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
|
||
let blk =
|
||
Block::new(cfg.embed_dim, cfg.num_heads, cfg.mlp_dim, cfg.norm_eps, vb).expect("build");
|
||
let xs = Tensor::randn(0f32, 1.0, (1, 16, cfg.embed_dim), &dev).unwrap();
|
||
let ys = <Block as Module>::forward(&blk, &xs).expect("forward");
|
||
assert_eq!(ys.dims(), &[1, 16, cfg.embed_dim]);
|
||
let v = ys.flatten_all().unwrap().to_vec1::<f32>().unwrap();
|
||
assert!(v.iter().all(|x| x.is_finite()));
|
||
}
|
||
|
||
#[test]
|
||
fn wav2vec2_random_init_end_to_end_shape() {
|
||
let dev = Device::Cpu;
|
||
let cfg = Wav2Vec2Config::base_960h();
|
||
let vm = VarMap::new();
|
||
let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
|
||
let model = Wav2Vec2::new(cfg.clone(), vb, dev.clone()).expect("build");
|
||
// 1 s of 16 kHz audio → expected output (1, ~50, 32).
|
||
let audio = Tensor::randn(0f32, 0.1, (1, 1, 16_000), &dev).unwrap();
|
||
let logits = model.forward(&audio).expect("forward");
|
||
let dims = logits.dims();
|
||
assert_eq!(dims.len(), 3);
|
||
assert_eq!(dims[0], 1);
|
||
assert_eq!(dims[2], cfg.vocab_size);
|
||
assert!(dims[1] >= 48 && dims[1] <= 52);
|
||
let v = logits.flatten_all().unwrap().to_vec1::<f32>().unwrap();
|
||
assert!(v.iter().all(|x| x.is_finite()));
|
||
}
|
||
|
||
#[test]
|
||
fn transcript_to_token_ids_handles_spaces_and_unknowns() {
|
||
let ids = transcript_to_token_ids("HE LL", VOCAB_960H).unwrap();
|
||
// H=11 E=5 |=4 L=15 L=15
|
||
assert_eq!(ids, vec![11, 5, 4, 15, 15]);
|
||
|
||
// Unknown char → <unk>=3
|
||
let ids = transcript_to_token_ids("@", VOCAB_960H).unwrap();
|
||
assert_eq!(ids, vec![3]);
|
||
|
||
// Multiple spaces collapse to a single separator.
|
||
let ids = transcript_to_token_ids("A B", VOCAB_960H).unwrap();
|
||
assert_eq!(ids, vec![7, 4, 24]);
|
||
}
|
||
|
||
#[test]
|
||
fn viterbi_align_recovers_obvious_alignment() {
|
||
let dev = Device::Cpu;
|
||
let v = 32usize;
|
||
// Synthesize log-probs so that the optimal alignment of "HE"
|
||
// (tokens [11, 5]) is unambiguous. Frames: [pad pad H H E E pad].
|
||
// We give a strong score to the intended state at each frame and
|
||
// -10 elsewhere — small enough to keep blank fallback feasible.
|
||
let layout: [u32; 7] = [0, 0, 11, 11, 5, 5, 0];
|
||
let mut data = vec![-10.0f32; layout.len() * v];
|
||
for (t, id) in layout.iter().enumerate() {
|
||
data[t * v + *id as usize] = 0.0;
|
||
}
|
||
let lp = Tensor::from_vec(data, (layout.len(), v), &dev).unwrap();
|
||
let aligned = viterbi_align(&lp, &[11, 5], CTC_BLANK_ID, VOCAB_960H).unwrap();
|
||
assert_eq!(aligned.len(), 2);
|
||
assert_eq!(aligned[0].token, "H");
|
||
assert_eq!(aligned[0].frame_start, 2);
|
||
assert_eq!(aligned[0].frame_end, 3);
|
||
assert_eq!(aligned[1].token, "E");
|
||
assert_eq!(aligned[1].frame_start, 4);
|
||
assert_eq!(aligned[1].frame_end, 5);
|
||
}
|
||
|
||
#[test]
|
||
fn group_into_words_splits_on_separator() {
|
||
let toks = vec![
|
||
AlignedToken {
|
||
token: "H".into(),
|
||
frame_start: 0,
|
||
frame_end: 1,
|
||
},
|
||
AlignedToken {
|
||
token: "I".into(),
|
||
frame_start: 2,
|
||
frame_end: 3,
|
||
},
|
||
AlignedToken {
|
||
token: "|".into(),
|
||
frame_start: 4,
|
||
frame_end: 5,
|
||
},
|
||
AlignedToken {
|
||
token: "Y".into(),
|
||
frame_start: 6,
|
||
frame_end: 6,
|
||
},
|
||
AlignedToken {
|
||
token: "O".into(),
|
||
frame_start: 7,
|
||
frame_end: 8,
|
||
},
|
||
];
|
||
let words = group_into_words(&toks);
|
||
assert_eq!(words.len(), 2);
|
||
assert_eq!(words[0].word, "HI");
|
||
assert_eq!(words[0].frame_start, 0);
|
||
assert_eq!(words[0].frame_end, 3);
|
||
assert_eq!(words[1].word, "YO");
|
||
assert_eq!(words[1].frame_start, 6);
|
||
assert_eq!(words[1].frame_end, 8);
|
||
}
|
||
|
||
#[test]
|
||
fn frame_to_ms_50hz_grid() {
|
||
// 50 Hz = 20 ms per frame.
|
||
assert_eq!(frame_to_ms(0), 0.0);
|
||
assert_eq!(frame_to_ms(50), 1000.0);
|
||
}
|
||
|
||
#[test]
|
||
fn ctc_greedy_decode_collapses_repeats_and_drops_blanks() {
|
||
let dev = Device::Cpu;
|
||
// Synthesize logits such that argmax sequence is:
|
||
// [pad pad H H E L L L O pad O pad pad |]
|
||
// (pad=0, H=11, E=5, L=15, O=8, |=4)
|
||
//
|
||
// Standard CTC collapse:
|
||
// - blanks (pad) drop out
|
||
// - consecutive same-id with NO blank between → collapsed
|
||
// - same-id with a blank between → both emitted
|
||
//
|
||
// So:
|
||
// H H → H ; E ; L L L → L ; O ; (blank) ; O ; (blanks) ; | → space
|
||
// final: "HELOO " (one H, one E, one L, two O's separated by
|
||
// blank, trailing space from | word-separator)
|
||
let argmax_seq: [u32; 14] = [0, 0, 11, 11, 5, 15, 15, 15, 8, 0, 8, 0, 0, 4];
|
||
let v = 32usize;
|
||
let mut data = vec![0.0f32; argmax_seq.len() * v];
|
||
for (t, id) in argmax_seq.iter().enumerate() {
|
||
data[t * v + *id as usize] = 1.0;
|
||
}
|
||
let logits = Tensor::from_vec(data, (1, argmax_seq.len(), v), &dev).unwrap();
|
||
let text = ctc_greedy_decode(&logits, VOCAB_960H).unwrap();
|
||
assert_eq!(text, "HELOO ");
|
||
}
|
||
}
|