Spaces:
Sleeping
Sleeping
# Phoneme-Level BERT for Enhanced Prosody of Text-to-Speech with Grapheme Predictions | |
### Yinghao Aaron Li, Cong Han, Xilin Jiang, Nima Mesgarani | |
> Large-scale pre-trained language models have been shown to be helpful in improving the naturalness of text-to-speech (TTS) models by enabling them to produce more naturalistic prosodic patterns. However, these models are usually word-level or sup-phoneme-level and jointly trained with phonemes, making them inefficient for the downstream TTS task where only phonemes are needed. In this work, we propose a phoneme-level BERT (PL-BERT) with a pretext task of predicting the corresponding graphemes along with the regular masked phoneme predictions. Subjective evaluations show that our phoneme-level BERT encoder has significantly improved the mean opinion scores (MOS) of rated naturalness of synthesized speech compared with the state-of-the-art (SOTA) StyleTTS baseline on out-of-distribution (OOD) texts. | |
Paper: [https://arxiv.org/abs/2301.08810](https://arxiv.org/abs/2301.08810) | |
Audio samples: [https://pl-bert.github.io/](https://pl-bert.github.io/) | |
## Pre-requisites | |
1. Python >= 3.7 | |
2. Clone this repository: | |
```bash | |
git clone https://github.com/yl4579/PL-BERT.git | |
cd PL-BERT | |
``` | |
3. Create a new environment (recommended): | |
```bash | |
conda create --name BERT python=3.8 | |
conda activate BERT | |
python -m ipykernel install --user --name BERT --display-name "BERT" | |
``` | |
4. Install python requirements: | |
```bash | |
pip install pandas singleton-decorator datasets "transformers<4.33.3" accelerate nltk phonemizer sacremoses pebble | |
``` | |
## Preprocessing | |
Please refer to the notebook [preprocess.ipynb](https://github.com/yl4579/PL-BERT/blob/main/preprocess.ipynb) for more details. The preprocessing is for English Wikipedia dataset only. I will make a new branch for Japanese if I have extra time to demostrate training on other languages. You may also refer to [#6](https://github.com/yl4579/PL-BERT/issues/6#issuecomment-1797869275) for preprocessing in other languages like Japanese. | |
## Trianing | |
Please run each cell in the notebook [train.ipynb](https://github.com/yl4579/PL-BERT/blob/main/train.ipynb). You will need to change the line | |
`config_path = "Configs/config.yml"` in cell 2 if you wish to use a different config file. The training code is in Jupyter notebook primarily because the initial epxeriment was conducted in Jupyter notebook, but you can easily make it a Python script if you want to. | |
## Finetuning | |
Here is an example of how to use it for StyleTTS finetuning. You can use it for other TTS models by replacing the text encoder with the pre-trained PL-BERT. | |
1. Modify line 683 in [models.py](https://github.com/yl4579/StyleTTS/blob/main/models.py#L683) with the following code to load BERT model in to StyleTTS: | |
```python | |
from transformers import AlbertConfig, AlbertModel | |
log_dir = "YOUR PL-BERT CHECKPOINT PATH" | |
config_path = os.path.join(log_dir, "config.yml") | |
plbert_config = yaml.safe_load(open(config_path)) | |
albert_base_configuration = AlbertConfig(**plbert_config['model_params']) | |
bert = AlbertModel(albert_base_configuration) | |
files = os.listdir(log_dir) | |
ckpts = [] | |
for f in os.listdir(log_dir): | |
if f.startswith("step_"): ckpts.append(f) | |
iters = [int(f.split('_')[-1].split('.')[0]) for f in ckpts if os.path.isfile(os.path.join(log_dir, f))] | |
iters = sorted(iters)[-1] | |
checkpoint = torch.load(log_dir + "/step_" + str(iters) + ".t7", map_location='cpu') | |
state_dict = checkpoint['net'] | |
from collections import OrderedDict | |
new_state_dict = OrderedDict() | |
for k, v in state_dict.items(): | |
name = k[7:] # remove `module.` | |
if name.startswith('encoder.'): | |
name = name[8:] # remove `encoder.` | |
new_state_dict[name] = v | |
bert.load_state_dict(new_state_dict) | |
nets = Munch(bert=bert, | |
# linear projection to match the hidden size (BERT 768, StyleTTS 512) | |
bert_encoder=nn.Linear(plbert_config['model_params']['hidden_size'], args.hidden_dim), | |
predictor=predictor, | |
decoder=decoder, | |
pitch_extractor=pitch_extractor, | |
text_encoder=text_encoder, | |
style_encoder=style_encoder, | |
text_aligner = text_aligner, | |
discriminator=discriminator) | |
``` | |
2. Modify line 126 in [train_second.py](https://github.com/yl4579/StyleTTS/blob/main/train_second.py#L126) with the following code to adjust the learning rate of BERT model: | |
```python | |
# for stability | |
for g in optimizer.optimizers['bert'].param_groups: | |
g['betas'] = (0.9, 0.99) | |
g['lr'] = 1e-5 | |
g['initial_lr'] = 1e-5 | |
g['min_lr'] = 0 | |
g['weight_decay'] = 0.01 | |
``` | |
3. Modify line 211 in [train_second.py](https://github.com/yl4579/StyleTTS/blob/main/train_second.py#L211) with the following code to replace text encoder with BERT encoder: | |
```python | |
bert_dur = model.bert(texts, attention_mask=(~text_mask).int()).last_hidden_state | |
d_en = model.bert_encoder(bert_dur).transpose(-1, -2) | |
d, _ = model.predictor(d_en, s, | |
input_lengths, | |
s2s_attn_mono, | |
m) | |
``` | |
[line 257](https://github.com/yl4579/StyleTTS/blob/main/train_second.py#L257): | |
```python | |
_, p = model.predictor(d_en, s, | |
input_lengths, | |
s2s_attn_mono, | |
m) | |
``` | |
and [line 415](https://github.com/yl4579/StyleTTS/blob/main/train_second.py#L415): | |
```python | |
bert_dur = model.bert(texts, attention_mask=(~text_mask).int()).last_hidden_state | |
d_en = model.bert_encoder(bert_dur).transpose(-1, -2) | |
d, p = model.predictor(d_en, s, | |
input_lengths, | |
s2s_attn_mono, | |
m) | |
``` | |
4. Modify line 347 in [train_second.py](https://github.com/yl4579/StyleTTS/blob/main/train_second.py#L347) with the following code to make sure parameters of BERT model are updated: | |
```python | |
optimizer.step('bert_encoder') | |
optimizer.step('bert') | |
``` | |
The pre-trained PL-BERT on Wikipedia for 1M steps can be downloaded at: [PL-BERT link](https://drive.google.com/file/d/19gzPmWKdmakeVszSNuUtVMMBaFYMQqJ7/view?usp=sharing). | |
The demo on LJSpeech dataset along with the pre-modified StyleTTS repo and pre-trained models can be downloaded here: [StyleTTS Link](https://drive.google.com/file/d/18DU4JrW1rhySrIk-XSxZkXt2MuznxoM-/view?usp=sharing). This zip file contains the code modification above, the pre-trained PL-BERT model listed above, pre-trained StyleTTS w/ PL-BERT, pre-trained StyleTTS w/o PL-BERT and pre-trained HifiGAN on LJSpeech from the StyleTTS repo. | |
## References | |
- [NVIDIA/NeMo-text-processing](https://github.com/NVIDIA/NeMo-text-processing) | |
- [tomaarsen/TTSTextNormalization](https://github.com/tomaarsen/TTSTextNormalization) | |