speechbrain.lobes.models.huggingface_transformers.discrete_ssl 模块

该叶节点支持集成预训练的离散SSL（hubert、wavlm、wav2vec），用于从SSL层的输出中提取语义标记。

需要安装来自HuggingFace的Transformer： https://huggingface.co/transformers/installation.html

Author

Pooneh Mousavi 2024
贾罗德·杜雷 2024

摘要

类：

DiscreteSSL

该模块支持集成HuggingFace和SpeechBrain预训练的Discrete SSL模型。

参考

class speechbrain.lobes.models.huggingface_transformers.discrete_ssl.DiscreteSSL(save_path, ssl_model, kmeans_dataset, vocoder_repo_id='speechbrain/hifigan-wavlm-k1000-LibriTTS', num_clusters=1000, layers_num=None, device='cpu', sample_rate=16000)[source]

基础：Module

该模块支持集成HuggingFace和SpeechBrain预训练的离散SSL模型。

需要安装来自HuggingFace的Transformer： https://huggingface.co/transformers/installation.html

该模型可以用作固定的离散特征提取器，也可以进行微调。它将自动从HuggingFace下载模型或使用本地路径。

下表总结了兼容的SSL模型、它们各自的HF编码器、k-means训练细节、支持的层以及预训练的声码器：

|------------|—————————————-|-----------------|————–|----------------------|———————————————| | WavLM | microsoft/wavlm-large | LibriSpeech960 | 1000 | 1, 3, 7, 12, 18, 23 | speechbrain/hifigan-wavlm-k1000-LibriTTS | | HuBERT | facebook/hubert-large-ll60k | LibriSpeech960 | 1000 | 1, 3, 7, 12, 18, 23 | speechbrain/hifigan-hubert-k1000-LibriTTS | | Wav2Vec2 | facebook/wav2vec2-large | LibriSpeech960 | 1000 | 1, 3, 7, 12, 18, 23 | speechbrain/hifigan-wav2vec2-k1000-LibriTTS |

Parameters:

save_path (str) – 下载模型的路径（目录）。
ssl_model (str) – 用于从其层输出中提取语义标记的SSL模型。请注意，应设置output_all_hiddens为True以启用多层离散化。
kmeans_dataset (str) – Kmeans模型在HF仓库上训练所使用的数据集名称。
vocoder_repo_id (str) – 包含预训练HiFi-GAN模型的Huggingface仓库。
num_clusters (int 或 List[int] (默认值: 1000)) – 确定要下载的目标kmeans模型的聚类数量。每一层的数量可能不同。
layers_num (List[int] (Optional)) – 确定从HF仓库下载的层。如果未提供，则从HF仓库加载所有具有num_clusters(int)的层。如果num_clusters是一个列表，则应提供layers_num以确定每层的聚类数。
device (str (默认 'cpu')) – 用于计算的设备（‘cpu’ 或 ‘cuda’）。
sample_rate (int (默认值: 16000)) – 输入音频的采样率。

Example

>>> import torch
>>> from speechbrain.lobes.models.huggingface_transformers.wavlm import (WavLM)
>>> inputs = torch.rand([3, 2000])
>>> model_hub = "microsoft/wavlm-large"
>>> save_path = "savedir"
>>> ssl_layer_num = [7,23]
>>> deduplicate =[False, True]
>>> bpe_tokenizers=[None, None]
>>> vocoder_repo_id = "speechbrain/hifigan-wavlm-k1000-LibriTTS"
>>> kmeans_dataset = "LibriSpeech"
>>> num_clusters = 1000
>>> ssl_model = WavLM(model_hub, save_path,output_all_hiddens=True)
>>> model = DiscreteSSL(save_path, ssl_model, vocoder_repo_id=vocoder_repo_id, kmeans_dataset=kmeans_dataset,num_clusters=num_clusters)
>>> tokens, _, _ = model.encode(inputs,SSL_layers=ssl_layer_num, deduplicates=deduplicate, bpe_tokenizers=bpe_tokenizers)
>>> print(tokens.shape)
torch.Size([3, 6, 2])
>>> sig = model.decode(tokens, ssl_layer_num)
>>> print(sig.shape)
torch.Size([3, 1, 1920])

