init

.gitignore

@@ -0,0 +1 @@
+__pycache__

app.py

@@ -0,0 +1,44 @@
+import sys
+
+sys.path.append("..")
+
+import gradio
+import torch, torchaudio
+import numpy as np
+from transformers import (
+    Wav2Vec2ForPreTraining,
+    Wav2Vec2CTCTokenizer,
+    Wav2Vec2FeatureExtractor,
+)
+from finetuning.wav2vec2 import SpeechRecognizer
+
+
+def load_model(ckpt_path: str):
+    model_name = "nguyenvulebinh/wav2vec2-base-vietnamese-250h"
+
+    wav2vec2 = Wav2Vec2ForPreTraining.from_pretrained(model_name)
+    tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(model_name)
+    feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name)
+
+    model = SpeechRecognizer.load_from_checkpoint(
+        ckpt_path,
+        wav2vec2=wav2vec2,
+        tokenizer=tokenizer,
+        feature_extractor=feature_extractor,
+    )
+
+    return model
+
+model = load_model("checkpoints/last.ckpt")
+model.eval()
+
+def transcribe(audio):
+    sample_rate, waveform = audio
+    waveform = torch.from_numpy(waveform[:, 0]).float().unsqueeze_(0)
+    waveform = torchaudio.functional.resample(waveform, sample_rate, 16_000)
+
+    transcript = model.predict(waveform)[0]
+
+    return transcript
+
+gradio.Interface(fn=transcribe, inputs=gradio.Audio(source="microphone", type="numpy"), outputs="textbox").launch()

finetuning/preprocess.py

@@ -0,0 +1,27 @@
+import sys
+
+sys.path.append("..")
+
+import os
+import argparse
+from torch.utils.data import random_split
+from src.datamodule import VLSP2020TarDataset, VLSP2020Dataset
+
+
+def prepare_tar_dataset(data_dir: str, dest_dir: str):
+    dts = VLSP2020Dataset(data_dir)
+    train_set, val_set = random_split(dts, [42_000, 14_427])
+
+    VLSP2020TarDataset(os.path.join(dest_dir, "vlsp2020_train_set.tar")).convert(
+        train_set
+    )
+    VLSP2020TarDataset(os.path.join(dest_dir, "vlsp2020_val_set.tar")).convert(val_set)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_dir", type=str, required=True)
+    parser.add_argument("--dest_dir", type=str, required=True)
+    args = parser.parse_args()
+
+    prepare_tar_dataset(args.data_dir, args.dest_dir)

finetuning/run.sh

@@ -0,0 +1,13 @@
+python3 main.py \
+    --batch_size 2 \
+    --num_workers 2 \
+    --classifier_lr 1e-4 \
+    --wav2vec2_lr 1e-5 \
+    --max_epochs 10 \
+    --accelerator cpu \
+    --weight_decay 0.001 \
+    --warmup_steps 0.1 \
+    --constant_steps 0.4 \
+    --scheduler_factor 0.001 \
+    --data_dir data \
+    --ckpt_dir ckpt

finetuning/train.py

@@ -0,0 +1,135 @@
+import sys
+
+sys.path.append("..")
+
+from argparse import ArgumentParser
+import os, string
+from transformers import (
+    Wav2Vec2ForPreTraining,
+    Wav2Vec2CTCTokenizer,
+    Wav2Vec2FeatureExtractor,
+)
+from pytorch_lightning import seed_everything
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
+from pytorch_lightning.loggers import WandbLogger
+
+from src.datamodule import VLSP2020TarDataset
+from src.datamodule.vlsp2020 import get_dataloader
+from finetuning.wav2vec2 import SpeechRecognizer
+
+
+def remove_punctuation(text: str):
+    return text.translate(str.maketrans("", "", string.punctuation)).lower()
+
+
+def prepare_dataloader(data_dir, batch_size, num_workers):
+    train_dataset = VLSP2020TarDataset(
+        os.path.join(data_dir, "vlsp2020_train_set.tar")
+    ).load()
+    val_dataset = VLSP2020TarDataset(
+        os.path.join(data_dir, "vlsp2020_val_set.tar")
+    ).load()
+
+    train_dataloader = get_dataloader(
+        train_dataset,
+        return_transcript=True,
+        target_transform=remove_punctuation,
+        batch_size=batch_size,
+        num_workers=num_workers,
+    )
+
+    val_dataloader = get_dataloader(
+        val_dataset,
+        return_transcript=True,
+        target_transform=remove_punctuation,
+        batch_size=batch_size,
+        num_workers=num_workers,
+    )
+
+    return train_dataloader, val_dataloader
+
+
+def prepare_model(adam_config: dict, tristate_scheduler_config: dict):
+    model_name = "nguyenvulebinh/wav2vec2-base-vietnamese-250h"
+
+    wav2vec2 = Wav2Vec2ForPreTraining.from_pretrained(model_name)
+    tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(model_name)
+    feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name)
+
+    model = SpeechRecognizer(
+        wav2vec2, tokenizer, feature_extractor, adam_config, tristate_scheduler_config
+    )
+
+    return model
+
+
+def main():
+    parser = ArgumentParser()
+
+    parser.add_argument("--batch_size", type=int, default=2)
+    parser.add_argument("--num_workers", type=int, default=0)
+    parser.add_argument("--classifier_lr", type=float, default=1e-4)
+    parser.add_argument("--wav2vec2_lr", type=float, default=1e-5)
+    parser.add_argument("--max_epochs", type=int, default=10)
+    parser.add_argument("--accelerator", type=str, default="gpu")
+    parser.add_argument("--weight_decay", type=float, default=0.0)
+    parser.add_argument("--warmup_steps", type=float, default=0.1)
+    parser.add_argument("--constant_steps", type=float, default=0.4)
+    parser.add_argument("--scheduler_factor", type=float, default=1e-3)
+    parser.add_argument("--data_dir", type=str, default="data")
+    parser.add_argument("--ckpt_dir", type=str, default="ckpt")
+    parser.add_argument("--ckpt_path", type=str, default=None)
+    parser.add_argument("--detect_anomaly", type=bool, default=False)
+    parser.add_argument("--grad_clip", type=float, default=None)
+    parser.add_argument("--wandb_id", type=str, default=None)
+
+    args = parser.parse_args()
+    print(args)
+
+    train_loader, val_loader = prepare_dataloader(
+        args.data_dir, args.batch_size, args.num_workers
+    )
+
+    total_steps = args.max_epochs * 42_000 // args.batch_size
+    warmup_steps = int(total_steps * args.warmup_steps)
+    constant_steps = int(total_steps * args.constant_steps)
+
+    model = prepare_model(
+        {
+            "wav2vec2_lr": args.wav2vec2_lr,
+            "classifier_lr": args.classifier_lr,
+            "weight_decay": args.weight_decay,
+        },
+        {
+            "warmup_steps": warmup_steps,
+            "constant_steps": constant_steps,
+            "total_steps": total_steps,
+            "factor": args.scheduler_factor,
+        },
+    )
+
+    trainer = Trainer(
+        accelerator=args.accelerator,
+        callbacks=[
+            ModelCheckpoint(
+                args.ckpt_dir,
+                monitor="val/wer",
+                mode="min",
+                save_top_k=1,
+                save_last=True,
+            ),
+            LearningRateMonitor(logging_interval="step"),
+        ],
+        logger=WandbLogger(project="Wav2Vec2", id=args.wandb_id),
+        max_epochs=args.max_epochs,
+        detect_anomaly=args.detect_anomaly,
+        gradient_clip_val=args.grad_clip,
+    )
+
+    trainer.fit(model, train_loader, val_loader)
+
+
+if __name__ == "__main__":
+    seed_everything(188)
+    main()

