espnet2.tok.asr_loss.ASRLoss

source

class espnet2.tok.asr_loss.ASRLoss(asr_model: ESPnetASRModel, embedding: TokenEmbedding)

Bases: Module

Apply an ESPnet ASR model to embedded differentiable speech tokens.

This loss reuses the joint CTC/attention ASR components while leaving SSL feature extraction and tokenization under the shared tokenizer.

The wrapped ASR model must be built with frontend=None and with its encoder input size equal to embedding.output_size(). The loss embeds hard straight-through token assignments before calling the standard ESPnetASRModel; no modification to the generic ASR model is needed.

• Parameters:
  • asr_model – Standard ESPnet ASR model containing the encoder, CTC, and optional attention or transducer decoder.
  • embedding – Loss-specific token embedding.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

encode(tokenizer_output: SpeechTokenizerOutput) → Tuple[Tensor, Tensor]

Encode shared tokens for ASR inference without computing a loss.

forward(tokenizer_output: SpeechTokenizerOutput, text: Tensor, text_lengths: Tensor, speech: Tensor | None = None, speech_lengths: Tensor | None = None, **kwargs: Tensor) → Tuple[Tensor, Dict[str, Tensor], Tensor]

Compute the standard ESPnet ASR loss from differentiable tokens.

• Parameters:
  • tokenizer_output – Shared tokenizer output. Its hard straight-through assignment is used rather than detached integer token IDs.
  • text – Reference token IDs of shape (B, L).
  • text_lengths – Valid reference lengths of shape (B,).
  • speech – Original waveform, unused by the ASR loss.
  • speech_lengths – Original waveform lengths, unused by the ASR loss.
  • kwargs – Optional auxiliary inputs consumed by ESPnetASRModel, such as intermediate-CTC targets.
• Returns: The (loss, stats, weight) tuple returned by ESPnetASRModel.