How to apply SpecAugment to a Whisper?

🤗Transformers
3

I’m not if this is the right way to do it but it seem to work:

from transformers import WhisperForConditionalGeneration

model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")

model.config.forced_decoder_ids = None
model.config.suppress_tokens = []

model.config.apply_spec_augment = True
model.config.mask_time_prob = 0.05
model.config.mask_feature_prob = 0.05

The last three lines enable the SpecAug for the Whisper model. The right values to set need to be researched.