finetuning/wav2vec2.py

@@ -0,0 +1,200 @@
+from typing import Tuple
+import torch
+from pytorch_lightning import LightningModule
+from torchmetrics import MeanMetric
+from transformers import (
+    Wav2Vec2ForPreTraining,
+    Wav2Vec2CTCTokenizer,
+    Wav2Vec2FeatureExtractor,
+)
+
+from src.utils.metrics import character_error_rate, word_error_rate
+from src.utils.scheduler import TriStateScheduler
+
+
+class SpeechRecognizer(LightningModule):
+    def __init__(
+        self,
+        wav2vec2: Wav2Vec2ForPreTraining,
+        tokenizer: Wav2Vec2CTCTokenizer,
+        feature_extractor: Wav2Vec2FeatureExtractor,
+        adam_config: dict,
+        tristate_scheduler_config: dict,
+    ):
+        super().__init__()
+
+        self.hidden_size = wav2vec2.config.proj_codevector_dim
+        self.vocab_size = tokenizer.vocab_size
+
+        self.wav2vec2 = wav2vec2
+        self.wav2vec2.freeze_feature_encoder()
+        self.tokenizer = tokenizer
+        self.feature_extractor = feature_extractor
+
+        self.adam_config = adam_config
+        self.tristate_scheduler_config = tristate_scheduler_config
+
+        self.dropout = torch.nn.Dropout(0.1)
+        self.fc = torch.nn.Sequential(
+            torch.nn.Linear(self.hidden_size, self.hidden_size // 2),
+            torch.nn.ReLU(inplace=True),
+            torch.nn.Linear(self.hidden_size // 2, self.vocab_size),
+        )
+
+        self.criterion = torch.nn.CTCLoss(blank=tokenizer.pad_token_id, zero_infinity=True)
+
+        self.train_loss = MeanMetric()
+
+        self.save_hyperparameters(ignore=["wav2vec2", "tokenizer", "feature_extractor"])
+
+    def forward(self, waveforms: Tuple[torch.Tensor], transcripts: Tuple[str] = None):
+        # convert torch.Tensor to numpy.ndarray
+        waveforms = tuple(waveform.cpu().numpy() for waveform in waveforms)
+
+        input_values, attention_mask = self.feature_extractor(
+            waveforms,
+            sampling_rate=16000,
+            padding=True,
+            return_tensors="pt",
+            return_attention_mask=True,
+        ).values()
+
+        input_values = input_values.to(self.device)
+        attention_mask = attention_mask.to(self.device)
+
+        # hidden_states.shape == (batch_size, sequence_length, hidden_size)
+        hidden_states = self.wav2vec2(
+            input_values,
+            attention_mask=attention_mask,
+        )[0]
+
+        hidden_states = self.dropout(hidden_states)
+
+        # logits.shape == (batch_size, sequence_length, vocab_size)
+        logits = self.fc(hidden_states)
+
+        # get the length of valids sequence
+        input_lengths = self.wav2vec2._get_feat_extract_output_lengths(
+            attention_mask.sum(-1)
+        ).long()
+
+        if transcripts is not None:
+            # tokenize transcripts
+            target_ids, target_lengths = self.tokenizer(
+                transcripts,
+                padding=True,
+                return_length=True,
+                return_attention_mask=False,
+                return_tensors="pt",
+            ).values()
+
+            target_ids = target_ids.to(self.device)
+            assert (
+                target_ids < self.tokenizer.vocab_size
+            ).all(), "target_ids is out of range"
+
+            target_lengths = target_lengths.to(self.device)
+            assert (
+                target_lengths <= logits.size(1)
+            ).all(), "target_lengths is out of range"
+
+            # (batch_size, sequence_length, vocab_size) -> (sequence_length, batch_size, vocab_size)
+            log_probs = torch.nn.functional.log_softmax(logits, dim=-1).transpose(0, 1)
+
+            # compute loss
+            loss = self.criterion(log_probs, target_ids, input_lengths, target_lengths)
+
+            return loss, logits, input_lengths
+        else:
+            return logits, input_lengths
+
+    @staticmethod
+    def _get_predicted_ids(logits: torch.Tensor, lengths: torch.Tensor):
+        # logits.shape == (batch_size, sequence_length, vocab_size)
+        # lengths.shape == (batch_size, )
+
+        # get the max value of logits
+        predicted_ids = torch.argmax(logits, dim=-1)
+
+        # remove the padding
+        predicted_ids = [
+            predicted_id[:length]
+            for predicted_id, length in zip(predicted_ids, lengths)
+        ]
+
+        return predicted_ids
+
+    def training_step(self, batch, batch_idx):
+        transcripts, waveforms = batch
+
+        loss = self(waveforms, transcripts)[0]
+
+        self.train_loss(loss)
+
+        if self.global_step % 500 == 0:
+            self.log("train/loss", self.train_loss, on_step=True, on_epoch=True)
+
+        return loss
+
+    def on_train_epoch_end(self) -> None:
+        self.train_loss.reset()
+
+    def validation_step(self, batch, batch_idx):
+        transcripts, waveforms = batch
+
+        logits, seq_lengths = self(waveforms)
+
+        predicted_ids = self._get_predicted_ids(logits, seq_lengths)
+        predicted_texts = self.tokenizer.batch_decode(
+            predicted_ids, skip_special_tokens=True
+        )
+
+        wer = word_error_rate(predicted_texts, transcripts)
+        cer = character_error_rate(predicted_texts, transcripts)
+
+        return wer, cer
+
+    def validation_epoch_end(self, outputs):
+        wer, cer = zip(*outputs)
+
+        wer = sum(wer) / len(wer)
+        cer = sum(cer) / len(cer)
+
+        self.log("val/wer", wer, on_epoch=True)
+        self.log("val/cer", cer, on_epoch=True)
+
+    @torch.no_grad()
+    def predict(self, waveforms: Tuple[torch.Tensor]):
+        logits, seq_lengths = self(waveforms)
+
+        predicted_ids = self._get_predicted_ids(logits, seq_lengths)
+        predicted_texts = self.tokenizer.batch_decode(
+            predicted_ids, skip_special_tokens=True
+        )
+
+        return predicted_texts
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(
+            params=[
+                {
+                    "params": self.wav2vec2.parameters(),
+                    "lr": self.adam_config["wav2vec2_lr"],
+                },
+                {
+                    "params": self.fc.parameters(),
+                    "lr": self.adam_config["classifier_lr"],
+                },
+            ],
+            weight_decay=self.adam_config["weight_decay"],
+        )
+
+        scheduler = TriStateScheduler(optimizer, **self.tristate_scheduler_config)
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": scheduler,
+                "interval": "step",
+                "frequency": 1,
+            },
+        }

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,9 @@
+torch
+torchaudio
+pytorch-lightning
+einops
+easydict
+webdataset
+transformers
+gradio
+altair

