espnet2.speechlm.tokenizer.codec_tokenizer.CodecTokenizer

source

class espnet2.speechlm.tokenizer.codec_tokenizer.CodecTokenizer(codec_choice: str, codec_fs: int, device: str = 'cpu', dump_audio: bool = False, checkpoint_path: str | None = None, config_path: str | None = None, max_token_per_frame: int = 32)

Bases: AbsTokenizer

Codec Tokenizer implementation

Use cases: : - use encode and decode for discrete (de)tokenization

• use encode_continuous and decode_continuous for continuous (de)tokenization
• use forward and detokenization for discrete (de)tokenization with flatten sequence style, which is more friendly for speechlm task

Codec Tokenizer initialization

Each of the codec implementation should assign all following features: : self.n_codebook (int): the number of codec codebooks. self.size_codebook (int): the dimension of codebooks. self.sample_rate (int): the sample rate the model trained on. self.subsample (int): the subsample rate, a.k.a., frame shift.

decode(codes)

Recover the waveform from the codes. Input:

codes (torch.Tensor): Int tensor in shape [B, T, n_codebook]

Output: : waveform (torch.Tensor): float tensor in shape [B, n_sample]

decode_continuous(z)

Recover the waveform from the continuous representations of codec Input:

z (torch.Tensor): Float tensor in shape [B, T, D], codec : continuous representations

Output: : waveform (torch.Tensor): float tensor in shape [B, n_sample]

detokenize(codes, n_codebook=None)

Convert flatten codec codes into resynthesis the audio Input:

codes (torch.Tensor): int tensor in shape [B, T * n_codebook], : or [T * n_codebook]

Output: : waveform (torch.Tensor): float tensor in shape [B, n_sample], : or [n_sample]

encode(wavs)

Convert audio waveforms into codec codes Input:

wavs (torch.Tensor): float tensor in shape [B, 1, n_sample],

Output: : codes (torch.Tensor): Int tensor in shape [B, T, n_codebook]

encode_continuous(wavs)

Convert audio waveforms into continuous codec encoding results Input:

wavs (torch.Tensor): float tensor in shape [B, 1, n_sample],

Output: : z (torch.Tensor): float tensor in shape [B, T, D]

forward(wavs)

Convert audio waveforms into flatten codec codes and resynthesis the audio Input:

wavs (torch.Tensor): float tensor in shape [B, 1, n_sample],

Output: : codes (torch.Tensor): Int tensor in shape [B, T * n_codebook], resyn_audio (torch.Tensor): float tensor in shape [B, n_samples]