Update README.md

README.md

@@ -2,20 +2,21 @@
 library_name: transformers
 license: other
 license_name: meralion-public-license
-license_link: https://huggingface.co/MERaLiON/MERaLiON-AudioLLM-Whisper-SEA-LION/blob/main/MERaLiON-Public-Licence-v1.pdf
+license_link: >-
+  https://huggingface.co/MERaLiON/MERaLiON-AudioLLM-Whisper-SEA-LION/blob/main/MERaLiON-Public-Licence-v1.pdf
 tags:
-- speech
-- best-rq
-- meralion
-- meralion-2
+  - speech
+  - best-rq
+  - meralion
+  - meralion-2
 language:
-- en
-- zh
-- ms
-- ta
-- id
-- th
-- vi
+  - en
+  - zh
+  - ms
+  - ta
+  - id
+  - th
+  - vi
 ---
 
 <h1 align="center">🎧 MERaLiON-SpeechEncoder-2 🎧</h1>
@@ -44,10 +45,10 @@ With only 630M parameters (≈2.5 GB in memory), the model is easily deployable
 Building on [MERaLiON-SpeechEncoder-v1](https://huggingface.co/MERaLiON/MERaLiON-SpeechEncoder-v1) (which focused on English and Singlish), this version expands to include English, Chinese, Malay, Tamil, Thai, Indonesian, and Vietnamese, along with codeswitching support across these languages. Given the wide coverage of languages in the training corpus, it may also be applicable beyond the officially supported languages.
   
 #### Competitive performance on downstream speech tasks
-The model retains near state-of-the-art results on the SUPERB benchmark for English, and showcases strong multilingual capabilities deomnstrated through its integration into a [high-performance ASR system shown below](#automatic-speech-recognition).
+The model retains near state-of-the-art results on the SUPERB benchmark for English, and showcases strong multilingual capabilities deomnstrated through its integration into a [high-performance ASR system shown below](#automatic-speech-recognition-asr).
 
 #### Innovative pre-training techniques
-MERaLiON-SpeechEncoder-2 was trained from scratch with an novel extension of the BEST-RQ self-supervised objective, by using more informative latent targets. We also adopted the Muon optimizer, which has previously only been shown to outperform the popular AdamW for LLM training. We find its advantages also carry over to speech-based models. 
+MERaLiON-SpeechEncoder-2 was trained from scratch with a novel extension of the BEST-RQ self-supervised objective, by using more informative latent targets. We also adopted the Muon optimizer, which has previously only been shown to outperform the popular AdamW for LLM training. We find its advantages also carry over to speech-based models. 
  
 ## Model Summary
 
@@ -65,7 +66,7 @@ For details on background, pre-training, tuning experiments and evaluation, plea
 |----------------------------------|---------------|------|------|-------|--------|-------|------|------|-------|-------|----------------------|
 | HuBERT large                    | 82.25         | 3.53 | 3.62 | 95.29 | 0.0354 | 90.33 | 5.98 | 5.75 | 67.62 | 98.76 | 89.91 / 21.76        |
 | WavLM large                     | 84.77         | 3.06 | 3.44 | 97.86 | 0.0886 | 95.49 | 3.77 | 3.24 | 70.62 | 99.31 | 92.21 / 18.36        |
-| MERaLiON-SpeechEncoder-V1       | 82.62         | 3.14 | 4.16 | 97.63 | 0.0590 | 91.09 | 5.18 | 5.06 | 68.02 | 98.60 | 88.99 / 23.89        |
+| MERaLiON-SpeechEncoder-v1       | 82.62         | 3.14 | 4.16 | 97.63 | 0.0590 | 91.09 | 5.18 | 5.06 | 68.02 | 98.60 | 88.99 / 23.89        |
 | MERaLiON-SpeechEncoder-2        | 82.72         | 3.40 | 4.96 | 97.57 | 0.0575 | 88.96 | 3.93 | 3.90 | 68.80 | 98.95 | 89.50 / 23.46        |
 
 SUPERB is an English-based benchmark for speech encoders covering a wide range of downstream speech tasks across domains such as recognition, detection, semantics, speaker, and paralinguistics, where each task is finetuned separately with a frozen encoder.
@@ -76,14 +77,135 @@ MERaLiON-SpeechEncoder-2 is competitive to state-of-the-art, improving slightly
 ### Automatic Speech Recognition (ASR)
 
 <p align="center">
-  <img src="overall_wer.svg" width="650"/> 
-  <img src="audiobench_wer.svg" width="650"/> 
-  <img src="fleurs_wer.svg" width="650"/>
+  <img src="overall_wer.svg" width="680"/> 
+  <img src="audiobench_wer.svg" width="680"/> 
+  <img src="fleurs_wer.svg" width="680"/>
 </p>
 
-Leveraging on the multilingual capabilities of MERaLiON-SpeechEncoder-2, we further finetuned the model for on supervised speech data to produce a lightweight MERaLiON-SpeechEncoder-2-ASR-CTC, which is competitive to models many times its size in transcribing the target languages, while offering much faster inference speeds. It outperforms the popular Whisper large v3 across most ASR benchmarks, including [Audiobench](https://huggingface.co/spaces/MERaLiON/AudioBench-Leaderboard) and FLEURS. Our internal benchmarking, shown in the 'Overall ASR Performance', also contains several private datasets in addition to Audiobench and FLEURS.  
+Leveraging on the multilingual capabilities of MERaLiON-SpeechEncoder-2, we further finetuned the model for on supervised speech data to produce a lightweight MERaLiON-SpeechEncoder-2-ASR-CTC, which is competitive to models many times its size in transcribing the target languages, while offering much faster inference speeds. It outperforms the popular Whisper large v3 across most languages in [Audiobench](https://huggingface.co/spaces/MERaLiON/AudioBench-Leaderboard) and maintains close perofrmance in FLEURS. Our internal benchmarking, shown in the 'Overall ASR Performance', also contains several private datasets in addition to Audiobench and FLEURS.  
 
+### Direct Use
 
+The follwing code snippet can be used to directly obtain latent features i.e. encoded speech by forwarding through the model. 
+
+```python
+import torch
+from datasets import load_dataset
+from transformers import AutoModel, AutoFeatureExtractor
+
+repo_id = 'MERaLiON/MERaLiON-SpeechEncoder-2'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# load model and feature extractor
+model = AutoModel.from_pretrained(
+    repo_id,
+    trust_remote_code=True,
+)
+model = model.to(device)
+
+feature_extractor = AutoFeatureExtractor.from_pretrained(
+    repo_id,
+    trust_remote_code=True
+)
+
+# prepare data
+data = load_dataset("distil-whisper/librispeech_long", "clean",
+                split="validation")
+
+def batch_collater(data):
+    tensors = []
+    for idx, sample in enumerate(data):
+        tensors.append(sample['audio']['array'])
+    return tensors
+
+audio_array = batch_collater(data)
+inputs = feature_extractor(audio_array, sampling_rate=16_000,
+                        return_attention_mask=True,
+                        return_tensors='pt', do_normalize=False)
+inputs = inputs.to(device)
+
+# model inference to obtain features
+with torch.no_grad():
+    model.eval()
+    output = model(input_values=inputs['input_values'],
+                attention_mask=inputs['attention_mask'],
+                output_hidden_states=True)
+```
+
+### Downstream Use
+
+Speech encoders are normally used in finetuning setups to provide the frontend to downstream speech applications. We provide an example below of an ASR finetuning setup with Huggingface. Please refer to this [blog](https://huggingface.co/blog/fine-tune-w2v2-bert) for the full ASR finetuning recipe with Huggingface Trainer. Alternatively, the Huggingface model can be loaded to any other frameworks such as Pytorch or ESPnet for custom finetuning loops.
+
+```python
+import torch
+import json
+from datasets import load_dataset
+from transformers import AutoModelForCTC, AutoFeatureExtractor, Wav2Vec2CTCTokenizer
+
+repo_id = 'MERaLiON/MERaLiON-SpeechEncoder-2'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# prepare data
+def pre_processing(batch):
+    batch["text"] = batch["text"].lower()
+    return batch
+
+def extract_all_chars(batch):
+  all_text = " ".join(batch["text"])
+  vocab = list(set(all_text))
+  return {"vocab": [vocab], "all_text": [all_text]}
+
+librispeech100h_train = load_dataset("openslr/librispeech_asr", split="train.clean.100")
+librispeech100h_test = load_dataset("openslr/librispeech_asr", split="validation.clean")
+librispeech100h_train = librispeech100h_train.remove_columns(
+                                    ['file', 'speaker_id', 'chapter_id', 'id'])
+librispeech100h_test = librispeech100h_test.remove_columns(
+                                    ['file', 'speaker_id', 'chapter_id', 'id'])
+
+librispeech100h_train = librispeech100h_train.map(pre_processing)
+librispeech100h_test = librispeech100h_test.map(pre_processing)
+
+vocab_train = librispeech100h_train.map(extract_all_chars, batched=True,
+                                    batch_size=-1, keep_in_memory=True,
+                                    remove_columns=librispeech100h_train.column_names)
+vocab_test = librispeech100h_test.map(extract_all_chars, batched=True,
+                                    batch_size=-1, keep_in_memory=True,
+                                    remove_columns=librispeech100h_test.column_names)
+vocab_list = list(set(vocab_train["vocab"][0]) | set(vocab_test["vocab"][0]))
+vocab_dict = {v: k for k, v in enumerate(sorted(vocab_list))}
+
+vocab_dict["|"] = vocab_dict[" "]
+del vocab_dict[" "]
+vocab_dict["[UNK]"] = len(vocab_dict)
+vocab_dict["[PAD]"] = len(vocab_dict)
+
+with open('ls_vocab.json', 'w') as vocab_file:
+    json.dump(vocab_dict, vocab_file)
+
+# load model, feature extractor and tokenizer
+feature_extractor = AutoFeatureExtractor.from_pretrained(
+    repo_id,
+    trust_remote_code = True,
+)
+
+tokenizer = Wav2Vec2CTCTokenizer("./ls_vocab.json",
+                            unk_token="[UNK]", pad_token="[PAD]",
+                            word_delimiter_token="|")
+
+model = AutoModelForCTC.from_pretrained(
+    repo_id,
+    trust_remote_code=True,
+    vocab_size=len(vocab_dict),
+    feat_proj_dropout=0.1,
+    activation_dropout=0.1,
+    hidden_dropout=0.1,
+    conformer_conv_dropout=0.1,
+    ctc_loss_reduction="mean",
+    pad_token_id=tokenizer.pad_token_id,
+    attention_dropout=0.1,
+)
+model = model.to(device)
+```