MU-NLPC / whisper-small-audio-captioning

Model Card for Whisper Audio Captioning

A transformer encoder-decoder model for automatic audio captioning. As opposed to speech-to-text, captioning describes the content of audio clips, such as prominent sounds or environmental noises. This task has numerous practical applications, e.g., for providing access to audio information for people with hearing impairments or improving the searchability of audio content.

• Model type: Whisper encoder-decoder transformer

• Language(s) (NLP): en

• License: cc-by-4.0

• Parent Model: openai/whisper-small

• Resources for more information:

  • GitHub Repo
  • Technical Report

Usage

The model expects an audio clip (up to 30s) to the encoder as an input and information about caption style as forced prefix to the decoder.

Minimal example:

# Load model
checkpoint = "MU-NLPC/whisper-small-audio-captioning"
model = WhisperForAudioCaptioning.from_pretrained(checkpoint)
tokenizer = transformers.WhisperTokenizer.from_pretrained(checkpoint, language="en", task="transcribe")
feature_extractor = transformers.WhisperFeatureExtractor.from_pretrained(checkpoint)

# Load and preprocess audio
input_file = "..."
audio, sampling_rate = librosa.load(input_file, sr=feature_extractor.sampling_rate)
features = feature_extractor(audio, sampling_rate=sampling_rate, return_tensors="pt").input_features

# Prepare caption style
style_prefix = "clotho > caption: "
style_prefix_tokens = tokenizer("", text_target=style_prefix, return_tensors="pt", add_special_tokens=False).labels

# Generate caption
model.eval()
outputs = model.generate(
    inputs=features.to(model.device),
    forced_ac_decoder_ids=style_prefix_tokens,
    max_length=100,
)

print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])

Example output: clotho > caption: Rain is pouring down and thunder is rumbling in the background.

The style prefix influences the style of the caption. Model knows 3 styles: audioset > keywords: , audiocaps > caption: , and clotho > caption: . It was finetuned on Clotho and that is the indended "default" style.

WhisperTokenizer must be initialized with language="en" and task="transcribe" .

Our model class WhisperForAudioCaptioning can be found in our git repository or here on the HuggingFace Hub in the model repository. The class overrides default Whisper generate method to support forcing decoder prefix.

Training details

The model was initialized by original speech-to-text openai/whisper-small weights. Then, it was pretrained on a mix of (1) subset of AudioSet with synthetic labels, (2) AudioCaps captioning dataset and (3) Clotho v2.1 captioning dataset. Finally, it was finetuned on Clotho v2.1 to focus the model on the specific style of the captions. For each traning input, the model was informed about the source of the data, so it can mimic the caption style in all 3 styles.

During pretraining, the ratio of samples in each batch was approximately 12:3:1 (AudioSet:AudioCaps:Clotho). The pretraining took 19800 steps with batch size 32 and learning rate 2e-5. Finetuning was done on Clotho only, and the model was trained for 1500 steps with batch size 32 and learning rate 4e-6. All layers except fc1 layers were frozen during finetuning.

For more information about the training regime, see the technical report .

Evaluation details

Metrics reported in the metadata were computed on Clotho v2.1 test split with captions generated using a beam search with 5 beams.

Limitations

The captions generated by the model can be misleading or not truthful, even if they appear convincing. The hallucination occurs especially in domains that were not present in the finetuning data.

While the original speech-to-text checkpoints by OpenAI were trained on multilingual data, our training contains only English captions, and therefore is not expected for the model to support other languages.

Licence

The model weights are published under non-commercial license CC BY-NC 4.0 as the model was finetuned on a dataset for non-commercial use.

Contact

If you'd like to chat about this, please get in touch with is via email at kadlcik <at> mail.muni.cz or ahajek <at> mail.muni.cz.