prj-beatrice/utmos22-torch-native

UTMOS22 strong learner

The UTMOS22 strong learner, ported to use only torch, torchaudio, and transformers, with reference to the SpeechMOS implementation.

import soundfile as sf
import torch
from transformers import AutoModel

device = "cuda" if torch.cuda.is_available() else "cpu"
wav, sr = sf.read("audio.wav", dtype="float32")
wav = torch.from_numpy(wav)
model = AutoModel.from_pretrained(
    "prj-beatrice/utmos22-torch-native", trust_remote_code=True
).eval().to(device)

with torch.inference_mode():
    score = model(wav.to(device), sampling_rate=sr).scores[0]
print(score)

The result matches this code:

reference_model = torch.hub.load(
    "tarepan/SpeechMOS:v1.2.0", "utmos22_strong", trust_repo=True
)
reference = reference_model(wav.unsqueeze(0), sr)[0]
torch.testing.assert_close(score.cpu(), reference)

Padded batches use a waveform tensor [batch_size, max_length] and an input_lengths tensor [batch_size].

License

The model and wrapper are MIT-licensed; see LICENSE.