Update README.md

README.md

@@ -31,7 +31,7 @@ We introduce **MERaLiON-SpeechEncoder-2**, our next-generation multilingual spee
 Unlike many existing models optimized for high-resource, Western languages, MERaLiON-SpeechEncoder-2 is designed from the ground up to reflect the linguistic diversity and complexity of Southeast Asia. Our training data was curated to contain a substantial amount originating from Singapore and SEA, including 60,000 hours of Singapore-accented speech, with a further 160,000 hours covering Singapore’s official languages Chinese, Malay and Tamil, along with a smaller portion of dialects like Hokkien and Cantonese. SEA data amounts to 200,000 hours, including significant proportions of Malay, Thai, Indonesian, Vietnamese, with smaller amounts of Tagalog, Burmese, Javanese, Sundanese, Khmer and Lao. See below for a regional breakdown of the language coverage of our pre-training data.
 
 <p align="center">
-  <img src="data2.svg" width="620"/> 
+  <img src="data2.svg" width="600"/> 
 </p>
 
 ## Model Highlights
@@ -46,7 +46,7 @@ Building on [MERaLiON-SpeechEncoder-v1](https://huggingface.co/MERaLiON/MERaLiON
 The model retains near state-of-the-art results on the SUPERB benchmark for English, and showcases strong multilingual capabilities demonstrated through its integration into a [high-performance ASR system](#automatic-speech-recognition-asr).
 
 #### Innovative pre-training techniques
-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 widely-used AdamW optimizer 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 widely-used AdamW optimizer for LLM training. We find its advantages also carry over to speech-based models. 
  
 ## Model Summary
 
@@ -84,7 +84,7 @@ Leveraging on the multilingual capabilities of MERaLiON-SpeechEncoder-2, we furt
 
 ## Direct Use
 
-The following code snippet can be used to directly obtain latent features i.e. encoded speech by forwarding through the model. Inputs into the model are expected to be 80-dimensional Mel-spectrogram features transformed from 16kHz sampled audio. The AutoFeatureExtractor method can carry out the conversion. 
+The following code snippet can be used to directly obtain latent features i.e. encoded speech by forwarding through the model. Inputs into the model are expected to be 80-dimensional Mel-spectrogram features transformed from audio sample at 16kHz. The AutoFeatureExtractor method can carry out this conversion. 
 
 ```python
 import torch
@@ -129,15 +129,19 @@ with torch.no_grad():
                 attention_mask=inputs['attention_mask'],
                 output_hidden_states=True)
 
-    # output is a Wav2Vec2BaseModelOutput or tuple containing:
-    # last_hidden_state: torch.FloatTensor containing hidden states of the last layer of the model
-    # extract_features: torch.FloatTensor containing extracted features from the convolution downsampling layers
-    # hidden_states: tuple(torch.FloatTensor) containing hidden states of each layer of the model
-    # attentions: tuple(torch.FloatTensor) containing attention states of each layer of the model
+# output is a Wav2Vec2BaseModelOutput or tuple containing:
+# last_hidden_state: torch.FloatTensor containing hidden states of the last layer of the model
+# extract_features: torch.FloatTensor containing extracted features from the convolution downsampling layers
+# hidden_states: tuple(torch.FloatTensor) containing hidden states of each layer of the model
+# attentions: tuple(torch.FloatTensor) containing attention states of each layer of the model
 ```
 
 ## Downstream Use
 
+<p align="center">
+  <img src="downstream.svg" width="600"/> 
+</p>
+
 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 using Huggingface Trainer. Alternatively, the Huggingface model can be loaded to any other frameworks such as Pytorch or ESPnet for custom finetuning loops.
 
 ```python
@@ -209,13 +213,15 @@ model = AutoModelForCTC.from_pretrained(
     attention_dropout=0.1,
 )
 model = model.to(device)
+
+# Continue with training loop...
 ```
 
 ### Compute and Infrastructure
 
 MERaLiON-SpeechEncoder-2 was trained on the [**ASPIRE 2A+**](https://help.nscc.sg/aspire2aplus/about/) Supercomputer Cluster, provided by [**National Supercomputing Centre (NSCC)**](https://www.nscc.sg/), Singapore. 
 
-MERaLiON-SpeechEncoder-2 was trained with 64 H100 GPUs across 8 nodes for collectively around 3.5 million steps. Training time took approximately 15 days.
+The model was trained with 64 H100 GPUs across 8 nodes for collectively around 3.5 million steps. The SEA data split was upsampled in the last phase of training. Training time took approximately 15 days. 
 
 
 ## Citation