walston/whisaid-medium-uniform

WhisAID Medium Uniform

WhisAID Mandarin accent classifier and accent encoder based on Whisper Medium. This model was trained with a uniform speaker-distribution adversarial loss (uniform MSE, alpha 10) for 10 epochs. The released checkpoint is epoch 9.

Configuration

• Base model: Whisper Medium
• Mel bins: 80
• Accent classes: 9
• Training speakers: 336
• Uniform loss alpha: 10

Usage

Install the WhisAID source package and its dependencies, then load the model:

import torch
from transformers import AutoModel
from whisper import load_audio, log_mel_spectrogram, pad_or_trim
from whisAID import WhisAIDConfig

repo_id = "walston/whisaid-medium-uniform"
model = AutoModel.from_config(
    WhisAIDConfig.from_pretrained(repo_id)
).cuda().eval()

audio = torch.from_numpy(load_audio("/path/to/audio.wav"))
mel = log_mel_spectrogram(
    pad_or_trim(audio), n_mels=model.config.n_mels
).unsqueeze(0).cuda()

with torch.no_grad():
    output = model(input_ids=mel)

accent_embedding = output.features[0].cpu().numpy()
accent_id = output.logits.argmax(dim=-1).item()

The model repository stores the Lightning checkpoint. Model integration code is provided by the WhisAID project rather than through remote Hub code.