check_if_input_is_compatible(layers_num, num_clusters)[source]

检查 layer_number 和 num_clusters 是否相互一致。

Parameters:

layers_num (List[int] (Optional)) – 如果num_clusters是一个列表，应提供layers_num以确定每层的聚类数量。
num_clusters (int 或 List[int]) – 确定要下载的目标kmeans模型的聚类数量。每一层的数量可能不同。

load_kmeans(repo_id, kmeans_dataset, encoder_name, num_clusters, cache_dir, layers_num=None)[source]

从HF加载一个预训练的kmeans模型。

Parameters:

repo_id (str) – 包含模型的hugingface仓库id。
kmeans_dataset (str) – 需要在HF仓库中下载的用于训练Kmeans模型的数据集名称。
encoder_name (str) – 用于定位文件的编码器名称。
num_clusters (int 或 List[int]) – 确定要下载的目标kmeans模型的聚类数量。每一层的数量可能不同。
cache_dir (str) – 下载模型的路径（目录）。
layers_num (List[int] (Optional)) – 如果 num_clusters 是一个列表，则应提供 layers_num 以确定每层的聚类数量。

Returns:

kmeans_model (MiniBatchKMeans) – 从HF加载的预训练Kmeans模型。
layer_ids (List[int]) – 支持用于kmeans的层数（从kmeans模型的名称中提取）。

forward(wav, wav_lens=None, SSL_layers=None, deduplicates=None, bpe_tokenizers=None)[source]

接收输入波形并返回其对应的标记和重建信号。

Parameters:

wav (torch.Tensor (signal)) – 一批要转换为特征的音频信号。
wav_lens (tensor) – 以SpeechBrain格式给出的wav的相对长度。
SSL_layers (List[int]:) – 确定应使用SSL的哪些层来提取信息。
deduplicates (List[boolean]:) – 确定是否对提取的相应层的标记应用去重（移除重复的后续标记）。
bpe_tokenizers (List[int]:) – 如果为该层训练了sentencePiece分词器，则确定是否对为该层提取的标记应用子词处理。

Returns:

tokens (torch.Tensor) – 一个 (Batch x Seq x num_SSL_layers) 的音频令牌张量
waveforms (torch.tensor) – 一批梅尔波形 [batch, time]

encode(wav, wav_lens=None, SSL_layers=None, deduplicates=None, bpe_tokenizers=None)[source]

获取输入波形并返回其对应的编码。

Parameters:

wav (torch.Tensor (signal)) – 一批要转换为特征的音频信号。
wav_lens (tensor) – 以SpeechBrain格式给出的wav的相对长度。
SSL_layers (List[int]:) – 确定应使用SSL的哪些层来提取信息。
deduplicates (List[boolean]:) – 确定是否对提取的相应层的标记应用去重（移除重复的后续标记）。
bpe_tokenizers (List[int]:) – 如果为该层训练了sentencePiece分词器，则确定是否对为该层提取的标记应用子词处理。

Returns:

tokens (torch.Tensor) – 一个 (Batch x Seq x num_SSL_layers) 的音频令牌张量
emb (torch.Tensor) – 一个 (Batch x Seq x num_SSL_layers x embedding_dim ) 的每个令牌的聚类中心嵌入
processed_tokens (torch.Tensor) – 一个 (Batch x Seq x num_SSL_layers) 的音频令牌张量，在必要时应用去重和子词处理后的结果。

decode(tokens, SSL_layers=None)[source]

接收一个输入波形并返回其对应的波形。原始来源： https://github.com/speechbrain/benchmarks/blob/c87beb61d4747909a133d3e1b3a3df7c8eda1f08/ benchmarks/DASB/Libri2Mix/separation/conformer/train_discrete_ssl.py#L44

Parameters:

tokens (torch.Tensor) – 一个 (Batch, codes, layers) 的张量，包含离散单元
SSL_layers (List[int]:) – 确定声码器应使用哪些SSL层。

Returns:

waveforms – 一批梅尔波形 [批次, 时间]

Return type:

torch.tensor