walston/whisaid-medium-grl

WhisAID Medium GRL

WhisAID Mandarin accent classifier and accent encoder based on Whisper Medium. This checkpoint was trained with a gradient reversal layer (GRL, lambda 0.05) for 10 epochs. The released checkpoint is epoch 9.

Configuration

• Base model: Whisper Medium
• Mel bins: 80
• Accent classes: 9
• Training speakers: 336
• GRL lambda: 0.05

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-grl"
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.