ArTST

This repository contains the implementation of the paper:

ArTST: Arabic Text and Speech Transformer

     

Hawau Olamide Toyin ^{* 1}   Amirbek Djanibekov ^{* 1}  Ajinkya Kulkarni ¹  Hanan Aldarmaki ¹ 

^* equal contribution   ¹ MBZUAI  

ArabicNLP 2023

ArTST

ArTST, a pre-trained Arabic text and speech transformer for supporting open-source speech technologies for the Arabic language. The model architecture in this first edition follows the unified-modal framework, SpeechT5, that was recently released for English, and is focused on Modern Standard Arabic (MSA), with plans to extend the model for dialectal and code-switched Arabic in future editions. We pre-trained the model from scratch on MSA speech and text data, and fine-tuned it for the following tasks: Automatic Speech Recognition (ASR), Text-To-Speech synthesis (TTS), and spoken dialect identification.

Update

• February, 2024: Bug fix with checkpoint loading
• December, 2023: Released ArTST ASR demo HF-Spaces
• November, 2023: Released ArTST TTS demo HF-Spaces
• October, 2023: Open-sourced model's weight to HuggingFace
• October, 2023: ArTST was accepted by EMNLP (ArabicNLP conference) 2023.

Checkpoints

Pre-Trained Models

Model	Pre-train Dataset	Model	Tokenizer
ArTST base	MGB2	Hugging Face	Hugging Face

Finetuned Models

Model	FInetune Dataset	Model	Tokenizer
ArTST ASR	MGB2	Hugging Face	Hugging Face
ArTST TTS	ClArTTS	Hugging Face	Hugging Face
ArTST* TTS	ClArTTS	Hugging Face	Hugging Face

Environment & Installation

Python version: 3.8+

1. Clone this repo

cd ArTST
conda create -n artst python=3.8
conda activate artst
pip install -r requirements.txt

2. Install fairseq

git clone https://github.com/pytorch/fairseq
cd fairseq
pip install --editable ./
python setup.py build_ext --inplace

3. Download Checkpoints

# Make sure you have git-lfs installed (https://git-lfs.com)
git lfs install
git clone https://huggingface.co/MBZUAI/ArTST

Loading Model

import torch
from artst.tasks.artst import ArTSTTask
from artst.models.artst import ArTSTTransformerModel

checkpoint = torch.load('checkpoint.pt')
checkpoint['cfg']['task'].t5_task = 't2s' # or "s2t" for asr
task = ArTSTTask.setup_task(checkpoint['cfg']['task'])

model = ArTSTTransformerModel.build_model(checkpoint['cfg']['model'], task)
model.load_state_dict(checkpoint['model'])

Data Preparation

Speech

For pretraining, follow the steps for preparing wav2vec 2.0 manifest here and preparing HuBERT label here.

For finetuning TTS task, an extra column is required in the speech manifest file for speaker embedding. To generate speaker embedding, we use speech brain. Here is a DATA_ROOT sample folder structure that contains manifest samples.

Text

Pretrain:

Please use fairseq-preprocess to generate the index and bin files for the text data. We use sentencepiece to pre-process the text, we've provided our SPM models and dictionary in this repo. You need to use the SPM model to process the text and then use fairseq-preprocess with the provided dictionary to get the index and bin files. Note that after SPM processes sentences, the resulting text should have individual characters separated by space.

For Finetuning, a simple text file containing corresponding texts on each line suffices. See here for sample manifest. Normalize the texts as we did for training/evaluation using this script.

Training

The bash files contain the parameters and hyperparameters used for pretraining and finetuning. Find more details on training arguments here

Pretrain

bash /scripts/pretrain/train.sh

Finetune

ASR

bash /scripts/ASR/finetune.sh

TTS

bash /scripts/TTS/finetune.sh

Inference

ASR

bash /scripts/ASR/inference.sh

TTS

bash /scripts/TTS/inference.sh

Acknowledgements

ArTST is built on SpeechT5 Architecture. If you use any of ArTST models, please cite

@inproceedings{toyin2023artst,
  title={ArTST: Arabic Text and Speech Transformer},
  author={Toyin, Hawau and Djanibekov, Amirbek and Kulkarni, Ajinkya and Aldarmaki, Hanan},
  booktitle={Proceedings of ArabicNLP 2023},
  pages={41--51},
  year={2023}
}