This repository contains an implementation of Google's BERT-based Speech pre-Training with Random-projection Quantizer (BEST-RQ) (Arxiv) pre-training objective. Both the original method as well as the more recent adaption of BEST-RQ in the Universal Speech Model (Arxiv) are implemented.

Further, the repository extends the OpenAI Whisper repository by implementing an IA^3 (Arxiv) trainable Whisper implementation. At the moment, IA^3 is only implemented for the AudioEncoder.

The repository is currently work in progress with coming features still to be implemented.

An easy-to-use Dockerfile is provided. After cloning the repository, the Docker image can be built using

docker build -t best_rq .

and a corresponding container can be started using

docker run -e <WANDB_API_KEY> --gpus all -it best_rq

where <WANDB_API_KEY> should be replaced with a Weights and Biases API key. If a valid API key is provided, metrics and (intermediate) checkpoints are logged automatically to Weights and Biases. As IA^3 requires only a tiny amount of trainable parameters, checkpoints are (individually) negligible in space requirements.

The Dockerfile contains a micromamba installation, so further libraries/requirements can easily be installed on the fly. As a GPU is highly recommended for running the training, the image is only tested on a machine with an installed GPU.

Alternatively, the environment can be directly installed using an existing conda (or mamba/micromamba) installation on a unix system using

conda create -y --file environment.yaml

IA^3-based training of a (pretrained) Whisper model using BEST-RQ can be launched using the train command within a Docker container. E.g., a training using a batch size of 4, 16 BEST-RQ codebooks, and gradient accumulation over 256 batches can be launched like this:

train --batch_size 4 --num_codebooks 16 --accumulate_gradients 256

The full list of possible hyper-parameters can be obtained via

train --help