espnet2.tok.espnet_model.ESPnetGANSpeechTokenizerModel

source

class espnet2.tok.espnet_model.ESPnetGANSpeechTokenizerModel(tokenizer: SpeechTokenizer, asr_loss: ASRLoss, reconstruction_loss: HiFiGANReconstructionLoss, reconstruction_weight: float = 0.1, extract_feats_in_collect_stats: bool = False)

Bases: AbsGANESPnetModel

Optimize a shared speech tokenizer with downstream losses.

The model combines ASR and speech reconstruction. GAN discriminator updates are exposed separately from the shared main loss.

The discriminator step reuses the waveform cached by the generator step of the same batch, so training must run generator-first. With discriminator-first order, the second batch finds the previous batch’s cache and raises an error.

Initialize a shared tokenizer with weighted downstream losses.

• Parameters:
  • tokenizer – Shared SSL and differentiable-k-means tokenizer.
  • asr_loss – Standard ESPnet ASR loss.
  • reconstruction_loss – Adversarial waveform reconstruction.
  • reconstruction_weight – Loss interpolation weight. ASR receives weight 1 - alpha and reconstruction receives alpha.
  • extract_feats_in_collect_stats – Whether collect-stats must build this model. Speech-tokenizer inputs need shapes only.

property asr_loss : [ASRLoss](ASRLoss.md#espnet2.tok.asr_loss.ASRLoss)

Return the ASR loss.

clear_generator_cache() → None

Discard cached generator outputs when the discriminator step is skipped.

collect_feats(**batch: Tensor) → Dict[str, Tensor]

Return no offline statistics; the SSL frontend runs inside the model.

encode_asr(speech: Tensor, speech_lengths: Tensor)

Tokenize speech and run the ASR encoder.

forward(speech: Tensor, speech_lengths: Tensor, spembs: Tensor, text: Tensor | None = None, text_lengths: Tensor | None = None, forward_generator: bool = True, **kwargs: Tensor) → Dict[str, Any]

Return main-generator or discriminator loss for its optimizer.

property reconstruction_loss : [HiFiGANReconstructionLoss](HiFiGANReconstructionLoss.md#espnet2.tok.hifigan_loss.HiFiGANReconstructionLoss)

Return the adversarial reconstruction loss.

synthesize(speech: Tensor, speech_lengths: Tensor, spembs: Tensor)

Tokenize speech and reconstruct waveform with speaker conditioning.

tokenize(speech: Tensor, speech_lengths: Tensor) → SpeechTokenizerOutput

Deterministically tokenize speech for inference.