S2ST-Transformer

Direct Speech-to-Speech Translation using a unit-based Transformer model with a pre-trained Conformer encoder.

The pretrained model is trained on a Persian-English dataset and can translate speech from Persian to English, but you can train the model for any language pair with a CVSS-like dataset.

The model is built using Speechbrain toolkit.

🤗 Model repo on Huggingface

Inference

You can use the Inference Notebook or directly use the inference class implemented in the Inference File for inference with the pre-trained model.

import torch
import torchaudio
from inference import S2STInference
from speechbrain.inference.vocoders import UnitHIFIGAN

device = 'cuda' if torch.cuda.is_available() else 'cpu'

s2st = S2STInference.from_hparams(source="sinarashidi/s2st_fa-en_augmented", savedir="tmpdir_s2st", run_opts={'device': device})
hifi_gan_unit = UnitHIFIGAN.from_hparams(source="sinarashidi/unit_hifigan_ljspeech", savedir="tmpdir_vocoder", run_opts={'device': device})
codes = s2st.translate_file("audio/test2.mp3")
codes = torch.IntTensor(codes)

waveforms = hifi_gan_unit.decode_unit(codes)
torchaudio.save(f"translated_file.wav",waveforms.squeeze(1).cpu(), 16000)

Training

Steps for training:

1. Setup environment:

    python -m venv .venv
    source .venv/bin/activate
    pip install -r requirements.txt

2. Download the Common Voice dataset
3. Download the CVSS dataset
4. Change the src_data_folder and tgt_data_folder in the hyperparams.yaml file according to your setup
5. Change other training parameters or model configs in the hyperparams.yaml file if needed
6. Start the training procedure:

    python train.py hyperparams.yaml --precision=bf16

With default settings, each epoch takes about 2 hours on the RTX 3090 GPU.

To run only the evaluation stage on a checkpoint, run the following command:

python train.py hyperparams.yaml --precision=bf16 --test_only