src/config/model.py

@@ -0,0 +1,57 @@
+from easydict import EasyDict as dict
+
+D_MODEL = 768
+HIDDEN_SIZE = 512
+
+
+
+context_encoder = dict(
+    feature_projection=dict(
+        in_features=HIDDEN_SIZE,
+        out_features=D_MODEL,
+        dropout=0.1,
+    ),
+    encoder=dict(
+        d_model=D_MODEL,
+        num_layers=12,
+        layer_drop=0.05,
+        pos_embedding=dict(
+            d_model=D_MODEL,
+            kernel_size=3,
+            groups=2,
+            dropout=0.1,
+        ),
+        layer=dict(
+            d_model=D_MODEL,
+            num_heads=8,
+            layer_norm_first=False,
+            feed_forward_dim=2048,
+            dropout=0.1,
+        ),
+    )
+)
+
+feature_extractor = dict(
+    num_channels=7 * (HIDDEN_SIZE,),
+    kernel_sizes=(10,) + 4 * (3,) + 2 * (2,),
+    strides=(5,) + 6 * (2,),
+)
+
+quantizer = dict(
+    in_features=HIDDEN_SIZE,
+    num_codebooks=2,
+    num_codewords=320,
+    d_model=D_MODEL,
+)
+
+wav2vec2_pretraining = dict(
+    context_encoder=context_encoder,
+    feature_extractor=feature_extractor,
+    quantizer=quantizer,
+    mask_prob=0.65,
+    mask_length=10,
+    min_masks=2,
+    num_negatives=100,
+    contrastive_logits_temperature=0.1,
+    diversity_loss_weight=0.2,
+)

src/datamodule/__init__.py

@@ -0,0 +1,4 @@
+from .vlsp2020 import (
+    VLSP2020TarDataset,
+    VLSP2020Dataset,
+)

src/datamodule/vlsp2020.py

@@ -0,0 +1,131 @@
+from typing import Callable, Optional, Union
+from tqdm import tqdm
+import os
+import torch
+import torchaudio
+import torchaudio.functional as F
+from torch.utils.data import Dataset, DataLoader, IterableDataset, random_split
+from pytorch_lightning import LightningDataModule
+import webdataset
+
+
+class VLSP2020Dataset(Dataset):
+    def __init__(self, root: str, sample_rate: int = 16000):
+        super().__init__()
+
+        self.sample_rate = sample_rate
+        self.memory = self._prepare_data(root)
+        self._memory = tuple(
+            (v["transcript"], v["audio"]) for v in self.memory.values()
+        )
+
+    @staticmethod
+    def _prepare_data(root: str):
+        memory = {}
+
+        for f in os.scandir(root):
+            file_name, file_ext = os.path.splitext(f.name)
+
+            if file_ext == ".txt":
+                if file_name not in memory:
+                    memory[file_name] = {"transcript": f.path}
+                elif "transcript" not in memory[file_name]:
+                    memory[file_name]["transcript"] = f.path
+                else:
+                    raise ValueError(f"Duplicate transcript for {f.path}")
+            else:
+                if file_name not in memory:
+                    memory[file_name] = {"audio": f.path}
+                elif "audio" not in memory[file_name]:
+                    memory[file_name]["audio"] = f.path
+                else:
+                    raise ValueError(f"Duplicate audio for {f.path}")
+
+        for key, value in memory.items():
+            if "audio" not in value:
+                raise ValueError(f"Missing audio for {key}")
+            elif "transcript" not in value:
+                raise ValueError(f"Missing transcript for {key}")
+
+        return memory
+
+    def __len__(self):
+        return len(self.memory)
+
+    def __getitem__(self, index: int):
+        transcript, audio = self._memory[index]
+
+        with open(transcript, "r") as f:
+            transcript = f.read()
+
+        audio, sample_rate = torchaudio.load(audio)
+        audio = F.resample(audio, sample_rate, self.sample_rate)
+
+        return transcript, audio
+
+
+class VLSP2020TarDataset:
+    def __init__(self, outpath: str):
+        self.outpath = outpath
+
+    def convert(self, dataset: VLSP2020Dataset):
+        writer = webdataset.TarWriter(self.outpath)
+
+        for idx, (transcript, audio) in enumerate(tqdm(dataset, colour="green")):
+            writer.write(
+                {
+                    "__key__": f"{idx:08d}",
+                    "txt": transcript,
+                    "pth": audio,
+                }
+            )
+
+        writer.close()
+
+    def load(self) -> webdataset.WebDataset:
+        self.data = (
+            webdataset.WebDataset(self.outpath)
+            .decode(
+                webdataset.handle_extension("txt", lambda x: x.decode("utf-8")),
+                webdataset.torch_audio,
+            )
+            .to_tuple("txt", "pth")
+        )
+
+        return self.data
+
+
+def get_dataloader(
+    dataset: Union[VLSP2020Dataset, webdataset.WebDataset],
+    return_transcript: bool = False,
+    target_transform: Optional[Callable] = None,
+    batch_size: int = 32,
+    num_workers: int = 2,
+):
+    def collate_fn(batch):
+        def get_audio(item):
+            audio = item[1]
+
+            assert (
+                isinstance(audio, torch.Tensor)
+                and audio.ndim == 2
+                and audio.size(0) == 1
+            )
+
+            return audio.squeeze(0)
+
+        audio = tuple(get_audio(item) for item in batch)
+
+        if return_transcript:
+            if target_transform is not None:
+                transcript = tuple(target_transform(item[0]) for item in batch)
+            else:
+                transcript = tuple(item[0] for item in batch)
+
+            return transcript, audio
+        else:
+            return audio
+
+    return DataLoader(
+        dataset, batch_size=batch_size, num_workers=num_workers, collate_fn=collate_fn
+    )

src/model/__init__.py

@@ -0,0 +1 @@
+from .wav2vec2 import Wav2Vec2PretrainingModule

src/model/modules/__init__.py

@@ -0,0 +1,4 @@
+from .context_encoder import ContextEncoder
+from .feature_extractor import FeatureExtractor
+from .quantization import QuantizationModule
+from .processor import Wav2Vec2Processor

src/model/modules/context_encoder.py

@@ -0,0 +1,149 @@
+from typing import Optional
+import torch
+from torch import nn
+import torch.nn.functional as F
+from .transformers import EncoderLayer
+
+
+class FeatureProjection(nn.Module):
+    def __init__(self, in_features: int, out_features: int, dropout: float = 0.1):
+        """
+        Projects the extracted features to the encoder dimension.
+
+        Args:
+            x (Tensor): The input features. Shape: (batch, num_frames, in_features)
+
+        Returns:
+            hiddens (Tensor): The latent features. Shape: (batch, num_frames, out_features)
+        """
+        super().__init__()
+
+        self.projection = nn.Linear(in_features, out_features)
+        self.layernorm = nn.LayerNorm(in_features)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor):
+
+        hiddens = self.layernorm(x)
+        hiddens = self.projection(x)
+        hiddens = self.dropout(hiddens)
+        return hiddens
+
+
+class RelativePositionalEmbedding(nn.Module):
+    def __init__(
+        self, d_model: int, kernel_size: int, groups: int, dropout: float = 0.1
+    ):
+        """
+        Args:
+            x (Tensor): The extracted features. Shape: (batch, num_frames, d_model)
+
+        Returns:
+            out (Tensor): The output which encoded the relative positional information. Shape: (batch, num_frames, d_model)
+        """
+        super().__init__()
+
+        self.conv = nn.Conv1d(
+            in_channels=d_model,
+            out_channels=d_model,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+            groups=groups,
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.num_remove = 1 if kernel_size % 2 == 0 else 0
+
+    def forward(self, x: torch.Tensor):
+        # (batch, channels=d_model, num_frames)
+        out = x.transpose(1, 2)
+
+        out = self.conv(out)
+
+        if self.num_remove > 0:
+            out = out[..., : -self.num_remove]
+
+        out = F.gelu(out)
+
+        # (batch, num_frames, channels=d_model)
+        out = out.transpose_(1, 2)
+        out = out + x
+        out = self.dropout(out)
+
+        return out
+
+
+class TranformerEncoder(nn.Module):
+    def __init__(self, config):
+        """
+        Args:
+            x (Tensor): The extracted features. Shape: (batch, num_frames, d_model)
+            mask (Tensor): The mask for the valid frames. Shape: (batch, num_frames)
+
+        Returns:
+            out (Tensor): The output of the transformer encoder. Shape: (batch, num_frames, d_model)
+        """
+        super().__init__()
+
+        self.pos_embedding = RelativePositionalEmbedding(**config.pos_embedding)
+        self.layernorm = nn.LayerNorm(config.d_model)
+        self.layer_drop = config.layer_drop
+
+        self.layers = nn.ModuleList(
+            EncoderLayer(**config.layer) for _ in range(config.num_layers)
+        )
+
+    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None):
+        out = self.pos_embedding(x)
+
+        for layer in self.layers:
+            skip_layer = self.training and torch.rand(1).item() < self.layer_drop
+
+            if skip_layer:
+                continue
+            else:
+                out, _ = layer(out, attention_mask=mask)
+
+        out = self.layernorm(out)
+
+        return out
+
+
+class ContextEncoder(nn.Module):
+    def __init__(self, config):
+        """
+        Args:
+            x (Tensor): The extracted features. Shape: (batch, num_frames, in_features)
+            attention_mask (BoolTensor): The mask for the valid frames. `True` is invalid. Shape: (batch, num_frames)
+        """
+        super().__init__()
+
+        self.feature_projection = FeatureProjection(**config.feature_projection)
+        self.encoder = TranformerEncoder(config.encoder)
+        self.masked_spec_embed = nn.Parameter(
+            torch.FloatTensor(config.feature_projection.out_features).uniform_()
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: torch.Tensor = None,
+        mask_time_indices: torch.Tensor = None,
+    ):
+        x = self.feature_projection(x)
+
+        if mask_time_indices is not None:
+            x[mask_time_indices] = self.masked_spec_embed.to(x.dtype)
+
+        if attention_mask is not None:
+            x[attention_mask] = 0.0  # turn invalid frames to zero
+
+            attention_mask = attention_mask[:, None, None, :]
+            # (batch, 1, num_frames, num_frames)
+            # mask = mask[:, None, None, :].repeat(1, 1, mask.size(1), 1) # TODO: check this
+            attention_mask = (
+                torch.maximum(attention_mask, attention_mask.transpose(2, 3)) * -1e6
+            )
+
+        x = self.encoder(x, mask=attention_mask)
+
+        return x

src/model/modules/feature_extractor.py

@@ -0,0 +1,103 @@
+from typing import Tuple
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class _Conv1DLayer(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int,
+    ):
+        """
+        Args:
+            x (Tensor): The ouput. Shape: (batch, in_channels, in_frames)
+            length (Tensor): The valid length of each sample. Shape: (batch)
+
+        Returns:
+            x (Tensor): The output. Shape: (batch, out_channels, out_frames)
+            length (Tensor): The valid length of each sample. Shape: (batch)
+        """
+        super().__init__()
+
+        self.kernel_size = kernel_size
+        self.stride = stride
+
+        self.conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            stride=stride,
+            kernel_size=kernel_size,
+            bias=False,
+        )
+
+        self.layernorm = nn.LayerNorm(out_channels)
+
+    def forward(self, x: torch.Tensor, length: torch.Tensor):
+        x = self.conv(x)
+        x = x.transpose_(1, 2)
+        x = self.layernorm(x)
+        x = x.transpose_(1, 2)
+        x = F.gelu(x)
+
+        length = (length - self.kernel_size) // self.stride + 1
+        length = length.clamp_min_(min=0)  # prevent negative lengths
+        return x, length
+
+
+class FeatureExtractor(nn.Module):
+    def __init__(self, config):
+        """
+        Extracts features from the waveform.
+
+        Args:
+            waveforms (Tensor): The waveform to extract features from. Shape: (batch, wavelength)
+            wavelength (Tensor): The valid length of each waveform. Shape: (batch)
+
+        Returns:
+            features (Tensor): The extracted features. Shape: (batch, num_frames, num_channels)
+            num_frames (Tensor): The valid length of each feature. Shape: (batch)
+        """
+        super().__init__()
+
+        num_channels = config.num_channels
+        kernel_sizes = config.kernel_sizes
+        strides = config.strides
+
+        assert (
+            len(num_channels) == len(kernel_sizes) == len(strides)
+        ), "The number of layers must be the same for all parameters"
+
+        self.conv_layers = nn.ModuleList(
+            (
+                _Conv1DLayer(
+                    in_channels=1,
+                    out_channels=num_channels[0],
+                    kernel_size=kernel_sizes[0],
+                    stride=strides[0],
+                ),
+            )
+        )
+
+        for i in range(1, len(num_channels)):
+            self.conv_layers.append(
+                _Conv1DLayer(
+                    in_channels=num_channels[i - 1],
+                    out_channels=num_channels[i],
+                    kernel_size=kernel_sizes[i],
+                    stride=strides[i],
+                )
+            )
+
+    def forward(self, waveforms: torch.Tensor, wavelength: torch.Tensor):
+        features = waveforms.unsqueeze(1)
+
+        for conv_layer in self.conv_layers:
+            features, wavelength = conv_layer(features, wavelength)
+
+        # (batch, num_channels, num_frames) -> (batch, num_frames, num_channels)
+        features = features.transpose(1, 2)
+        return features, wavelength

src/model/modules/processor.py

@@ -0,0 +1,42 @@
+from typing import Tuple
+import torch
+from torch import nn
+
+
+class Wav2Vec2Processor(nn.Module):
+    def __init__(self):
+        """
+        Convert tuple of waveforms whose length is different to a batch.
+
+        Args:
+            waveforms (Tuple[torch.Tensor]): The waveforms. Shape: (batch_size, wave_length).
+
+        Returns:
+            waveforms (torch.Tensor): The batched waveforms. Shape: (batch_size, max_wave_length).
+            wave_lengths (torch.Tensor): The wave length of each waveform. Shape: (batch_size,).
+        """
+        super().__init__()
+
+    def forward(self, waveforms: Tuple[torch.Tensor, ...]):
+        device = waveforms[0].device
+        wave_lengths = torch.tensor(
+            tuple(waveform.size(0) for waveform in waveforms), device=device
+        )
+
+        max_length = wave_lengths.max().item()
+
+        padded = []
+
+        for waveform in waveforms:
+            padded.append(
+                nn.functional.pad(
+                    waveform,
+                    (0, max_length - waveform.size(0)),
+                    mode="constant",
+                    value=0.0,
+                )
+            )
+
+        batched_waveforms = torch.stack(padded, dim=0)
+
+        return batched_waveforms, wave_lengths

src/model/modules/quantization.py

@@ -0,0 +1,103 @@
+from typing import Optional
+import torch
+from torch import nn
+import torch.nn.functional as F
+import einops
+
+
+class QuantizationModule(nn.Module):
+    def __init__(
+        self, config
+    ):
+        """
+        Args:
+            x (Tensor): The extracted features from waveforms. Shape: (batch, num_frames, in_features)
+            mask (BoolTensor): The mask for the valid frames. `True` is invalid. Shape: (batch, num_frames)
+
+        Returns:
+            out (Tensor): The quantized features. Shape: (batch, num_frames, d_model)
+            perplexity (Tensor): The perplexity of the quantized features. Shape: (1)
+        """
+        super().__init__()
+
+        assert (
+            config.d_model % config.num_codebooks == 0
+        ), "d_model must be divisible by num_codebooks"
+
+        self.num_codebooks = config.num_codebooks
+        self.num_codewords = config.num_codewords
+        self.d_model = config.d_model
+        self.codeword_dim = config.d_model // config.num_codebooks
+
+        self.codebooks = self._init_codebooks()
+
+        self.projection = nn.Linear(
+            config.in_features, self.num_codebooks * self.num_codewords
+        )
+
+        self.tau = 1  # temperature factor
+
+    def _init_codebooks(self):
+        codebooks = torch.randn(
+            1, 1, self.num_codebooks, self.num_codewords, self.codeword_dim
+        )
+        nn.init.xavier_uniform_(codebooks)
+
+        return nn.Parameter(codebooks)
+
+    @property
+    def total_codewords(self):
+        return self.num_codebooks * self.num_codewords
+
+    @staticmethod
+    def _compute_perplexity(probs: torch.Tensor, mask: Optional[torch.Tensor] = None):
+        """
+        Computes the perplexity of the quantized features. (Diversity loss)
+
+        Args:
+            probs (Tensor): The probability distribution of words in codebooks. Shape: (batch, num_frames, num_codebooks, num_codewords)
+            mask (BoolTensor): The mask for the valid frames. `True` is invalid. Shape: (batch, num_frames)
+        """
+        if mask is not None:
+            probs = (
+                probs * ~mask[..., None, None]
+            )  # Turn invalid frames' probability to 0
+            marginal_probs = (
+                einops.reduce(probs, "b nf nb nw -> nb nw", "sum") / mask.sum()
+            )
+        else:
+            marginal_probs = einops.reduce(probs, "b nf nb nw -> nb nw", "mean")
+
+        perplexity = torch.exp(
+            -torch.sum(marginal_probs * torch.log(marginal_probs + 1e-7), dim=-1)
+        ).sum()
+        return perplexity
+
+    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None):
+        batch_size, num_frames, _ = x.shape
+
+        logits = self.projection(x)
+        logits = logits.view(
+            batch_size, num_frames, self.num_codebooks, self.num_codewords
+        )
+
+        if self.training:
+            word_probs = F.gumbel_softmax(logits, tau=self.tau, hard=True, dim=-1)
+            word_soft_probs = F.softmax(logits, dim=-1)
+
+            perplexity = self._compute_perplexity(word_soft_probs, mask=mask)
+        else:
+            word_ids = torch.argmax(logits, dim=-1, keepdim=True)
+            word_probs = torch.zeros_like(logits).scatter_(-1, word_ids, 1.0)  # One-hot
+
+            perplexity = self._compute_perplexity(word_probs, mask=mask)
+
+        # (batch, num_frames, num_codebooks, num_codewords, 1) x (1, 1, num_codebooks, num_codewords, codeword_dim)
+        # -> (batch, num_frames, num_codebooks x codeword_dim)
+        quantized = einops.reduce(
+            word_probs.unsqueeze_(-1) * self.codebooks,
+            "b nf nb nw d -> b nf (nb d)",
+            reduction="sum",
+        )
+
+        return quantized, perplexity

src/model/modules/transformers.py

@@ -0,0 +1,200 @@
+"""
+This file contains the implementation of the Transformer Encoder layer.
+Source: https://github.com/pytorch/audio/blob/main/torchaudio/models/wav2vec2/components.py
+"""
+from typing import Optional, Tuple
+import torch
+from torch import nn, Tensor
+from torch.nn import Module
+
+
+class SelfAttention(Module):
+    """Multihead Self Attention module
+    Args:
+        embed_dim (int): Total dimension of the model.
+        num_heads (int): The number of heads.
+        dropout (float, optional):
+            Dropout probability on attn_output_weights. Default: ``0.0``
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        head_dim = embed_dim // num_heads
+        if head_dim * num_heads != embed_dim:
+            raise ValueError(
+                f"`embed_dim ({embed_dim})` is not divisible by `num_heads ({num_heads})`"
+            )
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = torch.nn.Dropout(dropout)
+        self.head_dim = head_dim
+
+        self.scaling = self.head_dim**-0.5
+
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+
+    def forward(
+        self,
+        x: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        position_bias: Optional[Tensor] = None,
+        key_padding_mask: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        """
+        Args:
+            x (Tensor): shape: ``[batch_size, sequence_length, embed_dim]``.
+            attention_mask (Tensor or ``None``, optional):
+                shape: ``[batch_size, 1, sequence_length, sequence_length]``
+            position_bias: Not used. Only for the compatibility with :py:class:`WavLMSelfAttention`.
+            key_padding_mask (Tensor or ``None``): Not used. Only for the compatibility with
+                :py:class:`WavLMSelfAttention`.
+        Returns:
+            (Tensor, ``None``): The resulting attention output and ``None`` (necessary for compatibility
+                with :py:class:`WavLMSelAttention`).
+                Attention output shape: ``[batch, sequence_length, embed_dim]``.
+        """
+        if x.ndim != 3 or x.shape[2] != self.embed_dim:
+            raise ValueError(
+                f"The expected input shape is (batch, sequence, embed_dim=={self.embed_dim}). "
+                f"Found {x.shape}."
+            )
+        batch_size, length, embed_dim = x.size()
+        if attention_mask is not None:
+            shape_ = (batch_size, 1, length, length)
+            if attention_mask.size() != shape_:
+                raise ValueError(
+                    f"The expected attention mask shape is {shape_}. "
+                    f"Found {attention_mask.size()}."
+                )
+
+        shape = (batch_size, length, self.num_heads, self.head_dim)
+        q = self.q_proj(x).view(*shape).transpose(2, 1)  # B, nH, L, Hd
+        k = self.k_proj(x).view(*shape).permute(0, 2, 3, 1)  # B, nH, Hd, L
+        v = self.v_proj(x).view(*shape).transpose(2, 1)  # B, nH, L, Hd
+
+        # scale down q to avoid value overflow.
+        weights = (self.scaling * q) @ k  # B, nH, L, L
+        if attention_mask is not None:
+            weights += attention_mask
+        # subtracting a constant value from the tensor won't change the output of softmax.
+        # apply the subtraction to avoid value overflow in torch.nn.functional.softmax.
+        # for more details, please see Equation 7 in https://arxiv.org/abs/2112.08778
+        weights = weights - weights.max(dim=-1, keepdim=True)[0]
+
+        weights = torch.nn.functional.softmax(weights, dim=-1)
+        weights = self.dropout(weights)
+
+        output = weights @ v  # B, nH, L, Hd
+        output = output.transpose(2, 1).reshape(batch_size, length, embed_dim)
+
+        output = self.out_proj(output)
+        return output, None  # Necessary for compatibility with WavLMSelAttention
+
+
+class FeedForward(Module):
+    """Layer that follows attention layer in encoder layer."""
+
+    def __init__(
+        self,
+        io_features: int,
+        intermediate_features: int,
+        intermediate_dropout: float,
+        output_dropout: float,
+    ):
+        super().__init__()
+        self.intermediate_dense = nn.Linear(io_features, intermediate_features)
+        self.intermediate_dropout = nn.Dropout(intermediate_dropout)
+        self.output_dense = nn.Linear(intermediate_features, io_features)
+        self.output_dropout = nn.Dropout(output_dropout)
+
+    def forward(self, x):
+        """
+        Args:
+            x (Tensor): shape: `(batch, sequence_length, io_features)`
+        Returns:
+            x (Tensor): shape: `(batch, sequence_length, io_features)`
+        """
+        x = self.intermediate_dense(x)
+        x = torch.nn.functional.gelu(x)
+        x = self.intermediate_dropout(x)
+
+        x = self.output_dense(x)
+        x = self.output_dropout(x)
+        return x
+
+
+class EncoderLayer(Module):
+    """A layer unit in encoder. Combines multihead self attention and feed forward."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        layer_norm_first: bool,
+        feed_forward_dim: int,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.attention = SelfAttention(
+            embed_dim=d_model,
+            num_heads=num_heads,
+            dropout=dropout,
+        )
+
+        self.dropout = nn.Dropout(dropout)
+        self.layer_norm = nn.LayerNorm(d_model)
+        self.layer_norm_first = layer_norm_first
+        self.feed_forward = FeedForward(d_model, feed_forward_dim, dropout, dropout)
+        self.final_layer_norm = nn.LayerNorm(d_model)
+
+    def forward(
+        self,
+        x: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        position_bias: Optional[Tensor] = None,
+        key_padding_mask: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        """
+        Args:
+            x (Tensor): Input of shape ``(batch, sequence_length, embed_dim)``.
+            attention_mask (Tensor or ``None``, optional): attention mask
+                of shape ``(batch, 1, sequence_length, sequence_length)``. (Default: ``None``)
+            position_bias (Tensor or ``None``, optional): position bias of shape
+                ``(batch_size * num_heads, src_len, src_len)``.
+                Only necessary for WavLM model, ``None`` otherwise. (Default: ``None``)
+            key_padding_mask (Tensor or ``None``, optional): key padding mask of shape ``(batch_size, src_len)``.
+                Only used for WavLM model, ignored otherwise. (Default: ``None``)
+        Returns:
+            (x, position_bias): Shapes are the same as in the input. Position bias is only relevant for WaLM model,
+                ``None`` otherwise.
+        """
+        residual = x
+
+        if self.layer_norm_first:
+            x = self.layer_norm(x)
+
+        x, position_bias = self.attention(
+            x,
+            attention_mask=attention_mask,
+            position_bias=position_bias,
+            key_padding_mask=key_padding_mask,
+        )
+
+        x = self.dropout(x)
+        x = residual + x
+
+        if self.layer_norm_first:
+            x = x + self.feed_forward(self.final_layer_norm(x))
+        else:
+            x = self.layer_norm(x)
+            x = self.final_layer_norm(x + self.feed_forward(x))
+        return x, position_bias

src/model/wav2vec2.py

@@ -0,0 +1,293 @@
+"""
+A wrapper of Wav2Vec2 for training phase.
+"""
+from typing import Tuple, Optional
+import torch
+from pytorch_lightning import LightningModule
+import einops
+from torchmetrics import MeanMetric
+
+from .modules import (
+    ContextEncoder,
+    FeatureExtractor,
+    QuantizationModule,
+    Wav2Vec2Processor,
+)
+from src.utils import init_module_weights
+
+
+class Wav2Vec2PretrainingModule(LightningModule):
+    def __init__(self, config):
+        super().__init__()
+
+        self.save_hyperparameters(config)
+
+        self.processor = Wav2Vec2Processor()
+        self.context_encoder = ContextEncoder(config.context_encoder)
+        self.feature_extractor = FeatureExtractor(config.feature_extractor)
+        self.quantizer = QuantizationModule(config.quantizer)
+
+        self.train_loss = MeanMetric()
+
+    def forward(self, waveforms: Tuple[torch.Tensor, ...]):
+        """
+        Args:
+            waveforms (Tuple[torch.Tensor]): The waveforms. Shape: (batch_size, wave_length).
+
+        Returns:
+            loss: The loss of the model. Contrastive loss + Diversity loss.
+        """
+        waveforms, wave_lengths = self.processor(waveforms)
+
+        # features.shape == (batch_size, num_frames, hidden_size)
+        features, num_frames = self.feature_extractor(waveforms, wave_lengths)
+
+        attention_mask = self._compute_attention_mask(num_frames)
+        mask_time_indices = self._compute_mask_span(
+            shape=features.shape[:-1],
+            mask_prob=self.hparams.mask_prob,
+            mask_length=self.hparams.mask_length,
+            attention_mask=attention_mask,
+            device=features.device,
+            min_masks=self.hparams.min_masks,
+        )
+
+        context_features = self.context_encoder(
+            features, attention_mask=attention_mask, mask_time_indices=mask_time_indices
+        )
+
+        quantized_features, perplexity = self.quantizer(features, attention_mask)
+
+        negative_quantized_features = self._sample_negatives(
+            quantized_features,
+            num_negatives=self.hparams.num_negatives,
+            attention_mask=attention_mask,
+        )
+
+        # (batch_size, num_frames, num_negatives + 1)
+        contrastive_logits = self._compute_contrastive_logits(
+            context_features,
+            quantized_features,
+            negative_quantized_features,
+            self.hparams.contrastive_logits_temperature,
+        ).flatten(0, -2)
+
+        # compute contrastive loss
+        # positive indices are always the first one
+        targets = (1 - mask_time_indices.long().flatten()) * -100
+
+        contrastive_loss = torch.nn.functional.cross_entropy(
+            contrastive_logits, targets, reduction="sum"
+        )
+
+        # compute diversity loss
+        diversity_loss = 1 - perplexity / self.quantizer.total_codewords
+
+        loss = contrastive_loss + diversity_loss * self.hparams.diversity_loss_weight
+
+        return loss
+
+    @staticmethod
+    def _sample_negatives(
+        features: torch.Tensor,
+        num_negatives: int,
+        attention_mask: Optional[torch.Tensor] = None,
+    ):
+        """
+        Sampling negative features from quantized features to compute the contrastive loss.
+
+        Args:
+            features (torch.Tensor): The quantized features. Shape: (batch_size, num_frames, d_model).
+            num_negatives (int): The number of negative samples.
+            attention_mask (Optional[torch.Tensor]): The mask for valid frames. `True` is invalid. Shape: (batch_size, num_frames).
+
+        Returns:
+            sampled_negatives (torch.Tensor): The sampled negative features. Shape: (batch_size, num_frames, num_negatives, d_model).
+        """
+
+        batch_size, num_frames, d_model = features.shape
+
+        features = features.view(-1, d_model)  # (batch_size * num_frames, d_model)
+
+        with torch.no_grad():
+            sampled_ids = []
+
+            for batch_idx in range(batch_size):
+                num_valid_frames = (
+                    features.size(1)
+                    if attention_mask is None
+                    else (1 - attention_mask[batch_idx].long()).sum()
+                ).item()
+
+                sampled_ids.append(
+                    torch.randint(
+                        0,
+                        num_valid_frames - 1,
+                        (num_frames * num_negatives,),
+                        device=features.device,
+                    )
+                )
+
+            sampled_ids = torch.stack(
+                sampled_ids, dim=0
+            )  # (batch_size, num_frames * num_negatives)
+
+            feature_ids = einops.repeat(
+                torch.arange(num_frames, device=features.device),
+                "f -> (f n)",
+                n=num_negatives,
+            )
+
+            # avoid sampling the same positive vector, but keep the distribution uniform
+            sampled_ids[sampled_ids >= feature_ids] += 1
+
+        # correct for batch size
+        # E.g [[0, 1, 2], [0, 1, 2]] -> [0, 1, 2, 3, 4, 5]
+        sampled_ids += torch.arange(
+            0, batch_size * num_frames, num_frames, device=features.device
+        ).unsqueeze_(-1)
+
+        sampled_negatives = features[sampled_ids.view(-1)]
+        sampled_negatives = einops.rearrange(
+            sampled_negatives,
+            "(b f n) d -> b f n d",
+            b=batch_size,
+            f=num_frames,
+            n=num_negatives,
+        )
+
+        return sampled_negatives
+
+    @staticmethod
+    def _compute_contrastive_logits(
+        predicted_features: torch.Tensor,
+        target_features: torch.Tensor,
+        negative_features: torch.Tensor,
+        temperature: int = 1,
+    ):
+        """
+        Compute the logits for contrastive loss.
+
+        Args:
+            predicted_features (torch.Tensor): The predicted features. Shape: (batch_size, num_frames, d_model).
+            target_features (torch.Tensor): The target features. Shape: (batch_size, num_frames, d_model).
+            negative_features (torch.Tensor): The negative features. Shape: (batch_size, num_frames, num_negatives, d_model).
+            temperature (int): The temperature for contrastive loss.
+
+        Returns:
+            logits (torch.Tensor): The logits for contrastive loss. Shape: (batch_size, num_frames, num_negatives + 1).
+        """
+
+        # (batch_size, num_frames, num_negatives + 1, d_model)
+        target_features = torch.cat(
+            (target_features.unsqueeze_(2), negative_features), dim=2
+        )
+
+        # (batch_size, num_frames, 1, d_model)
+        predicted_features = predicted_features.unsqueeze_(2)
+
+        # (batch_size, num_frames, num_negatives + 1)
+        logits = torch.cosine_similarity(predicted_features, target_features, dim=-1)
+        logits /= temperature
+
+        return logits
+
+    @staticmethod
+    def _compute_mask_span(
+        shape: Tuple[int, int],
+        mask_prob: float = 0.065,
+        mask_length: int = 10,
+        attention_mask: Optional[torch.Tensor] = None,
+        device: torch.device = torch.device("cpu"),
+        min_masks: int = 0,
+    ):
+        """
+        Compute the mask span for contrastive task.
+
+        Args:
+            shape (Tuple[int, int]): The shape of the mask span. Shape: (batch_size, num_frames).
+            mask_prob (float): The probability of choosing a frame to be the start of masking position.
+            mask_length (int): The length of the mask span.
+            attention_mask (Optional[torch.Tensor]): The mask for valid frames. `True` is invalid. Shape: (batch_size, num_frames).
+            device (torch.device): The device of the mask span.
+            min_masks (int): The minimum number of masks.
+
+        Returns:
+            mask_span (torch.Tensor): The mask span. Shape: (batch_size, num_frames).
+        """
+
+        batch_size, num_frames = shape
+
+        # NOTE: num_frames / mask_length: the number of spans in one waveform
+        num_masked_spans = int(
+            mask_prob * num_frames / mask_length + torch.rand(1).item()
+        )
+        num_masked_spans = max(num_masked_spans, min_masks)
+
+        # make sure num masked indices <= num frames
+        if num_masked_spans * mask_length > num_frames:
+            num_masked_spans = num_frames // mask_length
+
+        # uniform distribution to sample from
+        # NOTE: num_frames - (mask_length - 1): the number of start positions of the span
+        uniform_dist = torch.ones(
+            (batch_size, num_frames - (mask_length - 1)), device=device
+        )
+
+        # (batch_size, num_masked_spans)
+        mask_span_ids = torch.multinomial(uniform_dist, num_masked_spans)
+
+        # (batch_size, num_masked_spans * mask_length)
+        mask_span_ids = einops.repeat(mask_span_ids, "b n -> b (n l)", l=mask_length)
+
+        offsets = einops.repeat(
+            torch.arange(mask_length, device=device),
+            "l -> b (n l)",
+            b=batch_size,
+            n=num_masked_spans,
+        )
+
+        mask_span_ids = mask_span_ids + offsets
+
+        mask_span = torch.zeros(shape, device=device, dtype=torch.bool)
+        mask_span = mask_span.scatter_(1, mask_span_ids, True)
+
+        if attention_mask is not None:
+            # Make sure the invalid frames are not masked
+            mask_span = torch.where(attention_mask.bool(), mask_span, False)
+
+        return mask_span
+
+    @staticmethod
+    def _compute_attention_mask(length: torch.Tensor):
+        """
+        Args:
+            length (Tensor): The length of valid frames. Shape: (batch)
+            max_length (int): The maximum length of the frames.
+
+        Returns:
+            attention_mask (BoolTensor): The mask for the valid frames. `True` is invalid. Shape: (batch, num_frames)
+        """
+        max_length = length.max().item()
+
+        mask = (
+            torch.arange(max_length, device=length.device).expand(
+                length.size(0), max_length
+            )
+            >= length[:, None]
+        )
+
+        return mask
+
+    def training_step(self, batch, batch_idx):
+        loss = self(batch)
+
+        self.train_loss(loss)
+
+        if batch_idx % 100 == 0:
+            self.log("train/loss", self.train_loss, on_step=True, on_epoch=True)
+
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.AdamW(self.parameters(), lr=1e-4)

src/train.py

@@ -0,0 +1,27 @@
+import sys
+sys.path.append(".")
+
+from src.config import model as conf
+from src.model import Wav2Vec2PretrainingModule
+from src.datamodule import WebDatasetConverter, VLSP2020ForPretrainingDataModule
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+
+
+if __name__ == "__main__":
+
+    model = Wav2Vec2PretrainingModule(conf.wav2vec2_pretraining)
+    dts = WebDatasetConverter(conf.dataset.path).get_dataset()
+    dtm = VLSP2020ForPretrainingDataModule(dts, **conf.dataset)
+    trainer = Trainer(
+        callbacks=[
+            ModelCheckpoint(
+                monitor="val/loss",
+                dirpath=conf["checkpoint_dir"],
+            )
+        ],
+        gradient_clip_val=1.0,
+        accelerator="gpu"
+    )
+
+    trainer.fit(model, dtm)

src/utils/__init__.py

@@ -0,0 +1 @@
+from .functional import init_module_weights

src/utils/functional.py

@@ -0,0 +1,28 @@
+import torch
+
+
+def init_module_weights(module):
+    """Initialize the weights"""
+
+    from src.model.modules import QuantizationModule
+
+    # gumbel softmax requires special init
+    if isinstance(module, QuantizationModule):
+        module.weight_proj.weight.data.normal_(mean=0.0, std=1)
+        module.weight_proj.bias.data.zero_()
+        torch.nn.init.uniform_(module.codebooks)
+    elif isinstance(module, torch.nn.Linear):
+        # Slightly different from the TF version which uses truncated_normal for initialization
+        # cf https://github.com/pytorch/pytorch/pull/5617
+        module.weight.data.normal_(mean=0.0, std=0.5)
+    elif isinstance(module, (torch.nn.LayerNorm, torch.nn.GroupNorm)):
+        module.bias.data.zero_()
+        module.weight.data.fill_(1.0)
+    elif isinstance(module, torch.nn.Conv1d):
+        torch.nn.init.kaiming_normal_(module.weight.data)
+
+    if (
+        isinstance(module, (torch.nn.Linear, torch.nn.Conv1d))
+        and module.bias is not None
+    ):
+        module.bias.data.zero_()

src/utils/metrics.py

@@ -0,0 +1,72 @@
+from typing import Tuple, Union
+import re
+
+
+def levenshtein_distance(source: Tuple[str], target: Tuple[str]):
+    """
+    Compute the Levenshtein distance between two sequences.
+    """
+
+    n, m = len(source), len(target)
+    if n > m:
+        # Make sure n <= m, to use O(min(n,m)) space
+        source, target = target, source
+        n, m = m, n
+
+    current_row = range(n + 1)  # Keep current and previous row, not entire matrix
+    for i in range(1, m + 1):
+        previous_row, current_row = current_row, [i] + [0] * n
+        for j in range(1, n + 1):
+            add, delete, change = (
+                previous_row[j] + 1,
+                current_row[j - 1] + 1,
+                previous_row[j - 1],
+            )
+            if source[j - 1] != target[i - 1]:
+                change += 1
+            current_row[j] = min(add, delete, change)
+
+    distance = current_row[n]
+
+    del current_row
+    del previous_row
+
+    return distance
+
+
+def word_error_rate(
+    predicted: Union[str, Tuple[str]], transcript: Union[str, Tuple[str]]
+):
+    if isinstance(predicted, str):
+        predicted = (predicted,)
+    if isinstance(transcript, str):
+        transcript = (transcript,)
+
+    pattern = r"\W+"
+
+    err, total = 0, 0
+
+    for pred, tgt in zip(predicted, transcript):
+        pred_tokens = re.split(pattern, pred)
+        tgt_tokens = re.split(pattern, tgt)
+        err += levenshtein_distance(pred_tokens, tgt_tokens)
+        total += len(tgt_tokens)
+
+    return err / total
+
+
+def character_error_rate(
+    predicted: Union[str, Tuple[str]], transcript: Union[str, Tuple[str]]
+):
+    if isinstance(predicted, str):
+        predicted = (predicted,)
+    if isinstance(transcript, str):
+        transcript = (transcript,)
+
+    err, total = 0, 0
+
+    for pred, tgt in zip(predicted, transcript):
+        err += levenshtein_distance(pred, tgt)
+        total += len(tgt)
+
+    return err / total

src/utils/scheduler.py

@@ -0,0 +1,83 @@
+import math
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+class WarmUpScheduler(_LRScheduler):
+    def __init__(
+        self,
+        optimizer,
+        warmup_steps: int,
+        feature_size: int,
+        factor: float = 1.0,
+        last_epoch=-1,
+    ):
+        self.warmup_steps = warmup_steps
+        self.feature_size = feature_size
+        self.factor = factor
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        lr = self._compute_lr()
+        return [lr] * len(self.base_lrs)
+
+    def _compute_lr(self):
+        if self.last_epoch == 0:
+            return 0.0
+
+        lr = (self.feature_size ** (-0.5)) * min(
+            self.last_epoch ** (-0.5), self.last_epoch * self.warmup_steps ** (-1.5)
+        )
+
+        return lr * self.factor
+
+
+class TriStateScheduler(_LRScheduler):
+    def __init__(
+        self,
+        optimizer,
+        total_steps: int,
+        warmup_steps: int,
+        constant_steps: int,
+        factor: float = 0.3,
+        last_epoch: int = -1,
+    ):
+        self.warmup_steps = warmup_steps
+        self.constant_steps = constant_steps
+        self.total_steps = total_steps
+        self.factor = factor
+
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if not hasattr(self, "eta_min"):
+            self.eta_max = self.base_lrs.copy()
+            self.eta_min = [eta_max * self.factor for eta_max in self.eta_max]
+
+        return [
+            self._compute_lr(group["lr"], eta_min, eta_max)
+            for group, eta_min, eta_max in zip(
+                self.optimizer.param_groups, self.eta_min, self.eta_max
+            )
+        ]
+
+    def _compute_lr(self, prev_lr: float, eta_min: float, eta_max: float):
+        # first stage
+        if self.last_epoch <= self.warmup_steps:
+            lr = eta_max - 0.5 * (eta_max - eta_min) * (
+                1 + math.cos(math.pi * self.last_epoch / self.warmup_steps)
+            )
+        # second stage
+        elif self.last_epoch <= self.warmup_steps + self.constant_steps:
+            lr = prev_lr
+        else:
+            # third stage
+            decay_steps = self.total_steps - self.warmup_steps - self.constant_steps
+            k = self.last_epoch - self.warmup_steps - self.constant_steps
+            lr = eta_min + 0.5 * (eta_max - eta_min) * (
+                1 + math.cos(math.pi * k / decay_steps)
+            )
+
+        return lr
+
+    def state_dict(self) -> dict:
+        return super().state_dict()