balacoon/mhubert-147

This is a mHuBERT-147 model traced into a pytorch jit file. This allows to run and deploy the model without extra dependencies. Model after 2nd training iteration is combined with faiss index, to output hubert labels.

Usage:

import torch
import soundfile as sf

extractor = torch.jit.load("mhubert147_fp16_cuda.jit").to(torch.device("cuda"))
wav, sr = sf.read(
    "rms_arctic_a0001.wav",
    dtype="int16"
)
assert sr == 16000
x = torch.tensor(wav).unsqueeze(0).cuda()
x_len = torch.tensor([x.shape[1]]).cuda()
labels = traced_extractor(x, x_len)[1].detach().cpu().numpy()

More on tracing in the post