Skip to content

GitLab

Commit 5eae102f authored by Yoach Lacombe's avatar Yoach Lacombe
Browse files

Merge branch 'parler-tts-release' into main

parents 85b8cac7 86e4eb71
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 1207 additions and 1702 deletions
LICENSE
View file @ 5eae102f
......@@ -186,7 +186,7 @@
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Copyright [2024] [The HuggingFace Inc. team]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
......
README.md
View file @ 5eae102f
# Stable Speech
Work in-progress reproduction of the text-to-speech (TTS) model from the paper [Natural language guidance of high-fidelity text-to-speech with synthetic annotations](https://www.text-description-to-speech.com)
by Dan Lyth and Simon King, from Stability AI and Edinburgh University respectively.
Reproducing the TTS model requires the following 5 steps to be completed in order:
1. Train the Accent Classifier
2. Annotate the Training Set
3. Aggregate Statistics
4. Create Descriptions
5. Train the Model
## Step 1: Train the Accent Classifier
The script [`run_audio_classification.py`](run_audio_classification.py) can be used to train an audio encoder model from
the [Transformers library](https://github.com/huggingface/transformers) (e.g. Wav2Vec2, MMS, or Whisper) for the accent
classification task.
Starting with a pre-trained audio encoder model, a simple linear classifier is appended to the last hidden-layer to map the
audio embeddings to class label predictions. The audio encoder can either be frozen (`--freeze_base_model`) or trained.
The linear classifier is randomly initialised, and is thus always trained.
The script can be used to train on a single accent dataset, or a combination of datasets, which should be specified by
separating dataset names, configs and splits by the `+` character in the launch command (see below for an example).
In the proceeding example, we follow Stability's approach by taking audio embeddings from a frozen [MMS-LID](https://huggingface.co/facebook/mms-lid-126)
model, and training the linear classifier on a combination of three open-source datasets:
1. The English Accented (`en_accented`) subset of [Voxpopuli](https://huggingface.co/datasets/facebook/voxpopuli)
2. The train split of [VCTK](https://huggingface.co/datasets/vctk)
3. The dev split of [EdAcc](https://huggingface.co/datasets/edinburghcstr/edacc)
The model is subsequently evaluated on the test split of [EdAcc](https://huggingface.co/datasets/edinburghcstr/edacc)
to give the final classification accuracy.
```bash
#!/usr/bin/env bash
python run_audio_classification.py \
--model_name_or_path "facebook/mms-lid-126" \
--train_dataset_name "vctk+facebook/voxpopuli+edinburghcstr/edacc" \
--train_dataset_config_name "main+en_accented+default" \
--train_split_name "train+test+validation" \
--train_label_column_name "accent+accent+accent" \
--eval_dataset_name "edinburghcstr/edacc" \
--eval_dataset_config_name "default" \
--eval_split_name "test" \
--eval_label_column_name "accent" \
--output_dir "./" \
--do_train \
--do_eval \
--overwrite_output_dir \
--remove_unused_columns False \
--fp16 \
--learning_rate 1e-4 \
--max_length_seconds 20 \
--attention_mask False \
--warmup_ratio 0.1 \
--num_train_epochs 5 \
--per_device_train_batch_size 32 \
--per_device_eval_batch_size 32 \
--preprocessing_num_workers 16 \
--dataloader_num_workers 4 \
--logging_strategy "steps" \
--logging_steps 10 \
--evaluation_strategy "epoch" \
--save_strategy "epoch" \
--load_best_model_at_end True \
--metric_for_best_model "accuracy" \
--save_total_limit 3 \
--freeze_base_model \
--push_to_hub \
--trust_remote_code
# Parler-TTS
Parler-TTS is a lightweight text-to-speech (TTS) model that can generate high-quality, natural sounding speech in the style of a given speaker (gender, pitch, speaking style, etc). It is a reproduction of work from the paper [Natural language guidance of high-fidelity text-to-speech with synthetic annotations](https://www.text-description-to-speech.com) by Dan Lyth and Simon King, from Stability AI and Edinburgh University respectively.
Contrarily to other TTS models, Parler-TTS is a **fully open-source** release. All of the datasets, pre-processing, training code and weights are released publicly under permissive license, enabling the community to build on our work and develop their own powerful TTS models.
This repository contains the inference and training code for Parler-TTS. It is designed to accompany the [Data-Speech](https://github.com/huggingface/dataspeech) repository for dataset annotation.
> [!IMPORTANT]
> We're proud to release Parler-TTS v0.1, our first 300M parameter model, trained on 10.5K hours of audio data.
> In the coming weeks, we'll be working on scaling up to 50k hours of data, in preparation for the v1 model.
## 📖 Quick Index
* [Installation](#installation)
* [Usage](#usage)
* [Training](#training)
* [Demo](https://huggingface.co/spaces/parler-tts/parler_tts_mini)
* [Model weights and datasets](https://huggingface.co/parler-tts)
## Usage
> [!TIP]
> You can directly try it out in an interactive demo [here](https://huggingface.co/spaces/parler-tts/parler_tts_mini)!
Using Parler-TTS is as simple as "bonjour". Simply use the following inference snippet.
```py
from parler_tts import ParlerTTSForConditionalGeneration
from transformers import AutoTokenizer
import soundfile as sf
import torch
device = "cuda:0" if torch.cuda.is_available() else "cpu"
model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler_tts_300M_v0.1").to(device)
tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler_tts_300M_v0.1")
prompt = "Hey, how are you doing today?"
description = "A female speaker with a slightly low-pitched voice delivers her words quite expressively, in a very confined sounding environment with clear audio quality. She speaks very fast."
input_ids = tokenizer(description, return_tensors="pt").input_ids.to(device)
prompt_input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
generation = model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids)
audio_arr = generation.cpu().numpy().squeeze()
sf.write("parler_tts_out.wav", audio_arr, model.config.sampling_rate)
```
## Installation
Parler-TTS has light-weight dependencies and can be installed in one line:
```sh
pip install git+https://github.com/huggingface/parler-tts.git
```
## Training
The [training folder](/training/) contains all the information to train or fine-tune your own Parler-TTS model. It consists of:
- [1. An introduction to the Parler-TTS architecture](/training/README.md#1-architecture)
- [2. The first steps to get started](/training/README.md#2-getting-started)
- [3. A training guide](/training/README.md#3-training)
> [!IMPORTANT]
> **TL;DR:** After having followed the [installation steps](/training/README.md#requirements), you can reproduce the Parler-TTS v0.1 training recipe with the following command line:
```sh
accelerate launch ./training/run_parler_tts_training.py ./helpers/training_configs/starting_point_0.01.json
```
Tips:
1. **Number of labels:** normalisation should be applied to the target class labels to group linguistically similar accents together (e.g. "Southern Irish" and "Irish" should both be "Irish"). This helps _balance_ the dataset by removing labels with very few examples. You can modify the function `preprocess_labels` to implement any custom normalisation strategy.
## Acknowledgements
## Step 2: Annotate the Training Set
This library builds on top of a number of open-source giants, to whom we'd like to extend our warmest thanks for providing these tools!
Annotate the training dataset with information on: SNR, C50, pitch and speaking rate.
Special thanks to:
- Dan Lyth and Simon King, from Stability AI and Edinburgh University respectively, for publishing such a promising and clear research paper: [Natural language guidance of high-fidelity text-to-speech with synthetic annotations](https://arxiv.org/abs/2402.01912).
- the many libraries used, namely [🤗 datasets](https://huggingface.co/docs/datasets/v2.17.0/en/index), [🤗 accelerate](https://huggingface.co/docs/accelerate/en/index), [jiwer](https://github.com/jitsi/jiwer), [wandb](https://wandb.ai/), and [🤗 transformers](https://huggingface.co/docs/transformers/index).
- Descript for the [DAC codec model](https://github.com/descriptinc/descript-audio-codec)
- Hugging Face 🤗 for providing compute resources and time to explore!
## Step 3: Aggregate Statistics
Aggregate statistics from Step 2. Convert continuous values to discrete labels.
## Citation
## Step 4: Create Descriptions
If you found this repository useful, please consider citing this work and also the original Stability AI paper:
Convert sequence of discrete labels to text description (using an LLM).
```
@misc{lacombe-etal-2024-parler-tts,
author = {Yoach Lacombe and Vaibhav Srivastav and Sanchit Gandhi},
title = {Parler-TTS},
year = {2024},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/huggingface/parler-tts}}
}
```
## Step 5: Train the Model
```
@misc{lyth2024natural,
title={Natural language guidance of high-fidelity text-to-speech with synthetic annotations},
author={Dan Lyth and Simon King},
year={2024},
eprint={2402.01912},
archivePrefix={arXiv},
primaryClass={cs.SD}
}
```
Train MusicGen-style model on the TTS task.
## Contribution
Contributions are welcome, as the project offers many possibilities for improvement and exploration.
Namely, we're looking at ways to improve both quality and speed:
- Datasets:
- Train on more data
- Add more features such as accents
- Training:
- Add PEFT compatibility to do Lora fine-tuning.
- Add possibility to train without description column.
- Add notebook training.
- Explore multilingual training.
- Explore mono-speaker finetuning.
- Explore more architectures.
- Optimization:
- Compilation and static cache
- Support to FA2 and SDPA
- Evaluation:
- Add more evaluation metrics
audio_classification_scripts/run_wav2vec2_dummy.sh deleted 100644 → 0
View file @ 85b8cac7
#!/usr/bin/env bash
python run_audio_classification.py \
--model_name_or_path "hf-internal-testing/tiny-random-wav2vec2" \
--train_dataset_name "facebook/voxpopuli" \
--train_dataset_config_name "en_accented" \
--train_split_name "test" \
--train_label_column_name "accent" \
--eval_dataset_name "facebook/voxpopuli" \
--eval_dataset_config_name "en_accented" \
--eval_split_name "test" \
--eval_label_column_name "accent" \
--trust_remote_code \
--output_dir "./" \
--do_train \
--do_eval \
--max_train_samples 100 \
--max_eval_samples 100 \
--overwrite_output_dir \
--remove_unused_columns False \
--fp16 \
--learning_rate 1e-4 \
--min_length_seconds 5 \
--max_length_seconds 10 \
--attention_mask False \
--warmup_ratio 0.1 \
--num_train_epochs 5 \
--per_device_train_batch_size 4 \
--per_device_eval_batch_size 4 \
--dataloader_num_workers 0 \
--logging_strategy "steps" \
--logging_steps 10 \
--evaluation_strategy "epoch" \
--save_strategy "epoch" \
--load_best_model_at_end True \
--metric_for_best_model "accuracy" \
--save_total_limit 3 \
--seed 0
helpers/gradio_demo/app.py 0 → 100644
View file @ 5eae102f
import gradio as gr
import torch
from parler_tts import ParlerTTSForConditionalGeneration
from transformers import AutoTokenizer, AutoFeatureExtractor, set_seed
device = "cuda:0" if torch.cuda.is_available() else "cpu"
repo_id = "parler-tts/parler_tts_300M_v0.1"
model = ParlerTTSForConditionalGeneration.from_pretrained(repo_id).to(device)
tokenizer = AutoTokenizer.from_pretrained(repo_id)
feature_extractor = AutoFeatureExtractor.from_pretrained(repo_id)
SAMPLE_RATE = feature_extractor.sampling_rate
SEED = 41
default_text = "Please surprise me and speak in whatever voice you enjoy."
title = "# Parler-TTS </div>"
examples = [
[
"'This is the best time of my life, Bartley,' she said happily.",
"A female speaker with a slightly low-pitched, quite monotone voice delivers her words at a slightly faster-than-average pace in a confined space with very clear audio.",
],
[
"Montrose also, after having experienced still more variety of good and bad fortune, threw down his arms, and retired out of the kingdom. ",
"A male speaker with a slightly high-pitched voice delivering his words at a slightly slow pace in a small, confined space with a touch of background noise and a quite monotone tone.",
],
[
"montrose also after having experienced still more variety of good and bad fortune threw down his arms and retired out of the kingdom",
"A male speaker with a low-pitched voice delivering his words at a fast pace in a small, confined space with a lot of background noise and an animated tone.",
],
]
def gen_tts(text, description):
inputs = tokenizer(description, return_tensors="pt").to(device)
prompt = tokenizer(text, return_tensors="pt").to(device)
set_seed(SEED)
generation = model.generate(
input_ids=inputs.input_ids, prompt_input_ids=prompt.input_ids, do_sample=True, temperature=1.0
)
audio_arr = generation.cpu().numpy().squeeze()
return (SAMPLE_RATE, audio_arr)
css = """
#share-btn-container {
display: flex;
padding-left: 0.5rem !important;
padding-right: 0.5rem !important;
background-color: #000000;
justify-content: center;
align-items: center;
border-radius: 9999px !important;
width: 13rem;
margin-top: 10px;
margin-left: auto;
flex: unset !important;
}
#share-btn {
all: initial;
color: #ffffff;
font-weight: 600;
cursor: pointer;
font-family: 'IBM Plex Sans', sans-serif;
margin-left: 0.5rem !important;
padding-top: 0.25rem !important;
padding-bottom: 0.25rem !important;
right:0;
}
#share-btn * {
all: unset !important;
}
#share-btn-container div:nth-child(-n+2){
width: auto !important;
min-height: 0px !important;
}
#share-btn-container .wrap {
display: none !important;
}
"""
with gr.Blocks(css=css) as block:
gr.Markdown(title)
with gr.Row():
with gr.Column():
input_text = gr.Textbox(label="Input Text", lines=2, value=default_text, elem_id="input_text")
description = gr.Textbox(label="Description", lines=2, value="", elem_id="input_description")
run_button = gr.Button("Generate Audio", variant="primary")
with gr.Column():
audio_out = gr.Audio(label="Parler-TTS generation", type="numpy", elem_id="audio_out")
inputs = [input_text, description]
outputs = [audio_out]
gr.Examples(examples=examples, fn=gen_tts, inputs=inputs, outputs=outputs, cache_examples=True)
run_button.click(fn=gen_tts, inputs=inputs, outputs=outputs, queue=True)
block.queue()
block.launch(share=True)
helpers/model_init_scripts/init_dummy_model.py 0 → 100644
View file @ 5eae102f
from parler_tts import ParlerTTSForCausalLM, ParlerTTSForConditionalGeneration, ParlerTTSDecoderConfig
from transformers import AutoConfig
import os
import argparse
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("save_directory", type=str, help="Directory where to save the model and the decoder.")
parser.add_argument("--text_model", type=str, help="Repository id or path to the text encoder.")
parser.add_argument("--audio_model", type=str, help="Repository id or path to the audio encoder.")
args = parser.parse_args()
text_model = args.text_model
encodec_version = args.audio_model
t5 = AutoConfig.from_pretrained(text_model)
encodec = AutoConfig.from_pretrained(encodec_version)
encodec_vocab_size = encodec.codebook_size
num_codebooks = encodec.num_codebooks
print("num_codebooks", num_codebooks)
decoder_config = ParlerTTSDecoderConfig(
vocab_size=encodec_vocab_size + 1,
max_position_embeddings=2048,
num_hidden_layers=4,
ffn_dim=512,
num_attention_heads=8,
layerdrop=0.0,
use_cache=True,
activation_function="gelu",
hidden_size=512,
dropout=0.0,
attention_dropout=0.0,
activation_dropout=0.0,
pad_token_id=encodec_vocab_size,
eos_token_id=encodec_vocab_size,
bos_token_id=encodec_vocab_size + 1,
num_codebooks=num_codebooks,
)
decoder = ParlerTTSForCausalLM(decoder_config)
decoder.save_pretrained(os.path.join(args.save_directory, "decoder"))
model = ParlerTTSForConditionalGeneration.from_sub_models_pretrained(
text_encoder_pretrained_model_name_or_path=text_model,
audio_encoder_pretrained_model_name_or_path=encodec_version,
decoder_pretrained_model_name_or_path=os.path.join(args.save_directory, "decoder"),
vocab_size=t5.vocab_size,
)
# set the appropriate bos/pad token ids
model.generation_config.decoder_start_token_id = encodec_vocab_size + 1
model.generation_config.pad_token_id = encodec_vocab_size
model.generation_config.eos_token_id = encodec_vocab_size
# set other default generation config params
model.generation_config.max_length = int(30 * model.audio_encoder.config.frame_rate)
model.generation_config.do_sample = True # True
model.generation_config.guidance_scale = 1 # 3.0
model.config.pad_token_id = encodec_vocab_size
model.config.decoder_start_token_id = encodec_vocab_size+1
model.save_pretrained(os.path.join(args.save_directory, "tiny-model"))
helpers/model_init_scripts/init_dummy_model_with_encodec.py 0 → 100644
View file @ 5eae102f
from parler_tts import ParlerTTSForCausalLM, ParlerTTSForConditionalGeneration, ParlerTTSDecoderConfig
from transformers import AutoConfig
import os
import argparse
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("save_directory", type=str, help="Directory where to save the model and the decoder.")
args = parser.parse_args()
text_model = "google-t5/t5-small"
encodec_version = "facebook/encodec_24khz"
t5 = AutoConfig.from_pretrained(text_model)
encodec = AutoConfig.from_pretrained(encodec_version)
encodec_vocab_size = encodec.codebook_size
num_codebooks = 8
print("num_codebooks", num_codebooks)
decoder_config = ParlerTTSDecoderConfig(
vocab_size=encodec_vocab_size + 1,
max_position_embeddings=2048,
num_hidden_layers=4,
ffn_dim=512,
num_attention_heads=8,
layerdrop=0.0,
use_cache=True,
activation_function="gelu",
hidden_size=512,
dropout=0.0,
attention_dropout=0.0,
activation_dropout=0.0,
pad_token_id=encodec_vocab_size,
eos_token_id=encodec_vocab_size,
bos_token_id=encodec_vocab_size + 1,
num_codebooks=num_codebooks,
)
decoder = ParlerTTSForCausalLM(decoder_config)
decoder.save_pretrained(os.path.join(args.save_directory, "decoder"))
model = ParlerTTSForConditionalGeneration.from_sub_models_pretrained(
text_encoder_pretrained_model_name_or_path=text_model,
audio_encoder_pretrained_model_name_or_path=encodec_version,
decoder_pretrained_model_name_or_path=os.path.join(args.save_directory, "decoder"),
vocab_size=t5.vocab_size,
)
# set the appropriate bos/pad token ids
model.generation_config.decoder_start_token_id = encodec_vocab_size + 1
model.generation_config.pad_token_id = encodec_vocab_size
model.generation_config.eos_token_id = encodec_vocab_size
# set other default generation config params
model.generation_config.max_length = int(30 * model.audio_encoder.config.frame_rate)
model.generation_config.do_sample = True # True
model.generation_config.guidance_scale = 1 # 3.0
model.save_pretrained(os.path.join(args.save_directory, "tiny-model"))
helpers/model_init_scripts/init_model_300M.py 0 → 100644
View file @ 5eae102f
from parler_tts import ParlerTTSForCausalLM, ParlerTTSForConditionalGeneration, ParlerTTSDecoderConfig
from transformers import AutoConfig
import os
import argparse
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("save_directory", type=str, help="Directory where to save the model and the decoder.")
parser.add_argument("--text_model", type=str, help="Repository id or path to the text encoder.")
parser.add_argument("--audio_model", type=str, help="Repository id or path to the audio encoder.")
args = parser.parse_args()
text_model = args.text_model
encodec_version = args.audio_model
t5 = AutoConfig.from_pretrained(text_model)
encodec = AutoConfig.from_pretrained(encodec_version)
encodec_vocab_size = encodec.codebook_size
num_codebooks = encodec.num_codebooks
print("num_codebooks", num_codebooks)
decoder_config = ParlerTTSDecoderConfig(
vocab_size=encodec_vocab_size + 64, # + 64 instead of +1 to have a multiple of 64
max_position_embeddings=4096, # 30 s = 2580
num_hidden_layers=24,
ffn_dim=4096,
num_attention_heads=16,
layerdrop=0.0,
use_cache=True,
activation_function="gelu",
hidden_size=1024,
dropout=0.1,
attention_dropout=0.0,
activation_dropout=0.0,
pad_token_id=encodec_vocab_size,
eos_token_id=encodec_vocab_size,
bos_token_id=encodec_vocab_size + 1,
num_codebooks=num_codebooks,
)
decoder = ParlerTTSForCausalLM(decoder_config)
decoder.save_pretrained(os.path.join(args.save_directory, "decoder"))
model = ParlerTTSForConditionalGeneration.from_sub_models_pretrained(
text_encoder_pretrained_model_name_or_path=text_model,
audio_encoder_pretrained_model_name_or_path=encodec_version,
decoder_pretrained_model_name_or_path=os.path.join(args.save_directory, "decoder"),
vocab_size=t5.vocab_size,
)
# set the appropriate bos/pad token ids
model.generation_config.decoder_start_token_id = encodec_vocab_size + 1
model.generation_config.pad_token_id = encodec_vocab_size
model.generation_config.eos_token_id = encodec_vocab_size
# set other default generation config params
model.generation_config.max_length = int(30 * model.audio_encoder.config.frame_rate)
model.generation_config.do_sample = True # True
model.generation_config.guidance_scale = 1 # 3.0
model.config.pad_token_id = encodec_vocab_size
model.config.decoder_start_token_id = encodec_vocab_size+1
model.save_pretrained(os.path.join(args.save_directory, "parler-tts-untrained-300M/"))
helpers/push_to_hub_scripts/push_dac_to_hub.py 0 → 100644
View file @ 5eae102f
import dac
from parler_tts import DACConfig, DACModel
from transformers import AutoConfig, AutoModel
from transformers import EncodecFeatureExtractor
AutoConfig.register("dac", DACConfig)
AutoModel.register(DACConfig, DACModel)
# Download a model
model_path = dac.utils.download(model_type="44khz")
model = dac.DAC.load(model_path)
hf_dac = DACModel(DACConfig())
hf_dac.model.load_state_dict(model.state_dict())
hf_dac.push_to_hub("parler-tts/dac_44khZ_8kbps")
EncodecFeatureExtractor(sampling_rate=44100).push_to_hub("parler-tts/dac_44khZ_8kbps")
helpers/push_to_hub_scripts/push_trained_parler_tts_to_hub.py 0 → 100644
View file @ 5eae102f
from parler_tts import ParlerTTSForConditionalGeneration
from transformers import AutoTokenizer, AutoFeatureExtractor
path = "TODO"
repo_id = "parler_tts_300M"
AutoFeatureExtractor.from_pretrained("ylacombe/dac_44khZ_8kbps").push_to_hub(repo_id)
AutoTokenizer.from_pretrained("google/t5-v1_1-base").push_to_hub(repo_id)
ParlerTTSForConditionalGeneration.from_pretrained(path).push_to_hub(repo_id)
helpers/training_configs/librispeech_tts_r_300M_dummy.json 0 → 100644
View file @ 5eae102f
{
"model_name_or_path": "./parler-tts-untrained-300M/parler-tts-untrained-300M/",
"save_to_disk": "./tmp_dataset_audio/",
"temporary_save_to_disk": "./audio_code_tmp/",
"feature_extractor_name":"ylacombe/dac_44khZ_8kbps",
"description_tokenizer_name":"google/flan-t5-base",
"prompt_tokenizer_name":"google/flan-t5-base",
"report_to": ["wandb"],
"overwrite_output_dir": true,
"output_dir": "./output_dir_training",
"train_dataset_name": "blabble-io/libritts_r",
"train_metadata_dataset_name": "parler-tts/libritts_r_tags_tagged_10k_generated",
"train_dataset_config_name": "clean",
"train_split_name": "test.clean",
"eval_dataset_name": "blabble-io/libritts_r",
"eval_metadata_dataset_name": "parler-tts/libritts_r_tags_tagged_10k_generated",
"eval_dataset_config_name": "clean",
"eval_split_name": "test.clean",
"target_audio_column_name": "audio",
"description_column_name": "text_description",
"prompt_column_name": "text",
"max_eval_samples": 48,
"max_train_samples": 96,
"max_duration_in_seconds": 20,
"min_duration_in_seconds": 2.0,
"add_audio_samples_to_wandb": true,
"id_column_name": "id",
"preprocessing_num_workers": 8,
"do_train": true,
"num_train_epochs": 50,
"gradient_accumulation_steps": 1,
"gradient_checkpointing": false,
"per_device_train_batch_size": 4,
"learning_rate": 1e-3,
"adam_beta1": 0.9,
"adam_beta2": 0.99,
"weight_decay": 0.01,
"lr_scheduler_type": "cosine",
"warmup_steps": 40,
"logging_steps": 2,
"freeze_text_encoder": true,
"do_eval": true,
"predict_with_generate": true,
"include_inputs_for_metrics": true,
"evaluation_strategy": "steps",
"eval_steps": 500,
"save_steps": 5000,
"per_device_eval_batch_size": 12,
"audio_encoder_per_device_batch_size":24,
"dtype": "bfloat16",
"seed": 456,
"dataloader_num_workers":8
}
helpers/training_configs/starting_point_0.01.json 0 → 100644
View file @ 5eae102f
{
"model_name_or_path": "./parler-tts-untrained-300M/parler-tts-untrained-300M/",
"save_to_disk": "./tmp_dataset_audio/",
"temporary_save_to_disk": "./audio_code_tmp/",
"feature_extractor_name":"ylacombe/dac_44khZ_8kbps",
"description_tokenizer_name":"google/flan-t5-base",
"prompt_tokenizer_name":"google/flan-t5-base",
"report_to": ["wandb"],
"overwrite_output_dir": true,
"output_dir": "./output_dir_training",
"train_dataset_name": "blabble-io/libritts_r+blabble-io/libritts_r+blabble-io/libritts_r+parler-tts/mls_eng_10k",
"train_metadata_dataset_name": "parler-tts/libritts_r_tags_tagged_10k_generated+parler-tts/libritts_r_tags_tagged_10k_generated+parler-tts/libritts_r_tags_tagged_10k_generated+parler-tts/mls-eng-10k-tags_tagged_10k_generated",
"train_dataset_config_name": "clean+clean+other+default",
"train_split_name": "train.clean.360+train.clean.100+train.other.500+train",
"eval_dataset_name": "blabble-io/libritts_r+parler-tts/mls_eng_10k",
"eval_metadata_dataset_name": "parler-tts/libritts_r_tags_tagged_10k_generated+parler-tts/mls-eng-10k-tags_tagged_10k_generated",
"eval_dataset_config_name": "other+default",
"eval_split_name": "test.other+test",
"target_audio_column_name": "audio",
"description_column_name": "text_description",
"prompt_column_name": "text",
"max_eval_samples": 96,
"max_duration_in_seconds": 30,
"min_duration_in_seconds": 2.0,
"max_text_length": 400,
"group_by_length": true,
"add_audio_samples_to_wandb": true,
"id_column_name": "id",
"preprocessing_num_workers": 8,
"do_train": true,
"num_train_epochs": 40,
"gradient_accumulation_steps": 8,
"gradient_checkpointing": false,
"per_device_train_batch_size": 3,
"learning_rate": 0.00095,
"adam_beta1": 0.9,
"adam_beta2": 0.99,
"weight_decay": 0.01,
"lr_scheduler_type": "constant_with_warmup",
"warmup_steps": 20000,
"logging_steps": 1000,
"freeze_text_encoder": true,
"do_eval": true,
"predict_with_generate": true,
"include_inputs_for_metrics": true,
"evaluation_strategy": "steps",
"eval_steps": 10000,
"save_steps": 10000,
"per_device_eval_batch_size": 12,
"audio_encoder_per_device_batch_size":20,
"dtype": "bfloat16",
"seed": 456,
"dataloader_num_workers":8
}
parler_tts/__init__.py 0 → 100644
View file @ 5eae102f
__version__ = "0.1"
from .configuration_parler_tts import ParlerTTSConfig, ParlerTTSDecoderConfig
from .modeling_parler_tts import (
ParlerTTSForCausalLM,
ParlerTTSForConditionalGeneration,
apply_delay_pattern_mask,
build_delay_pattern_mask,
)
from .dac_wrapper import DACConfig, DACModel
from transformers import AutoConfig, AutoModel
AutoConfig.register("dac", DACConfig)
AutoModel.register(DACConfig, DACModel)
stable_speech/configuration_stable_speech.py → parler_tts/configuration_parler_tts.py
View file @ 5eae102f
# coding=utf-8
# Copyright 2023 Meta AI and The HuggingFace Inc. team. All rights reserved.
# Copyright 2024 and The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
......@@ -12,7 +12,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" Stable Speech model configuration"""
""" Parler-TTS model configuration"""
from transformers import AutoConfig, logging
from transformers.configuration_utils import PretrainedConfig
......@@ -21,26 +21,26 @@ from transformers.configuration_utils import PretrainedConfig
logger = logging.get_logger(__name__)
MUSICGEN_PRETRAINED_CONFIG_ARCHIVE_MAP = {
"facebook/stable_speech-small": "https://huggingface.co/facebook/stable_speech-small/resolve/main/config.json",
# See all StableSpeech models at https://huggingface.co/models?filter=stable_speech
"facebook/parler_tts-small": "https://huggingface.co/facebook/parler_tts-small/resolve/main/config.json",
# See all ParlerTTS models at https://huggingface.co/models?filter=parler_tts
}
class StableSpeechDecoderConfig(PretrainedConfig):
class ParlerTTSDecoderConfig(PretrainedConfig):
r"""
This is the configuration class to store the configuration of an [`StableSpeechDecoder`]. It is used to instantiate a
Stable Speech decoder according to the specified arguments, defining the model architecture. Instantiating a
configuration with the defaults will yield a similar configuration to that of the Stable Speech
[facebook/stable_speech-small](https://huggingface.co/facebook/stable_speech-small) architecture.
This is the configuration class to store the configuration of an [`ParlerTTSDecoder`]. It is used to instantiate a
Parler-TTS decoder according to the specified arguments, defining the model architecture. Instantiating a
configuration with the defaults will yield a similar configuration to that of the Parler-TTS
[facebook/parler_tts-small](https://huggingface.co/facebook/parler_tts-small) architecture.
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.
Args:
vocab_size (`int`, *optional*, defaults to 2048):
Vocabulary size of the StableSpeechDecoder model. Defines the number of different tokens that can be
represented by the `inputs_ids` passed when calling [`StableSpeechDecoder`].
vocab_size (`int`, *optional*, defaults to 2049):
Vocabulary size of the ParlerTTSDecoder model. Defines the number of different tokens that can be
represented by the `inputs_ids` passed when calling [`ParlerTTSDecoder`].
hidden_size (`int`, *optional*, defaults to 1024):
Dimensionality of the layers and the pooler layer.
num_hidden_layers (`int`, *optional*, defaults to 24):
......@@ -76,12 +76,12 @@ class StableSpeechDecoderConfig(PretrainedConfig):
Whether input and output word embeddings should be tied.
"""
model_type = "stable_speech_decoder"
model_type = "parler_tts_decoder"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
vocab_size=2048,
vocab_size=2049, # vocab size = 2048 (encodec vocab size) + 1 (eos)
max_position_embeddings=2048,
num_hidden_layers=24,
ffn_dim=4096,
......@@ -97,8 +97,8 @@ class StableSpeechDecoderConfig(PretrainedConfig):
scale_embedding=False,
num_codebooks=4,
pad_token_id=2048,
bos_token_id=2048,
eos_token_id=None,
bos_token_id=2049,
eos_token_id=2048,
tie_word_embeddings=False,
**kwargs,
):
......@@ -127,16 +127,19 @@ class StableSpeechDecoderConfig(PretrainedConfig):
)
class StableSpeechConfig(PretrainedConfig):
class ParlerTTSConfig(PretrainedConfig):
r"""
This is the configuration class to store the configuration of a [`StableSpeechModel`]. It is used to instantiate a
Stable Speech model according to the specified arguments, defining the text encoder, audio encoder and Stable Speech decoder
This is the configuration class to store the configuration of a [`ParlerTTSModel`]. It is used to instantiate a
Parler-TTS model according to the specified arguments, defining the text encoder, audio encoder and Parler-TTS decoder
configs.
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.
Args:
vocab_size (`int`, *optional*, defaults to 1024):
Vocabulary size of the prompt token ids. Defines the number of different tokens that can be
represented by the `prompt_inputs_ids`.
kwargs (*optional*):
Dictionary of keyword arguments. Notably:
......@@ -151,24 +154,24 @@ class StableSpeechConfig(PretrainedConfig):
```python
>>> from transformers import (
... StableSpeechConfig,
... StableSpeechDecoderConfig,
... ParlerTTSConfig,
... ParlerTTSDecoderConfig,
... T5Config,
... EncodecConfig,
... StableSpeechForConditionalGeneration,
... ParlerTTSForConditionalGeneration,
... )
>>> # Initializing text encoder, audio encoder, and decoder model configurations
>>> text_encoder_config = T5Config()
>>> audio_encoder_config = EncodecConfig()
>>> decoder_config = StableSpeechDecoderConfig()
>>> decoder_config = ParlerTTSDecoderConfig()
>>> configuration = StableSpeechConfig.from_sub_models_config(
>>> configuration = ParlerTTSConfig.from_sub_models_config(
... text_encoder_config, audio_encoder_config, decoder_config
... )
>>> # Initializing a StableSpeechForConditionalGeneration (with random weights) from the facebook/stable_speech-small style configuration
>>> model = StableSpeechForConditionalGeneration(configuration)
>>> # Initializing a ParlerTTSForConditionalGeneration (with random weights) from the facebook/parler_tts-small style configuration
>>> model = ParlerTTSForConditionalGeneration(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config
......@@ -177,17 +180,17 @@ class StableSpeechConfig(PretrainedConfig):
>>> config_decoder = model.config.decoder
>>> # Saving the model, including its configuration
>>> model.save_pretrained("stable_speech-model")
>>> model.save_pretrained("parler_tts-model")
>>> # loading model and config from pretrained folder
>>> stable_speech_config = StableSpeechConfig.from_pretrained("stable_speech-model")
>>> model = StableSpeechForConditionalGeneration.from_pretrained("stable_speech-model", config=stable_speech_config)
>>> parler_tts_config = ParlerTTSConfig.from_pretrained("parler_tts-model")
>>> model = ParlerTTSForConditionalGeneration.from_pretrained("parler_tts-model", config=parler_tts_config)
```"""
model_type = "stable_speech"
model_type = "parler_tts"
is_composition = True
def __init__(self, **kwargs):
def __init__(self, vocab_size=1024, **kwargs):
super().__init__(**kwargs)
if "text_encoder" not in kwargs or "audio_encoder" not in kwargs or "decoder" not in kwargs:
raise ValueError("Config has to be initialized with text_encoder, audio_encoder and decoder config")
......@@ -200,9 +203,10 @@ class StableSpeechConfig(PretrainedConfig):
decoder_config = kwargs.pop("decoder")
self.vocab_size = vocab_size
self.text_encoder = AutoConfig.for_model(text_encoder_model_type, **text_encoder_config)
self.audio_encoder = AutoConfig.for_model(audio_encoder_model_type, **audio_encoder_config)
self.decoder = StableSpeechDecoderConfig(**decoder_config)
self.decoder = ParlerTTSDecoderConfig(**decoder_config)
self.is_encoder_decoder = True
@classmethod
......@@ -210,15 +214,15 @@ class StableSpeechConfig(PretrainedConfig):
cls,
text_encoder_config: PretrainedConfig,
audio_encoder_config: PretrainedConfig,
decoder_config: StableSpeechDecoderConfig,
decoder_config: ParlerTTSDecoderConfig,
**kwargs,
):
r"""
Instantiate a [`StableSpeechConfig`] (or a derived class) from text encoder, audio encoder and decoder
Instantiate a [`ParlerTTSConfig`] (or a derived class) from text encoder, audio encoder and decoder
configurations.
Returns:
[`StableSpeechConfig`]: An instance of a configuration object
[`ParlerTTSConfig`]: An instance of a configuration object
"""
return cls(
......
parler_tts/dac_wrapper/__init__.py 0 → 100644
View file @ 5eae102f
from .configuration_dac import DACConfig
from .modeling_dac import DACModel
parler_tts/dac_wrapper/configuration_dac.py 0 → 100644
View file @ 5eae102f
from transformers import PretrainedConfig
from typing import List
class DACConfig(PretrainedConfig):
model_type = "dac"
def __init__(
self,
num_codebooks: int = 9,
model_bitrate: int = 8, # kbps
codebook_size: int = 1024,
latent_dim: int = 1024,
frame_rate: int = 86,
sampling_rate: int = 44100,
**kwargs,
):
self.codebook_size = codebook_size
self.model_bitrate = model_bitrate
self.latent_dim = latent_dim
self.num_codebooks = num_codebooks
self.frame_rate = frame_rate
self.sampling_rate = sampling_rate
super().__init__(**kwargs)
parler_tts/dac_wrapper/modeling_dac.py 0 → 100644
View file @ 5eae102f
import torch
from transformers import PreTrainedModel
from transformers.models.encodec.modeling_encodec import EncodecEncoderOutput, EncodecDecoderOutput
from .configuration_dac import DACConfig
from dac.model import DAC
# model doesn't support batching yet
class DACModel(PreTrainedModel):
config_class = DACConfig
def __init__(self, config):
super().__init__(config)
self.model = DAC(
n_codebooks=config.num_codebooks,
latent_dim=config.latent_dim,
codebook_size=config.codebook_size,
)
def encode(
self, input_values, padding_mask=None, bandwidth=None, return_dict=None, n_quantizers=None, sample_rate=None
):
"""
Encodes the input audio waveform into discrete codes.
Args:
input_values (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`):
Float values of the input audio waveform.
padding_mask (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`):
Padding mask used to pad the `input_values`.
bandwidth (`float`, *optional*):
Not used, kept to have the same inferface as HF encodec.
n_quantizers (`int`, *optional*) :
Number of quantizers to use, by default None
If None, all quantizers are used.
sample_rate (`int`, *optional*) :
Signal sampling_rate
Returns:
A list of frames containing the discrete encoded codes for the input audio waveform, along with rescaling
factors for each chunk when `normalize` is True. Each frames is a tuple `(codebook, scale)`, with
`codebook` of shape `[batch_size, num_codebooks, frames]`.
Scale is not used here.
"""
_, channels, input_length = input_values.shape
if channels < 1 or channels > 2:
raise ValueError(f"Number of audio channels must be 1 or 2, but got {channels}")
audio_data = self.model.preprocess(input_values, sample_rate)
return_dict = return_dict if return_dict is not None else self.config.return_dict
# TODO: for now, no chunk length
chunk_length = None # self.config.chunk_length
if chunk_length is None:
chunk_length = input_length
stride = input_length
else:
stride = self.config.chunk_stride
if padding_mask is None:
padding_mask = torch.ones_like(input_values).bool()
encoded_frames = []
scales = []
step = chunk_length - stride
if (input_length % stride) - step != 0:
raise ValueError(
"The input length is not properly padded for batched chunked decoding. Make sure to pad the input correctly."
)
for offset in range(0, input_length - step, stride):
mask = padding_mask[..., offset : offset + chunk_length].bool()
frame = audio_data[:, :, offset : offset + chunk_length]
scale = None
_, encoded_frame, _, _, _ = self.model.encode(frame, n_quantizers=n_quantizers)
encoded_frames.append(encoded_frame)
scales.append(scale)
encoded_frames = torch.stack(encoded_frames)
if not return_dict:
return (encoded_frames, scales)
return EncodecEncoderOutput(encoded_frames, scales)
def decode(
self,
audio_codes,
audio_scales,
padding_mask=None,
return_dict=None,
):
"""
Decodes the given frames into an output audio waveform.
Note that the output might be a bit bigger than the input. In that case, any extra steps at the end can be
trimmed.
Args:
audio_codes (`torch.FloatTensor` of shape `(batch_size, nb_chunks, chunk_length)`, *optional*):
Discret code embeddings computed using `model.encode`.
audio_scales (`torch.Tensor` of shape `(batch_size, nb_chunks)`, *optional*):
Not used, kept to have the same inferface as HF encodec.
padding_mask (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`):
Padding mask used to pad the `input_values`.
Not used yet, kept to have the same inferface as HF encodec.
return_dict (`bool`, *optional*):
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""
return_dict = return_dict or self.config.return_dict
# TODO: for now, no chunk length
if len(audio_codes) != 1:
raise ValueError(f"Expected one frame, got {len(audio_codes)}")
audio_values = self.model.quantizer.from_codes(audio_codes.squeeze(0))[0]
audio_values = self.model.decode(audio_values)
if not return_dict:
return (audio_values,)
return EncodecDecoderOutput(audio_values)
def forward(self, tensor):
raise ValueError(f"`DACModel.forward` not implemented yet")
prompt_creation_scripts/run_prompt_creation_dummy.sh deleted 100644 → 0
View file @ 85b8cac7
#!/usr/bin/env bash
python run_prompt_creation.py \
--dataset_name "ylacombe/libritts_r_tags_and_text" \
--dataset_config_name "clean" \
--dataset_split_name "dev.clean" \
--model_name_or_path "mistralai/Mistral-7B-Instruct-v0.2" \
--per_device_eval_batch_size 2 \
--attn_implementation "sdpa" \
--dataloader_num_workers 0 \
--output_dir "./" \
--load_in_4bit
run_audio_classification.py deleted 100644 → 0
View file @ 85b8cac7
This diff is collapsed.
run_dataset_concatenation.py deleted 100644 → 0
View file @ 85b8cac7
import os
import sys
from dataclasses import dataclass, field
from pathlib import Path
import numpy as np
from datasets import Audio, concatenate_datasets, load_dataset
from huggingface_hub import get_full_repo_name
from transformers import HfArgumentParser, WhisperTokenizerFast
@dataclass
class DataTrainingArguments:
"""
Arguments pertaining to what data we are going to input our model for training and eval.
"""
dataset_name: str = field(
default=None,
metadata={"help": "The name of the dataset to use (via the datasets library)."},
)
dataset_config_name: str = field(
default=None,
metadata={"help": "The configuration name of the dataset to use (via the datasets library)."},
)
dataset_split_name: str = field(
default=None,
metadata={
"help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
},
)
label_column_name: str = field(
default="labels",
metadata={"help": "The name of the dataset column containing the labels in the dataset. Defaults to 'label'"},
)
text_column_name: str = field(
default="text",
metadata={
"help": "The name of the dataset column containing the text transcriptions in the dataset. Defaults to 'text'"
},
)
speaker_column_name: str = field(
default="speaker_id",
metadata={
"help": "The name of the dataset column containing the speaker ids in the dataset. Defaults to 'speaker_id'"
},
)
dataset_cache_dir: str = field(
default=None,
metadata={"help": "Path to cache directory for saving and loading datasets"},
)
preprocessing_num_workers: int = field(
default=None,
metadata={"help": "The number of processes to use for the preprocessing."},
)
batch_size: int = field(
default=500,
metadata={"help": "Number of examples per batch provided to the preprocessing function."},
)
download_only: bool = field(
default=False,
metadata={"help": "Whether to only do data download and skip pre-processing."},
)
audio_column_name: str = field(
default="audio",
metadata={"help": "The name of the dataset column containing the audio data. Defaults to 'audio'"},
)
max_duration_in_seconds: float = field(
default=20.0,
metadata={"help": "Filter audio files that are longer than `max_duration_in_seconds` seconds"},
)
sampling_rate: int = field(
default=16_000,
metadata={
"help": "Sampling rate at which to resample the audio data. Should be set to the same sampling rate as the target model."
},
)
max_samples: int = field(
default=None,
metadata={
"help": "For debugging purposes, truncate the number of examples in the dataset to this value if set."
},
)
output_dir: str = field(
default=None,
metadata={
"help": "Where to save the processed dataset to disk. If unspecified, uses a 'pretty' version of the "
"original dataset name. E.g. 'facebook/voxpopuli' will be saved under 'voxpopuli'."
},
)
push_to_hub: bool = field(
default=False,
metadata={"help": "Whether or not to push the processed dataset to the Hub."},
)
seed: int = field(
default=0,
metadata={"help": "RNG seed for reproducibility. Used during the final shuffling of the combined dataset."},
)
def convert_dataset_str_to_list(
dataset_names,
dataset_config_names,
splits=None,
label_column_names=None,
text_column_names=None,
speaker_column_names=None,
dataset_samples=None,
default_split="train",
):
if isinstance(dataset_names, str):
dataset_names = dataset_names.split("+")
dataset_config_names = dataset_config_names.split("+")
splits = splits.split("+") if splits is not None else None
label_column_names = label_column_names.split("+") if label_column_names is not None else None
text_column_names = text_column_names.split("+") if text_column_names is not None else None
speaker_column_names = speaker_column_names.split("+") if speaker_column_names is not None else None
dataset_samples = dataset_samples.split("+") if dataset_samples is not None else None
# basic checks to ensure we've got the right number of datasets/configs/splits/columns/probs
if len(dataset_names) != len(dataset_config_names):
raise ValueError(
f"Ensure one config is passed for each dataset, got {len(dataset_names)} datasets and"
f" {len(dataset_config_names)} configs."
)
if splits is not None and len(splits) != len(dataset_names):
raise ValueError(
f"Ensure one split is passed for each dataset, got {len(dataset_names)} datasets and {len(splits)} splits."
)
if label_column_names is not None and len(label_column_names) != len(dataset_names):
raise ValueError(
f"Ensure one label column name is passed for each dataset, got {len(dataset_names)} datasets and"
f" {len(label_column_names)} label column names."
)
if text_column_names is not None and len(text_column_names) != len(dataset_names):
raise ValueError(
f"Ensure one text column name is passed for each dataset, got {len(dataset_names)} datasets and"
f" {len(text_column_names)} text column names."
)
if speaker_column_names is not None and len(speaker_column_names) != len(dataset_names):
raise ValueError(
f"Ensure one text column name is passed for each dataset, got {len(dataset_names)} datasets and"
f" {len(speaker_column_names)} speaker column names."
)
if dataset_samples is not None:
if len(dataset_samples) != len(dataset_names):
raise ValueError(
f"Ensure one sample is passed for each dataset, got {len(dataset_names)} datasets and "
f"{len(dataset_samples)} samples."
)
dataset_samples = [float(ds_sample) for ds_sample in dataset_samples]
else:
dataset_samples = [None] * len(dataset_names)
label_column_names = (
label_column_names if label_column_names is not None else ["labels" for _ in range(len(dataset_names))]
)
text_column_names = (
text_column_names if text_column_names is not None else ["text" for _ in range(len(dataset_names))]
)
speaker_column_names = (
speaker_column_names if speaker_column_names is not None else ["speaker_id" for _ in range(len(dataset_names))]
)
splits = splits if splits is not None else [default_split for _ in range(len(dataset_names))]
dataset_names_dict = []
for i, ds_name in enumerate(dataset_names):
dataset_names_dict.append(
{
"name": ds_name,
"config": dataset_config_names[i],
"split": splits[i],
"label_column_name": label_column_names[i],
"text_column_name": text_column_names[i],
"speaker_column_name": speaker_column_names[i],
"samples": dataset_samples[i],
}
)
return dataset_names_dict
def main():
# 1. Parse input arguments
parser = HfArgumentParser(DataTrainingArguments)
if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
# If we pass only one argument to the script and it's the path to a json file,
# let's parse it to get our arguments.
data_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))[0]
else:
data_args = parser.parse_args_into_dataclasses()[0]
dataset_names_dict = convert_dataset_str_to_list(
data_args.dataset_name,
data_args.dataset_config_name,
splits=data_args.dataset_split_name,
label_column_names=data_args.label_column_name,
text_column_names=data_args.text_column_name,
speaker_column_names=data_args.speaker_column_name,
)
# load whisper tokenizer for normalisation
sampling_rate = data_args.sampling_rate
tokenizer = WhisperTokenizerFast.from_pretrained("openai/whisper-tiny.en")
max_input_length = int(data_args.max_duration_in_seconds * sampling_rate)
batch_size = data_args.batch_size
preprocessing_num_workers = data_args.preprocessing_num_workers
all_vectorized_datasets = []
for dataset_dict in dataset_names_dict:
print(10 * "=", dataset_dict["name"], 10 * "=")
raw_datasets = load_dataset(
dataset_dict["name"],
dataset_dict["config"],
split=dataset_dict["split"],
cache_dir=data_args.dataset_cache_dir,
num_proc=data_args.preprocessing_num_workers,
)
if data_args.download_only:
continue
features = raw_datasets.column_names
if dataset_dict["label_column_name"] not in features:
raise ValueError(
f"--label_column_name {dataset_dict['label_column_name']} not found in dataset '{dataset_dict['name']}'. "
"Make sure to set `--label_column_name` to the correct text column - one of "
f"{', '.join(features)}."
)
elif dataset_dict["label_column_name"] != "labels":
raw_datasets = raw_datasets.rename_column(dataset_dict["label_column_name"], "labels")
if dataset_dict["text_column_name"] not in features:
raise ValueError(
f"--text_column_name {dataset_dict['text_column_name']} not found in dataset '{dataset_dict['name']}'. "
"Make sure to set `--text_column_name` to the correct text column - one of "
f"{', '.join(features)}."
)
elif dataset_dict["text_column_name"] != "text":
raw_datasets = raw_datasets.rename_column(dataset_dict["text_column_name"], "text")
if dataset_dict["speaker_column_name"] not in features:
raise ValueError(
f"--speaker_column_name {dataset_dict['speaker_column_name']} not found in dataset '{dataset_dict['name']}'. "
"Make sure to set `--speaker_column_name` to the correct speaker id column - one of "
f"{', '.join(features)}."
)
elif dataset_dict["speaker_column_name"] != "speaker_id":
raw_datasets = raw_datasets.rename_column(dataset_dict["speaker_column_name"], "speaker_id")
raw_datasets = raw_datasets.remove_columns(
set(raw_datasets.features.keys()) - {"audio", "labels", "text", "speaker_id"}
)
if data_args.max_samples is not None:
raw_datasets = raw_datasets.select(range(data_args.max_samples))
raw_datasets = raw_datasets.cast_column(data_args.audio_column_name, Audio(sampling_rate=sampling_rate))
raw_datasets = raw_datasets.sort("speaker_id")
def filter_transcriptions(text):
normalized_text = tokenizer.normalize(text).strip()
return bool(normalized_text) and text.lower() != "ignore_time_segment_in_scoring"
raw_datasets = raw_datasets.filter(
filter_transcriptions, input_columns=["text"], desc="Filtering non-speech transcriptions"
)
def prepare_dataset(batch):
audio = [sample["array"] for sample in batch["audio"]]
input_lengths = [len(sample) for sample in audio]
concatenated_audio = []
concatenated_text = []
concatenated_speaker = []
concatenated_labels = []
audio_sample = audio[0]
text_sample = batch["text"][0]
label_sample = batch["labels"][0]
for idx in range(1, len(audio)):
prev_speaker = batch["speaker_id"][idx - 1]
speaker = batch["speaker_id"][idx]
if len(audio_sample) + input_lengths[idx] < max_input_length:
if speaker == prev_speaker:
# we have no information about whether the segments follow on sequentially
# so we just ensure the same speaker as we concatenate across files
audio_sample = np.append(audio_sample, audio[idx])
# extra spaces in the text transcription don't matter, since we only use it for the WER computation
text_sample += " " + batch["text"][idx]
else:
# segments do not follow sequentially, save the audio and start looping again
concatenated_audio.append(audio_sample)
concatenated_text.append(text_sample)
concatenated_labels.append(label_sample)
concatenated_speaker.append(speaker)
audio_sample = audio[idx]
text_sample = batch["text"][idx]
label_sample = batch["labels"][idx]
else:
# concatenated audio exceeds max length, save the audio and start looping again
concatenated_audio.append(audio_sample)
concatenated_text.append(text_sample)
concatenated_labels.append(label_sample)
concatenated_speaker.append(speaker)
audio_sample = audio[idx]
text_sample = batch["text"][idx]
label_sample = batch["labels"][idx]
batch["audio"] = [{"array": array, "sampling_rate": sampling_rate} for array in concatenated_audio]
batch["text"] = concatenated_text
batch["labels"] = concatenated_labels
batch["speaker_id"] = concatenated_speaker
return batch
raw_datasets = raw_datasets.map(
prepare_dataset,
batched=True,
batch_size=batch_size,
num_proc=preprocessing_num_workers,
desc="Concatenating dataset...",
)
pretty_name = dataset_dict["name"].split("/")[-1]
def postprocess_ids(speaker_id, idx):
formatted_idx = f"{pretty_name}-{speaker_id}-{idx}"
return {"id": formatted_idx}
raw_datasets = raw_datasets.map(
postprocess_ids,
input_columns=["speaker_id"],
with_indices=True,
desc="Setting sample idxs...",
num_proc=preprocessing_num_workers,
)
print(f"Final length {pretty_name}: ", len(raw_datasets))
# Re-format transcriptions and condition on prev as numpy arrays
raw_datasets = raw_datasets.with_format("np")
all_vectorized_datasets.append(raw_datasets)
all_vectorized_datasets = concatenate_datasets(all_vectorized_datasets)
dataset_features = all_vectorized_datasets.features.copy()
dataset_features["audio"] = Audio(sampling_rate=sampling_rate)
all_vectorized_datasets = all_vectorized_datasets.cast(
dataset_features, batch_size=batch_size, writer_batch_size=batch_size, num_proc=preprocessing_num_workers
)
all_vectorized_datasets = all_vectorized_datasets.shuffle(seed=data_args.seed)
all_vectorized_datasets.save_to_disk(data_args.output_dir)
repo_name = get_full_repo_name(Path(data_args.output_dir).absolute().name)
if data_args.push_to_hub:
all_vectorized_datasets.push_to_hub(repo_name, config_name="train", max_shard_size="1GB")
if __name__ == "__main__":
main()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment