mstar.model.omnilingual_asr.components.wav2vec2

The Omnilingual ASR CTC network: a wav2vec 2.0 encoder plus a linear head.

waveform (N, S), normalized per utterance, 16 kHz

-> 7 strided convs, each LayerNorm (in fp32) + GELU (x320: 50 frames/s) -> LayerNorm, Linear 512 -> model_dim -> + GELU(grouped conv positional encoding) (kernel 128, 16 groups) -> pre-norm Transformer encoder, final LayerNorm -> Linear model_dim -> vocab (CTC logits)

Module and parameter names follow the fairseq2 checkpoint (encoder_frontend, encoder, final_proj), so the state dict loads with only the positional conv’s weight norm folded. Every module here is plain torch with no M* imports, so the parity harness can run it inside the reference’s own environment.

Batches are right-padded. seq_lens gives each row’s sample count; the padding never reaches a valid frame: the convolutions only look backwards from a valid output frame over valid samples, padded frames are zeroed before the positional convolution, and attention keys past a row’s length are masked.

Functions

attention_mask(frame_mask)

(N, 1, T, T) boolean mask: a valid frame attends to its row's valid frames only.

Classes

Encoder(cfg)
EncoderLayer(cfg)	Pre-norm Transformer layer (fairseq2 TransformerNormOrder.PRE).
FeatureExtractionLayer(in_dim, out_dim, ...)	Conv1d -> LayerNorm over channels (computed in fp32) -> GELU.
FeatureExtractor(layer_descs, eps)
FeedForward(model_dim, inner_dim)
Frontend(cfg)
PositionEncoder(model_dim, kernel_size, ...)	Adds GELU(conv(x)) with a grouped, even-kernel conv (its last output column dropped, so the output keeps the input length).
SelfAttention(model_dim, num_heads)
Wav2Vec2CTC(cfg)	Omnilingual ASR CTC model (fairseq2 Wav2Vec2AsrModel at inference).

class mstar.model.omnilingual_asr.components.wav2vec2.Encoder(cfg)

Bases: Module

forward(x, attn_mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• x (Tensor)

• attn_mask (Tensor | None)

Return type:

Tensor

class mstar.model.omnilingual_asr.components.wav2vec2.EncoderLayer(cfg)

Bases: Module

Pre-norm Transformer layer (fairseq2 TransformerNormOrder.PRE).

forward(x, attn_mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• x (Tensor)

• attn_mask (Tensor | None)

Return type:

Tensor

class mstar.model.omnilingual_asr.components.wav2vec2.FeatureExtractionLayer(in_dim, out_dim, kernel_size, stride, eps)

Bases: Module

Conv1d -> LayerNorm over channels (computed in fp32) -> GELU.

Parameters:

• in_dim (int)

• out_dim (int)

• kernel_size (int)

• stride (int)

• eps (float)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Tensor)

Return type:

Tensor

class mstar.model.omnilingual_asr.components.wav2vec2.FeatureExtractor(layer_descs, eps)

Bases: Module

Parameters:

• layer_descs (list[tuple[int, int, int]])

• eps (float)

forward(waveforms)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

waveforms (Tensor)

Return type:

Tensor

output_length(num_samples)

Frames a row yields: floor((len - kernel) / stride) + 1 per conv.

Parameters:

num_samples (int)

Return type:

int

class mstar.model.omnilingual_asr.components.wav2vec2.FeedForward(model_dim, inner_dim)

Bases: Module

Parameters:

• model_dim (int)

• inner_dim (int)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Tensor)

Return type:

Tensor

class mstar.model.omnilingual_asr.components.wav2vec2.Frontend(cfg)

Bases: Module

class mstar.model.omnilingual_asr.components.wav2vec2.PositionEncoder(model_dim, kernel_size, num_groups)

Bases: Module

Adds GELU(conv(x)) with a grouped, even-kernel conv (its last output column dropped, so the output keeps the input length). The checkpoint’s weight norm (weight_g/weight_v) is folded into conv.weight at load.

Parameters:

• model_dim (int)

• kernel_size (int)

• num_groups (int)

forward(x, frame_mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• x (Tensor)

• frame_mask (Tensor | None)

Return type:

Tensor

class mstar.model.omnilingual_asr.components.wav2vec2.SelfAttention(model_dim, num_heads)

Bases: Module

Parameters:

• model_dim (int)

• num_heads (int)

forward(x, attn_mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• x (Tensor)

• attn_mask (Tensor | None)

Return type:

Tensor

class mstar.model.omnilingual_asr.components.wav2vec2.Wav2Vec2CTC(cfg)

Bases: Module

Omnilingual ASR CTC model (fairseq2 Wav2Vec2AsrModel at inference).

forward(waveforms, seq_lens=None)

(N, S) waveforms (and each row’s sample count, for a padded batch) -> (N, T, vocab) logits and each row’s valid frame count.

The lengths are host ints, so a padded batch is masked without reading anything back from the device.

Parameters:

• waveforms (Tensor)

• seq_lens (list[int] | None)

Return type:

tuple[Tensor, list[int]]

frame_lengths(seq_lens)

Parameters:

seq_lens (list[int])

Return type:

list[int]

mstar.model.omnilingual_asr.components.wav2vec2.attention_mask(frame_mask)

(N, 1, T, T) boolean mask: a valid frame attends to its row’s valid frames only. A padded frame attends to every frame, which keeps its (discarded) output finite instead of a fully masked NaN row.

Parameters:

frame_mask (Tensor)

Return type:

Tensor