Update README.md

README.md

@@ -1,3 +1,132 @@
 ---
+datasets:
+- AudioSet
+- AudioCaps
+- Clotho-v2.1
+metrics:
+- 
+
+model-index:
+- name: whisper-small-audio-captioning
+  results:
+  - task:
+      type: audio-captioning
+      name: Audio Captioning
+    dataset:
+      type: clotho-v2.1
+      name: Clotho
+      split: evaluation
+    metrics:
+    - type: SPICE
+      value: 0.1234
+    - type: CIDEr
+      value: 0.4142
+    - type: SPIDEr
+      value: 0.2687
+    - type: METEOR
+      value: 0.3781
+    - type: SacreBLEU
+      value: 15.76
 license: cc-by-nc-4.0
+language:
+- en
 ---
+
+
+
+# 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-
+
+- **Resources for more information:** 
+    - [GitHub Repo](https://github.com/prompteus/audio-captioning)
+    - [Technical Report](TODO)
+
+
+## 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:
+
+```python3
+# Load model
+architecture = "openai/whisper-small"
+checkpoint = "MU-NLPC/whisper-small-audio-captioning"
+model = audiocap.WhisperForAudioCaptioning.from_pretrained(checkpoint)
+tokenizer = transformers.WhisperTokenizer.from_pretrained(checkpoint, language="en", task="transcribe")
+feature_extractor = transformers.WhisperFeatureExtractor.from_pretrained(architecture)
+
+# 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](TODO).
+
+
+## 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.
+
+|                      | whisper-tiny | whisper-small | whisper-large-v2 |
+|----------------------|--------------|---------------|------------------|
+| SacreBLEU            | 13.77        | 15.76         | 16.50            |
+| METEOR               | 0.3452       | 0.3781        | 0.3782           |
+| CIDEr                | 0.3404       | 0.4142        | 0.4331           |
+| SPICE                | 0.1077       | 0.1234        | 0.1257           |
+| SPIDEr               | 0.2240       | 0.2687        | 0.2794           |
+
+
+## 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.