Update README.md

README.md

@@ -11,7 +11,7 @@ language:
 
 # MERaLiON-SpeechEncoder-v1
 
-The MERaLiON-SpeechEncoder is a speech foundation model that is designed to support a wide range of downstream speech applications, like speech recognition, intent classification and speaker identification, among others. This version is trained on 200,000 hours of predominantly English data including 10,000 hours of Singapore-based speech, to address the speech processing needs in Singapore and beyond.
+The MERaLiON-SpeechEncoder is a speech foundation model designed to support a wide range of downstream speech applications, like speech recognition, intent classification and speaker identification, among others. This version was trained on 200,000 hours of predominantly English data including 10,000 hours of Singapore-based speech, to cater to the speech processing needs in Singapore and beyond. Gradual support for other languages, starting with major South-East Asian ones are planned for subsequent releases.
 
 - **Developed by:** I<sup>2</sup>R, A\*STAR
 - **Funded by:** Singapore NRF
@@ -27,7 +27,6 @@ The model takes in speech as input in the form of mel-spectrograms and returns c
 
 For more details on background, pre-training, tuning experiments and evaluation, please refer to our [technical report]().
 
-
 ## Uses
 
 We have evaluated the MERaLiON-SpeechEncoder extensively on several speech recognition datasets, and fine-tuned the model on ten different tasks encompassing the SUPERB benchmark: `automatic speech recognition` (ASR), `automatic phoneme recognition` (PR), `keyword spotting` (KS), `query by example spoken term detection` (QbE), `intent classification` (IC), `slot filling` (SF), `speaker identification` (SID), `automatic speaker verification` (ASV), `speaker diarization` (SD), and `emotion recognition` (ER).
@@ -44,7 +43,7 @@ import torch
 from datasets import load_dataset
 from transformers import AutoModel, AutoFeatureExtractor
 
-repo_id = "MERaLiON/MERaLiON-SpeechEncoder-V1"
+repo_id = 'MERaLiON/MERaLiON-SpeechEncoder-v1'
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 # load model and feature extractor
@@ -52,19 +51,20 @@ 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
+    trust_remote_code=True
 )
 
 # prepare data
-data = load_dataset("openslr/librispeech_asr", split="validation.clean")
+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"])
+        tensors.append(sample['audio']['array'])
     return tensors
 
 audio_array = batch_collater(data)
@@ -80,24 +80,81 @@ with torch.no_grad():
     output = model(input_values=input_values, input_lengths=input_lengths, output_hidden_states=True)
 ```
 
-### Downstream Use [optional]
+### Downstream Use
 
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+```python
+import torch
+import json
+from datasets import load_dataset
+from transformers import AutoModelForCTC, AutoFeatureExtractor, Wav2Vec2CTCTokenizer
 
-[More Information Needed]
+repo_id = 'MERaLiON/MERaLiON-SpeechEncoder-v1'
+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)
+```
+Consult this [blog](https://huggingface.co/blog/fine-tune-w2v2-bert) for further training recipe with Huggingface Trainer. Alternatively, the Huggingface model can be loaded to any other frameworks such as Pytorch or ESPnet for custom fine-tuning loops.
 
 ## Technical Specifications 
 
 ### Training Data
 
-MERaLiON-SpeechEncoder has been trained on a diverse set of unsupervised speech data, predominantly in English. Our collection is derived and processed from various publicly available datasets. The combined dataset covers a wide range of conditions, encompassing factors such as domain, style, speaker, gender, and accent. Duration-wise it comprises around 170,000 hours of English, including 10,000 hours of Singapore-based English that incorporates code-switching, plus 30,000 additional hours of multilingual speech from 113 languages, totalling 200,000 hours.
-
-[More Information Needed]
+MERaLiON-SpeechEncoder has been trained on a diverse set of unsupervised speech data, primarily in English. Our collection is derived and processed from various publicly available datasets. The combined dataset covers a wide range of conditions, encompassing factors such as domain, style, speaker, gender, and accent. Duration-wise it comprises around 170,000 hours of English, including 10,000 hours of Singapore-based English that incorporates code-switching, plus 30,000 additional hours of multilingual speech from 113 languages, totalling 200,000 hours. Consult our technical report for the full breakdown.
 
 ### Training Procedure and Compute
 
-MERaLiON-SpeechEncoder was trained in two phases, initially on a 60,00 hours subset of data, before continued pre-trainining on the full 200,000 hours dataset using this prior checkpoint as initialisation. The initial model was trained on **ASPIRE 2A** Supercomputer Cluster, provided by the **National Supercomputing Centre (NSCC)** for 325K steps on 12 Nvidia A100 40GB GPUs. The full pre-training run was carried out on the **LUMI** Supercomputer Cluster with 128 AMD MI250x GPUs for a further 382K steps taking about 25 days of active GPU time.
+MERaLiON-SpeechEncoder was trained in two phases, initially on a 60,00 hours subset of data, before continued pre-trainining on the full 200,000 hours dataset using this prior checkpoint as initialisation. The initial model was trained on the **ASPIRE 2A** Supercomputer Cluster provided by the **National Supercomputing Centre (NSCC)** for 325K steps on 12 Nvidia A100 40GB GPUs. The full pre-training run was carried out on the **LUMI** Supercomputer Cluster with 128 AMD MI250x GPUs for a further 382K steps taking about 25 days of active GPU time.
 
 
 ## Citation [optional]