Upload 518 files

capspeech/ar/README.md

@@ -0,0 +1,44 @@
+# CapSpeech-AR
+
+## Pretrain
+
+```bash
+bash pretrain.sh
+```
+Make sure to change paths and keys in `pretrain.sh` to yours.
+
+## Finetune on CapTTS
+
+```bash
+bash finetune_captts.sh
+```
+Make sure to change paths and keys in `finetune_captts.sh` to yours.
+
+## Finetune on EmoCapTTS
+
+```bash
+bash finetune_emocaptts.sh
+```
+Make sure to change paths and keys in `finetune_emocaptts.sh` to yours.
+
+## Finetune on AccCapTTS
+
+```bash
+bash finetune_acccaptts.sh
+```
+Make sure to change paths and keys in `finetune_acccaptts.sh` to yours.
+
+## Finetune on CapTTS-SE
+
+```bash
+bash finetune_capttsse.sh
+```
+Make sure to change paths and keys in `finetune_capttsse.sh` to yours.
+
+
+## Finetune on AgentTTS
+
+```bash
+bash finetune_agenttts.sh
+```
+Make sure to change paths and keys in `finetune_agenttts.sh` to yours.

capspeech/ar/events.txt

@@ -0,0 +1,395 @@
+people whispering
+Microwave oven
+extending ladders
+mosquito buzzing
+dog whimpering
+coyote howling
+hair dryer drying
+Writing
+rapping
+machine gun shooting
+dog bow-wow
+dog howling
+barn swallow calling
+baby babbling
+Fireworks
+church bell ringing
+car horn
+cat caterwauling
+subway, metro, underground
+waterfall burbling
+lions roaring
+toilet flushing
+skateboarding
+wind
+ripping paper
+vacuum cleaner cleaning floors
+mouse squeaking
+keyboard typing
+playing timpani
+playing harp
+sheep bleating
+eletric blender running
+people slapping
+playing ukulele
+frog
+car engine knocking
+cat purring
+chainsaw
+Violin or fiddle
+people hiccup
+playing acoustic guitar
+donkey, ass braying
+playing french horn
+playing squash
+gibbon howling
+playing harmonica
+playing shofar
+hedge trimmer running
+playing washboard
+running electric fan
+splashing water
+playing bassoon
+people slurping
+playing accordion
+playing oboe
+popping popcorn
+glass breaking
+alarm clock ringing
+mouse click
+Laughter
+magpie calling
+playing snare drum
+people finger snapping
+ferret dooking
+tornado roaring
+Hi-hat
+lawn mowing
+church bells
+cat growling
+cheetah chirrup
+heart sounds, heartbeat
+firing muskets
+vehicle horn, car horn, honking
+turkey gobbling
+ice cream truck, ice cream van
+underwater bubbling
+footsteps on snow
+water drops
+people sobbing
+basketball bounce
+Applause
+playing sitar
+playing gong
+train
+coughing
+people screaming
+Gunshot or gunfire
+chinchilla barking
+cat hissing
+horse clip-clop
+engine
+people battle cry
+typing on computer keyboard
+playing clarinet
+driving motorcycle
+male singing
+singing bowl
+skiing
+driving buses
+alligators, crocodiles hissing
+people eating apple
+door slamming
+Flute
+raining
+Electric piano
+sliding door
+washing machine
+opening or closing car electric windows
+baby crying
+people babbling
+snake hissing
+brushing teeth
+playing tambourine
+Acoustic guitar
+clock tick
+playing castanets
+thunder
+playing didgeridoo
+playing synthesizer
+mouse clicking
+lathe spinning
+spraying water
+hen
+stream burbling
+door wood creaks
+sailing
+dog
+car engine idling
+bowling impact
+driving snowmobile
+toilet flush
+bird squawking
+playing timbales
+playing drum kit
+owl hooting
+striking pool
+Oboe
+duck quacking
+people belly laughing
+lighting firecrackers
+roller coaster running
+blowtorch igniting
+wood thrush calling
+Glockenspiel
+frog croaking
+playing harpsichord
+train horning
+plastic bottle crushing
+playing tabla
+fire crackling
+dog barking
+thunderstorm
+playing banjo
+swimming
+volcano explosion
+playing table tennis
+sea lion barking
+rowboat, canoe, kayak rowing
+Meow
+pouring water
+playing tympani
+rooster
+siren
+parrot talking
+Finger snapping
+playing steel guitar, slide guitar
+Trumpet
+tractor digging
+people coughing
+cat meowing
+Snare drum
+playing erhu
+crow cawing
+playing djembe
+whale calling
+mynah bird singing
+playing tennis
+chopping food
+golf driving
+tapping guitar
+playing cello
+dog growling
+elephant trumpeting
+sea waves
+police radio chatter
+lions growling
+playing lacrosse
+children shouting
+missile launch
+baby laughter
+air conditioning noise
+playing saxophone
+typing on typewriter
+printer printing
+race car, auto racing
+Bus
+pigeon, dove cooing
+playing violin, fiddle
+Double bass
+striking bowling
+fireworks banging
+Harmonica
+playing glockenspiel
+reversing beeps
+playing piano
+breathing
+people marching
+electric shaver, electric razor shaving
+chimpanzee pant-hooting
+cricket chirping
+bird chirping, tweeting
+using sewing machines
+crickets
+cow lowing
+playing cymbal
+vacuum cleaner
+playing zither
+train whistling
+goat bleating
+eating with cutlery
+black capped chickadee calling
+ambulance siren
+playing hockey
+dog baying
+Burping or eructation
+cupboard opening or closing
+air horn
+crying baby
+people eating crisps
+sloshing water
+goose honking
+orchestra
+people giggling
+warbler chirping
+child singing
+dinosaurs bellowing
+motorboat, speedboat acceleration
+airplane
+chicken clucking
+woodpecker pecking tree
+Drawer open or close
+people eating
+drinking sipping
+singing choir
+playing bass guitar
+playing bass drum
+car passing by
+playing tuning fork
+Squeak
+pig oinking
+Computer keyboard
+yodelling
+playing trombone
+clapping
+people sneezing
+pheasant crowing
+writing on blackboard with chalk
+Tambourine
+opening or closing car doors
+sharpen knife
+people whistling
+fireworks
+playing bagpipes
+chainsawing trees
+squishing water
+people farting
+playing electric guitar
+people booing
+female singing
+ocean burbling
+cattle mooing
+footsteps
+Knock
+wind rustling leaves
+cattle, bovinae cowbell
+Clarinet
+police car (siren)
+Fart
+cat
+sheep
+chopping wood
+tap dancing
+playing mandolin
+wind chime
+can opening
+playing hammond organ
+zebra braying
+scuba diving
+chirping birds
+playing steelpan
+playing theremin
+Keys jangling
+beat boxing
+firing cannon
+bouncing on trampoline
+door wood knock
+bathroom ventilation fan running
+snake rattling
+bull bellowing
+electric grinder grinding
+penguins braying
+otter growling
+civil defense siren
+wind noise
+people humming
+clock alarm
+disc scratching
+fire truck siren
+telephone bell ringing
+people sniggering
+playing bongo
+cap gun shooting
+opening or closing drawers
+cow
+hammering nails
+ice cracking
+foghorn
+rain
+playing badminton
+eagle screaming
+playing double bass
+insects
+people running
+planing timber
+cutting hair with electric trimmers
+Cello
+people clapping
+smoke detector beeping
+mouse pattering
+bee, wasp, etc. buzzing
+canary calling
+people burping
+Shatter
+baltimore oriole calling
+cuckoo bird calling
+snoring
+strike lighter
+people cheering
+playing bugle
+playing congas
+playing vibraphone
+hail
+rope skipping
+playing trumpet
+pig
+hand saw
+people gargling
+Scissors
+metronome
+chipmunk chirping
+playing flute
+fox barking
+crackling fire
+playing volleyball
+skidding
+Bass drum
+crow
+elk bugling
+Telephone
+Bark
+chicken crowing
+people nose blowing
+car engine starting
+pumping water
+Saxophone
+fly, housefly buzzing
+Cough
+people eating noodle
+francolin calling
+arc welding
+horse neighing
+Tearing
+helicopter
+playing electronic organ
+Cowbell
+railroad car, train wagon
+cell phone buzzing
+playing cornet
+sneezing
+engine accelerating, revving, vroom
+bird wings flapping
+playing marimba, xylophone
+playing guiro
+people crowd
+train wheels squealing
+slot machine
+laughing
+lip smacking
+forging swords
+Chime
+playing darts
+people shuffling
+Gong
+airplane flyby
+None

capspeech/ar/finetune_acccaptts.sh

@@ -0,0 +1,64 @@
+# Please log in to huggingface first
+
+LIBRITTSR_WAV_DIR='' # downloaded libritts-r wav dir
+OTHER_WAV_DIR='' # downloaded other wav dirs
+OUTPUT_DIR="./output_finetuning_acccaptts/" # output dir, to save checkpoints
+TEMPORY_SAVE_TO_DISK="./audio_code_finetuning_acccaptts/" # dac codec saved dir
+SAVE_TO_DISK="./dataset_finetuning_acccaptts/" # huggingface metadata saved dir
+WANDB_KEY='' # your wandb key for logging
+
+PRETRAINED_MODEL_PATH="" # your pretrained model path
+
+export CUDA_LAUNCH_BLOCKING=1
+export TORCH_USE_CUDA_DSA=1
+
+accelerate launch ./training/finetune_captts.py \
+    --model_name_or_path ${PRETRAINED_MODEL_PATH} \
+    --feature_extractor_name "parler-tts/dac_44khZ_8kbps" \
+    --description_tokenizer_name ${PRETRAINED_MODEL_PATH} \
+    --prompt_tokenizer_name ${PRETRAINED_MODEL_PATH} \
+    --report_to "wandb" \
+    --wandb_key ${WANDB_KEY} \
+    --overwrite_output_dir true \
+    --train_dataset_name "OpenSound/CapSpeech" \
+    --train_split_name "train_SFT_AccCapTTS" \
+    --eval_dataset_name "OpenSound/CapSpeech" \
+    --eval_split_name "validation_SFT_AccCapTTS" \
+    --librittsr_dir ${LIBRITTSR_WAV_DIR} \
+    --other_dir ${OTHER_WAV_DIR} \
+    --max_eval_samples 96 \
+    --per_device_eval_batch_size 32 \
+    --target_audio_column_name "audio_path" \
+    --description_column_name "caption" \
+    --source_column_name "source" \
+    --prompt_column_name "text" \
+    --max_duration_in_seconds 20 \
+    --min_duration_in_seconds 3 \
+    --max_text_length 600 \
+    --preprocessing_num_workers 32 \
+    --do_train true \
+    --num_train_epochs 5 \
+    --gradient_accumulation_steps 6 \
+    --gradient_checkpointing false \
+    --per_device_train_batch_size 4 \
+    --learning_rate 0.0001 \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.99 \
+    --weight_decay 0.01 \
+    --lr_scheduler_type "constant_with_warmup" \
+    --warmup_steps 1000 \
+    --logging_steps 200 \
+    --freeze_text_encoder true \
+    --per_device_eval_batch_size 4 \
+    --audio_encoder_per_device_batch_size 24 \
+    --dtype "float16" \
+    --seed 456 \
+    --output_dir ${OUTPUT_DIR} \
+    --temporary_save_to_disk ${TEMPORY_SAVE_TO_DISK} \
+    --save_to_disk ${SAVE_TO_DISK} \
+    --dataloader_num_workers 32 \
+    --do_eval \
+    --evaluation_strategy steps \
+    --eval_steps 500 \
+    --save_steps 500 \
+    --group_by_length true

capspeech/ar/finetune_agenttts.sh

@@ -0,0 +1,61 @@
+# Please log in to huggingface first
+
+OTHER_WAV_DIR='' # downloaded capspeech-agentdb wav dir
+OUTPUT_DIR="./output_finetuning_agenttts/" # output dir, to save checkpoints
+TEMPORY_SAVE_TO_DISK="./audio_code_finetuning_agenttts/" # dac codec saved dir
+SAVE_TO_DISK="./dataset_finetuning_agenttts/" # huggingface metadata saved dir
+WANDB_KEY='' # your wandb key for logging
+PRETRAINED_MODEL_PATH="" # your pretrained model path
+
+export CUDA_LAUNCH_BLOCKING=1
+export TORCH_USE_CUDA_DSA=1
+
+accelerate launch ./training/finetune_captts.py \
+    --model_name_or_path "/export/fs05/hwang258/parler-tts/parler-tts" \
+    --feature_extractor_name "parler-tts/dac_44khZ_8kbps" \
+    --description_tokenizer_name ${PRETRAINED_MODEL_PATH} \
+    --prompt_tokenizer_name ${PRETRAINED_MODEL_PATH} \
+    --report_to "wandb" \
+    --wandb_key ${WANDB_KEY} \
+    --overwrite_output_dir true \
+    --train_dataset_name "OpenSound/CapSpeech" \
+    --train_split_name "train_AgentDB" \
+    --eval_dataset_name "OpenSound/CapSpeech" \
+    --eval_split_name "test_AgentDB" \
+    --other_dir ${OTHER_WAV_DIR} \
+    --max_eval_samples 96 \
+    --per_device_eval_batch_size 32 \
+    --target_audio_column_name "audio_path" \
+    --description_column_name "caption" \
+    --source_column_name "source" \
+    --prompt_column_name "text" \
+    --max_duration_in_seconds 20 \
+    --min_duration_in_seconds 3 \
+    --max_text_length 600 \
+    --preprocessing_num_workers 32 \
+    --do_train true \
+    --num_train_epochs 50 \
+    --gradient_accumulation_steps 6 \
+    --gradient_checkpointing false \
+    --per_device_train_batch_size 4 \
+    --learning_rate 0.0001 \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.99 \
+    --weight_decay 0.01 \
+    --lr_scheduler_type "constant_with_warmup" \
+    --warmup_steps 500 \
+    --logging_steps 100 \
+    --freeze_text_encoder true \
+    --per_device_eval_batch_size 4 \
+    --audio_encoder_per_device_batch_size 24 \
+    --dtype "float16" \
+    --seed 456 \
+    --output_dir ${OUTPUT_DIR} \
+    --temporary_save_to_disk ${TEMPORY_SAVE_TO_DISK} \
+    --save_to_disk ${SAVE_TO_DISK} \
+    --dataloader_num_workers 32 \
+    --do_eval \
+    --evaluation_strategy steps \
+    --eval_steps 500 \
+    --save_steps 500 \
+    --group_by_length true

capspeech/ar/finetune_captts.sh

@@ -0,0 +1,64 @@
+# Please log in to huggingface first
+
+LIBRITTSR_WAV_DIR='' # downloaded libritts-r wav dir
+OTHER_WAV_DIR='' # downloaded other wav dirs
+OUTPUT_DIR="./output_finetuning_captts/" # output dir, to save checkpoints
+TEMPORY_SAVE_TO_DISK="./audio_code_finetuning_captts/" # dac codec saved dir
+SAVE_TO_DISK="./dataset_finetuning_captts/" # huggingface metadata saved dir
+WANDB_KEY='' # your wandb key for logging
+
+PRETRAINED_MODEL_PATH="" # your pretrained model path
+
+export CUDA_LAUNCH_BLOCKING=1
+export TORCH_USE_CUDA_DSA=1
+
+accelerate launch ./training/finetune_captts.py \
+    --model_name_or_path ${PRETRAINED_MODEL_PATH} \
+    --feature_extractor_name "parler-tts/dac_44khZ_8kbps" \
+    --description_tokenizer_name ${PRETRAINED_MODEL_PATH} \
+    --prompt_tokenizer_name ${PRETRAINED_MODEL_PATH} \
+    --report_to "wandb" \
+    --wandb_key ${WANDB_KEY} \
+    --overwrite_output_dir true \
+    --train_dataset_name "OpenSound/CapSpeech" \
+    --train_split_name "train_SFT_CapTTS" \
+    --eval_dataset_name "OpenSound/CapSpeech" \
+    --eval_split_name "validation_SFT_CapTTS" \
+    --librittsr_dir ${LIBRITTSR_WAV_DIR} \
+    --other_dir ${OTHER_WAV_DIR} \
+    --max_eval_samples 96 \
+    --per_device_eval_batch_size 32 \
+    --target_audio_column_name "audio_path" \
+    --description_column_name "caption" \
+    --source_column_name "source" \
+    --prompt_column_name "text" \
+    --max_duration_in_seconds 20 \
+    --min_duration_in_seconds 3 \
+    --max_text_length 600 \
+    --preprocessing_num_workers 32 \
+    --do_train true \
+    --num_train_epochs 5 \
+    --gradient_accumulation_steps 6 \
+    --gradient_checkpointing false \
+    --per_device_train_batch_size 4 \
+    --learning_rate 0.0001 \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.99 \
+    --weight_decay 0.01 \
+    --lr_scheduler_type "constant_with_warmup" \
+    --warmup_steps 1000 \
+    --logging_steps 200 \
+    --freeze_text_encoder true \
+    --per_device_eval_batch_size 4 \
+    --audio_encoder_per_device_batch_size 24 \
+    --dtype "float16" \
+    --seed 456 \
+    --output_dir ${OUTPUT_DIR} \
+    --temporary_save_to_disk ${TEMPORY_SAVE_TO_DISK} \
+    --save_to_disk ${SAVE_TO_DISK} \
+    --dataloader_num_workers 32 \
+    --do_eval \
+    --evaluation_strategy steps \
+    --eval_steps 2000 \
+    --save_steps 2000 \
+    --group_by_length true

capspeech/ar/finetune_capttsse.sh

@@ -0,0 +1,62 @@
+# Please log in to huggingface first
+
+LIBRITTSRMIX_WAV_DIR='' # downloaded capspeech-sedb wav dir
+OUTPUT_DIR="./output_finetuning_capttsse/" # output dir, to save checkpoints
+TEMPORY_SAVE_TO_DISK="./audio_code_finetuning_capttsse/" # dac codec saved dir
+SAVE_TO_DISK="./dataset_finetuning_capttsse/" # huggingface metadata saved dir
+WANDB_KEY='' # your wandb key for logging
+
+PRETRAINED_MODEL_PATH="" # your pretrained model path
+
+export CUDA_LAUNCH_BLOCKING=1
+export TORCH_USE_CUDA_DSA=1
+
+accelerate launch ./training/finetune_capttsse.py \
+    --model_name_or_path ${PRETRAINED_MODEL_PATH} \
+    --feature_extractor_name "parler-tts/dac_44khZ_8kbps" \
+    --description_tokenizer_name ${PRETRAINED_MODEL_PATH} \
+    --prompt_tokenizer_name ${PRETRAINED_MODEL_PATH} \
+    --report_to "wandb" \
+    --wandb_key ${WANDB_KEY} \
+    --overwrite_output_dir true \
+    --train_dataset_name "OpenSound/CapSpeech" \
+    --train_split_name "train_SEDB" \
+    --eval_dataset_name "OpenSound/CapSpeech" \
+    --eval_split_name "test_SEDB" \
+    --librittsrmix_dir ${LIBRITTSRMIX_WAV_DIR} \
+    --max_eval_samples 96 \
+    --per_device_eval_batch_size 32 \
+    --target_audio_column_name "audio_path" \
+    --description_column_name "caption" \
+    --source_column_name "source" \
+    --prompt_column_name "text" \
+    --max_duration_in_seconds 20 \
+    --min_duration_in_seconds 3 \
+    --max_text_length 600 \
+    --preprocessing_num_workers 32 \
+    --do_train true \
+    --num_train_epochs 50 \
+    --gradient_accumulation_steps 6 \
+    --gradient_checkpointing false \
+    --per_device_train_batch_size 4 \
+    --learning_rate 0.0001 \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.99 \
+    --weight_decay 0.01 \
+    --lr_scheduler_type "constant_with_warmup" \
+    --warmup_steps 50 \
+    --logging_steps 20 \
+    --freeze_text_encoder true \
+    --per_device_eval_batch_size 4 \
+    --audio_encoder_per_device_batch_size 24 \
+    --dtype "float16" \
+    --seed 456 \
+    --output_dir ${OUTPUT_DIR} \
+    --temporary_save_to_disk ${TEMPORY_SAVE_TO_DISK} \
+    --save_to_disk ${SAVE_TO_DISK} \
+    --dataloader_num_workers 32 \
+    --do_eval \
+    --evaluation_strategy steps \
+    --eval_steps 50 \
+    --save_steps 50 \
+    --group_by_length true

capspeech/ar/finetune_emocaptts.sh

@@ -0,0 +1,64 @@
+# Please log in to huggingface first
+
+LIBRITTSR_WAV_DIR='' # downloaded libritts-r wav dir
+OTHER_WAV_DIR='' # downloaded other wav dirs
+OUTPUT_DIR="./output_finetuning_emocaptts/" # output dir, to save checkpoints
+TEMPORY_SAVE_TO_DISK="./audio_code_finetuning_emocaptts/" # dac codec saved dir
+SAVE_TO_DISK="./dataset_finetuning_emocaptts/" # huggingface metadata saved dir
+WANDB_KEY='' # your wandb key for logging
+
+PRETRAINED_MODEL_PATH="" # your pretrained model path
+
+export CUDA_LAUNCH_BLOCKING=1
+export TORCH_USE_CUDA_DSA=1
+
+accelerate launch ./training/finetune_captts.py \
+    --model_name_or_path ${PRETRAINED_MODEL_PATH} \
+    --feature_extractor_name "parler-tts/dac_44khZ_8kbps" \
+    --description_tokenizer_name ${PRETRAINED_MODEL_PATH} \
+    --prompt_tokenizer_name ${PRETRAINED_MODEL_PATH} \
+    --report_to "wandb" \
+    --wandb_key ${WANDB_KEY} \
+    --overwrite_output_dir true \
+    --train_dataset_name "OpenSound/CapSpeech" \
+    --train_split_name "train_SFT_EmoCapTTS" \
+    --eval_dataset_name "OpenSound/CapSpeech" \
+    --eval_split_name "validation_SFT_EmoCapTTS" \
+    --librittsr_dir ${LIBRITTSR_WAV_DIR} \
+    --other_dir ${OTHER_WAV_DIR} \
+    --max_eval_samples 96 \
+    --per_device_eval_batch_size 32 \
+    --target_audio_column_name "audio_path" \
+    --description_column_name "caption" \
+    --source_column_name "source" \
+    --prompt_column_name "text" \
+    --max_duration_in_seconds 20 \
+    --min_duration_in_seconds 3 \
+    --max_text_length 600 \
+    --preprocessing_num_workers 32 \
+    --do_train true \
+    --num_train_epochs 5 \
+    --gradient_accumulation_steps 6 \
+    --gradient_checkpointing false \
+    --per_device_train_batch_size 4 \
+    --learning_rate 0.0001 \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.99 \
+    --weight_decay 0.01 \
+    --lr_scheduler_type "constant_with_warmup" \
+    --warmup_steps 1000 \
+    --logging_steps 200 \
+    --freeze_text_encoder true \
+    --per_device_eval_batch_size 4 \
+    --audio_encoder_per_device_batch_size 24 \
+    --dtype "float16" \
+    --seed 456 \
+    --output_dir ${OUTPUT_DIR} \
+    --temporary_save_to_disk ${TEMPORY_SAVE_TO_DISK} \
+    --save_to_disk ${SAVE_TO_DISK} \
+    --dataloader_num_workers 32 \
+    --do_eval \
+    --evaluation_strategy steps \
+    --eval_steps 400 \
+    --save_steps 400 \
+    --group_by_length true

capspeech/ar/parler_tts/__init__.py

@@ -0,0 +1,25 @@
+__version__ = "0.2.2"
+
+
+from transformers import AutoConfig, AutoModel
+
+from .configuration_parler_tts import ParlerTTSConfig, ParlerTTSDecoderConfig
+from .dac_wrapper import DACConfig, DACModel
+from .modeling_parler_tts import (
+    ParlerTTSForCausalLM,
+    ParlerTTSForConditionalGeneration,
+    apply_delay_pattern_mask,
+    build_delay_pattern_mask,
+)
+
+from .streamer import ParlerTTSStreamer
+
+from importlib.metadata import version
+from packaging.version import Version
+
+if Version(version("transformers"))<= Version("4.44.2dev"):
+    AutoConfig.register("dac", DACConfig)
+else:
+    AutoConfig.register("dac_on_the_hub", DACConfig)
+
+AutoModel.register(DACConfig, DACModel)

capspeech/ar/parler_tts/configuration_parler_tts.py

@@ -0,0 +1,291 @@
+# coding=utf-8
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+""" Parler-TTS model configuration"""
+
+from transformers import AutoConfig, logging
+from transformers.configuration_utils import PretrainedConfig
+
+from importlib.metadata import version
+from packaging.version import Version
+
+use_dac_on_the_hub = Version(version("transformers")) > Version("4.44.2dev")
+
+logger = logging.get_logger(__name__)
+
+PARLER_TTS_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "parler-tts/parler-tts-mini-v1": "https://huggingface.co/parler-tts/parler-tts-mini-v1/resolve/main/config.json",
+    # See all ParlerTTS models at https://huggingface.co/models?filter=parler_tts
+}
+
+
+class ParlerTTSDecoderConfig(PretrainedConfig):
+    r"""
+    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
+    [parler-tts/parler-tts-mini-v1](https://huggingface.co/parler-tts/parler-tts-mini-v1) 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 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):
+            Number of decoder layers.
+        num_attention_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer block.
+        num_key_value_heads (`int`, *optional*):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+            `num_attention_heads`.
+        num_cross_attention_key_value_heads (`int`, *optional*):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention in the cross-attention layers.
+            If it is not specified, will default to `num_key_value_heads`.
+        ffn_dim (`int`, *optional*, defaults to 4096):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer block.
+        activation_function (`str` or `function`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the decoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"silu"` and `"gelu_new"` are supported.
+        dropout (`float`, *optional*, defaults to 0.1):
+            The dropout probability for all fully connected layers in the embeddings, text_encoder, and pooler.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        activation_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for activations inside the fully connected layer.
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with. Typically, set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        initializer_factor (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        layerdrop (`float`, *optional*, defaults to 0.0):
+            The LayerDrop probability for the decoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
+            for more details.
+        scale_embedding (`bool`, *optional*, defaults to `False`):
+            Scale embeddings by diving by sqrt(hidden_size).
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether the model should return the last key/values attentions (not used by all models)
+        num_codebooks (`int`, *optional*, defaults to 4):
+            The number of parallel codebooks forwarded to the model.
+        tie_word_embeddings(`bool`, *optional*, defaults to `False`):
+            Whether input and output word embeddings should be tied.
+        rope_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether to use ROPE or absolute positional embeddings.
+        rope_theta (`float`, *optional*, defaults to 100000.0):
+            The base period of the RoPE embeddings.
+        cross_attention_implementation_strategy (`str`, *optional*):
+            If not specified, the cross-attention implementation will be the same as `_attn_implementation`. If `always_eager`, it will always be the eager implementation. If `always_sdpa`, it will always be the sdpa implementation.
+        use_fused_lm_heads(`bool`, *optional*, defaults to `False`):
+            Whether to fuse audio LM heads instead of applying them sequentially.
+        codebook_weights(`List[int]`, *optional*):
+            Weights applied to each codebook when computing the loss.
+    """
+
+    model_type = "parler_tts_decoder"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=2049,  # vocab size = 2048 (encodec vocab size) + 1 (eos)
+        max_position_embeddings=2048,
+        num_hidden_layers=24,
+        ffn_dim=4096,
+        num_attention_heads=16,
+        num_key_value_heads=None,
+        num_cross_attention_key_value_heads=None,
+        layerdrop=0.0,
+        use_cache=True,
+        activation_function="gelu",
+        hidden_size=1024,
+        dropout=0.1,
+        attention_dropout=0.0,
+        activation_dropout=0.0,
+        initializer_factor=0.02,
+        scale_embedding=False,
+        num_codebooks=4,
+        pad_token_id=2048,
+        bos_token_id=2049,
+        eos_token_id=2048,
+        tie_word_embeddings=False,
+        rope_embeddings=False,
+        rope_theta=10_000.0,
+        cross_attention_implementation_strategy=None,
+        use_fused_lm_heads=False,
+        codebook_weights=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.ffn_dim = ffn_dim
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        if num_cross_attention_key_value_heads is None:
+            num_cross_attention_key_value_heads = num_key_value_heads
+        self.num_cross_attention_key_value_heads = num_cross_attention_key_value_heads
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.activation_function = activation_function
+        self.initializer_factor = initializer_factor
+        self.layerdrop = layerdrop
+        self.use_cache = use_cache
+        self.scale_embedding = scale_embedding  # scale factor will be sqrt(d_model) if True
+        self.num_codebooks = num_codebooks
+        self.rope_embeddings = rope_embeddings
+        self.rope_theta = rope_theta
+        self.cross_attention_implementation_strategy = cross_attention_implementation_strategy
+        self.use_fused_lm_heads = use_fused_lm_heads
+        self.codebook_weights = codebook_weights
+
+        if codebook_weights is not None and len(codebook_weights) != num_codebooks:
+            raise ValueError(f"`codebook_weights` has length {len(codebook_weights)} when it should be of length {num_codebooks}.")
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+class ParlerTTSConfig(PretrainedConfig):
+    r"""
+    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`.
+        prompt_cross_attention (`bool`, *optional*, defaults to `False`):
+            Whether to use cross-attention conditioning for the prompt (as well as the description).
+        kwargs (*optional*):
+            Dictionary of keyword arguments. Notably:
+
+                - **text_encoder** ([`PretrainedConfig`], *optional*) -- An instance of a configuration object that
+                  defines the text encoder config.
+                - **audio_encoder** ([`PretrainedConfig`], *optional*) -- An instance of a configuration object that
+                  defines the audio encoder config.
+                - **decoder** ([`PretrainedConfig`], *optional*) -- An instance of a configuration object that defines
+                  the decoder config.
+
+    Example:
+
+    ```python
+    >>> from transformers import (
+    ...     ParlerTTSConfig,
+    ...     ParlerTTSDecoderConfig,
+    ...     T5Config,
+    ...     EncodecConfig,
+    ...     ParlerTTSForConditionalGeneration,
+    ... )
+
+    >>> # Initializing text encoder, audio encoder, and decoder model configurations
+    >>> text_encoder_config = T5Config()
+    >>> audio_encoder_config = EncodecConfig()
+    >>> decoder_config = ParlerTTSDecoderConfig()
+
+    >>> configuration = ParlerTTSConfig.from_sub_models_config(
+    ...     text_encoder_config, audio_encoder_config, decoder_config
+    ... )
+
+    >>> # Initializing a ParlerTTSForConditionalGeneration (with random weights) from the parler-tts/parler-tts-mini-v1 style configuration
+    >>> model = ParlerTTSForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    >>> config_text_encoder = model.config.text_encoder
+    >>> config_audio_encoder = model.config.audio_encoder
+    >>> config_decoder = model.config.decoder
+
+    >>> # Saving the model, including its configuration
+    >>> model.save_pretrained("parler_tts-model")
+
+    >>> # loading model and config from pretrained folder
+    >>> parler_tts_config = ParlerTTSConfig.from_pretrained("parler_tts-model")
+    >>> model = ParlerTTSForConditionalGeneration.from_pretrained("parler_tts-model", config=parler_tts_config)
+    ```"""
+
+    model_type = "parler_tts"
+    is_composition = True
+
+    def __init__(self, vocab_size=1024, prompt_cross_attention=False, **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")
+
+        text_encoder_config = kwargs.pop("text_encoder")
+        text_encoder_model_type = text_encoder_config.pop("model_type")
+
+        audio_encoder_config = kwargs.pop("audio_encoder")
+        audio_encoder_model_type = audio_encoder_config.pop("model_type")
+
+        model_version = kwargs.get("transformers_version", None)
+        if model_version is not None and Version(model_version) <= Version("4.44.2dev") and use_dac_on_the_hub and audio_encoder_model_type=="dac":
+            # here we have to manually change model type if DAC based on transformers version
+            audio_encoder_model_type = "dac_on_the_hub"
+
+        decoder_config = kwargs.pop("decoder")
+
+        self.vocab_size = vocab_size
+        self.prompt_cross_attention = prompt_cross_attention
+        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 = ParlerTTSDecoderConfig(**decoder_config)
+        self.is_encoder_decoder = True
+
+    @classmethod
+    def from_sub_models_config(
+        cls,
+        text_encoder_config: PretrainedConfig,
+        audio_encoder_config: PretrainedConfig,
+        decoder_config: ParlerTTSDecoderConfig,
+        **kwargs,
+    ):
+        r"""
+        Instantiate a [`ParlerTTSConfig`] (or a derived class) from text encoder, audio encoder and decoder
+        configurations.
+
+        Returns:
+            [`ParlerTTSConfig`]: An instance of a configuration object
+        """
+
+        return cls(
+            text_encoder=text_encoder_config.to_dict(),
+            audio_encoder=audio_encoder_config.to_dict(),
+            decoder=decoder_config.to_dict(),
+            **kwargs,
+        )
+
+    @property
+    # This is a property because you might want to change the codec model on the fly
+    def sampling_rate(self):
+        return self.audio_encoder.sampling_rate

capspeech/ar/parler_tts/dac_wrapper/__init__.py

@@ -0,0 +1,2 @@
+from .configuration_dac import DACConfig
+from .modeling_dac import DACModel

capspeech/ar/parler_tts/dac_wrapper/configuration_dac.py

@@ -0,0 +1,27 @@
+
+from transformers import PretrainedConfig
+from importlib.metadata import version
+from packaging.version import Version
+
+
+class DACConfig(PretrainedConfig):
+    model_type = "dac" if Version(version("transformers"))<= Version("4.44.2dev") else "dac_on_the_hub"
+
+    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)

capspeech/ar/parler_tts/dac_wrapper/modeling_dac.py

@@ -0,0 +1,164 @@
+import torch
+from dac.model import DAC
+from torch import nn
+
+from transformers import PreTrainedModel
+from transformers.models.encodec.modeling_encodec import EncodecDecoderOutput, EncodecEncoderOutput
+
+from .configuration_dac import DACConfig
+
+
+# model doesn't support batching yet
+
+
+class DACModel(PreTrainedModel):
+    config_class = DACConfig
+    main_input_name = "input_values"
+
+    # Set main input to 'input_values' for voice steering
+    main_input_name = "input_values"
+
+    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,
+        )
+        
+        self.remove_weight_norm()
+        self.apply_weight_norm()
+
+    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("`DACModel.forward` not implemented yet")
+    
+
+    def apply_weight_norm(self):
+        weight_norm = nn.utils.weight_norm
+        if hasattr(nn.utils.parametrizations, "weight_norm"):
+            weight_norm = nn.utils.parametrizations.weight_norm
+
+        def _apply_weight_norm(module):
+            if isinstance(module, nn.Conv1d) or isinstance(module, nn.ConvTranspose1d):
+                weight_norm(module)
+
+        self.apply(_apply_weight_norm)
+
+
+    def remove_weight_norm(self):
+        def _remove_weight_norm(module):
+            if isinstance(module, nn.Conv1d) or isinstance(module, nn.ConvTranspose1d):
+                nn.utils.remove_weight_norm(module)
+        self.apply(_remove_weight_norm)

capspeech/ar/parler_tts/logits_processors.py

@@ -0,0 +1,54 @@
+from transformers import LogitsProcessor, LogitsProcessorList
+from transformers.pytorch_utils import isin_mps_friendly
+import math
+import torch
+
+class ParlerTTSLogitsProcessor(LogitsProcessor):
+    r"""This processor ensures that the delayed pattern mask constraints are respected.
+
+    <Tip warning={true}>
+
+    This logits processor is exclusively compatible with Parler-TTS. 
+    See the model documentation for examples.
+
+    </Tip>
+
+    Args:
+        eos_token_id (`Union[int, List[int], torch.Tensor]`):
+            The id(s) of the *end-of-sequence* token.
+        min_eos_p (`float`, *optional*):
+            Minimum end of speech threshold.
+    """
+
+    def __init__(self, eos_token_id, num_codebooks: int, batch_size: int, device: str = "cpu"):
+        if not isinstance(eos_token_id, torch.Tensor):
+            if isinstance(eos_token_id, int):
+                eos_token_id = [eos_token_id]
+            eos_token_id = torch.tensor(eos_token_id, device=device)
+        self.eos_token_id = eos_token_id
+        self.batch_size = batch_size
+
+        if torch.is_floating_point(eos_token_id) or (eos_token_id < 0).any():
+            raise ValueError(f"`eos_token_id` has to be a list of positive integers, but is {eos_token_id}")
+
+        self.num_codebooks = num_codebooks
+        self.device = device
+
+
+        self.codebook_idx = torch.arange(self.batch_size*self.num_codebooks, device=self.device)
+        self.first_codebooks_unfinished = torch.arange(batch_size, device=device)*num_codebooks
+        
+        max_codebooks = torch.arange(self.batch_size, device=self.device)*self.num_codebooks + self.num_codebooks -1
+        self.max_codebooks = max_codebooks
+        
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        
+        is_eos = isin_mps_friendly(input_ids, self.eos_token_id).sum(1)
+        
+        self.first_codebooks_unfinished = torch.where((is_eos[self.first_codebooks_unfinished]>0) & (self.first_codebooks_unfinished<self.max_codebooks), self.first_codebooks_unfinished+1, self.first_codebooks_unfinished)
+                
+        # every codebook higher than the first one unfinished will never be eos
+        eos_token_mask = self.codebook_idx > self.first_codebooks_unfinished.repeat_interleave(self.num_codebooks)
+        scores[eos_token_mask, self.eos_token_id] = -math.inf
+        
+        return scores

capspeech/ar/parler_tts/streamer.py

@@ -0,0 +1,147 @@
+
+from .modeling_parler_tts import ParlerTTSForConditionalGeneration
+from transformers.generation.streamers import BaseStreamer
+from typing import Optional
+import torch
+import numpy as np
+import math
+from queue import Queue
+
+
+class ParlerTTSStreamer(BaseStreamer):
+    def __init__(
+        self,
+        model: ParlerTTSForConditionalGeneration,
+        device: Optional[str] = None,
+        play_steps: Optional[int] = 10,
+        stride: Optional[int] = None,
+        timeout: Optional[float] = None,
+    ):
+        """
+        Streamer that stores playback-ready audio in a queue, to be used by a downstream application as an iterator. This is
+        useful for applications that benefit from accessing the generated audio in a non-blocking way (e.g. in an interactive
+        Gradio demo).
+        Parameters:
+            model (`ParlerTTSForConditionalGeneration`):
+                The Parler-TTS model used to generate the audio waveform.
+            device (`str`, *optional*):
+                The torch device on which to run the computation. If `None`, will default to the device of the model.
+            play_steps (`int`, *optional*, defaults to 10):
+                The number of generation steps with which to return the generated audio array. Using fewer steps will
+                mean the first chunk is ready faster, but will require more codec decoding steps overall. This value
+                should be tuned to your device and latency requirements.
+            stride (`int`, *optional*):
+                The window (stride) between adjacent audio samples. Using a stride between adjacent audio samples reduces
+                the hard boundary between them, giving smoother playback. If `None`, will default to a value equivalent to
+                play_steps // 6 in the audio space.
+            timeout (`int`, *optional*):
+                The timeout for the audio queue. If `None`, the queue will block indefinitely. Useful to handle exceptions
+                in `.generate()`, when it is called in a separate thread.
+        """
+        self.decoder = model.decoder
+        self.audio_encoder = model.audio_encoder
+        self.generation_config = model.generation_config
+        self.device = device if device is not None else model.device
+        self.use_audio_scales = model.use_audio_scales
+        self.use_4dim_audio_codes = model.use_4dim_audio_codes
+        self.audio_kwargs = {}
+        if self.use_audio_scales:
+            self.audio_kwargs["audio_scales"] = [None]
+
+        # variables used in the streaming process
+        self.play_steps = play_steps
+        if stride is not None:
+            self.stride = stride
+        else:
+            hop_length = math.floor(self.audio_encoder.config.sampling_rate / self.audio_encoder.config.frame_rate)
+            self.stride = hop_length * (play_steps - self.decoder.num_codebooks) // 6
+        self.token_cache = None
+        self.to_yield = 0
+
+        # varibles used in the thread process
+        self.audio_queue = Queue()
+        self.stop_signal = None
+        self.timeout = timeout
+
+    def apply_delay_pattern_mask(self, input_ids):
+        # build the delay pattern mask for offsetting each codebook prediction by 1 (this behaviour is specific to Parler)
+        _, delay_pattern_mask = self.decoder.build_delay_pattern_mask(
+            input_ids[:, :1],
+            bos_token_id=self.generation_config.bos_token_id,
+            pad_token_id=self.generation_config.decoder_start_token_id,
+            max_length=input_ids.shape[-1],
+        )
+        # apply the pattern mask to the input ids
+        input_ids = self.decoder.apply_delay_pattern_mask(input_ids, delay_pattern_mask)
+
+        # revert the pattern delay mask by filtering the pad token id
+        mask = (delay_pattern_mask != self.generation_config.bos_token_id) & (delay_pattern_mask != self.generation_config.pad_token_id)
+        input_ids = input_ids[mask].reshape(1, self.decoder.num_codebooks, -1)
+
+        if self.use_4dim_audio_codes:
+            # append the frame dimension back to the audio codes
+            input_ids = input_ids[None, ...]
+
+        # send the input_ids to the correct device
+        input_ids = input_ids.to(self.audio_encoder.device)
+
+        decode_sequentially = (
+            self.generation_config.bos_token_id in input_ids
+            or self.generation_config.pad_token_id in input_ids
+            or self.generation_config.eos_token_id in input_ids
+        )
+        if not decode_sequentially:
+            sample = self.audio_encoder.decode(
+                audio_codes=input_ids,
+                **self.audio_kwargs,
+            ).audio_values
+            output_values = sample if sample.ndim == 3 else sample.unsqueeze(0)
+        else:
+            sample = input_ids[:, 0] if self.use_4dim_audio_codes else input_ids[0]
+            sample_mask = ((sample >= self.audio_encoder.config.codebook_size).sum(dim=(0, 1)) == 0) if self.use_4dim_audio_codes else ((sample >= self.audio_encoder.config.codebook_size).sum(dim=0) == 0)
+            sample = sample[:, :, sample_mask] if self.use_4dim_audio_codes else sample[:, sample_mask]
+            sample = self.audio_encoder.decode(audio_codes=sample[None, ...], **self.audio_kwargs).audio_values
+            output_values = sample if sample.ndim == 3 else sample.unsqueeze(0)
+
+        audio_values = output_values[0, 0]
+        return audio_values.cpu().float().numpy()
+
+    def put(self, value):
+        batch_size = value.shape[0] // self.decoder.num_codebooks
+        if batch_size > 1:
+            raise ValueError("ParlerTTSStreamer only supports batch size 1")
+
+        if self.token_cache is None:
+            self.token_cache = value
+        else:
+            self.token_cache = torch.concatenate([self.token_cache, value[:, None]], dim=-1)
+
+        if self.token_cache.shape[-1] % self.play_steps == 0:
+            audio_values = self.apply_delay_pattern_mask(self.token_cache)
+            self.on_finalized_audio(audio_values[self.to_yield : -self.stride])
+            self.to_yield += len(audio_values) - self.to_yield - self.stride
+
+    def end(self):
+        """Flushes any remaining cache and appends the stop symbol."""
+        if self.token_cache is not None:
+            audio_values = self.apply_delay_pattern_mask(self.token_cache)
+        else:
+            audio_values = np.zeros(self.to_yield)
+
+        self.on_finalized_audio(audio_values[self.to_yield :], stream_end=True)
+
+    def on_finalized_audio(self, audio: np.ndarray, stream_end: bool = False):
+        """Put the new audio in the queue. If the stream is ending, also put a stop signal in the queue."""
+        self.audio_queue.put(audio, timeout=self.timeout)
+        if stream_end:
+            self.audio_queue.put(self.stop_signal, timeout=self.timeout)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        value = self.audio_queue.get(timeout=self.timeout)
+        if not isinstance(value, np.ndarray) and value == self.stop_signal:
+            raise StopIteration()
+        else:
+            return value

capspeech/ar/pretrain.sh

@@ -0,0 +1,68 @@
+# Please log in to huggingface first
+
+MLS_WAV_DIR='' # downloaded mls wav path
+LIBRITTSRMIX_WAV_DIR='' # downloaded librittsrmix wav path
+GIGASPEECH_WAV_DIR='' # downloaded gigaspeech wav path
+COMMONVOICE_WAV_DIR='' # downloaded commonvoice wav path
+EMILIA_WAV_DIR='' # downloaded emilia wav path
+OUTPUT_DIR="./output_pretraining/" # output dir, to save checkpoints
+TEMPORY_SAVE_TO_DISK="./audio_code_pretraining/" # dac codec saved dir
+SAVE_TO_DISK="./dataset_pretraining/" # huggingface metadata saved dir
+WANDB_KEY='' # your wandb key for logging
+
+export CUDA_LAUNCH_BLOCKING=1
+export TORCH_USE_CUDA_DSA=1
+
+accelerate launch ./training/run_parler_tts_training.py \
+    --model_name_or_path "parler-tts/parler-tts-mini-v1" \
+    --feature_extractor_name "parler-tts/dac_44khZ_8kbps" \
+    --description_tokenizer_name "google/flan-t5-large" \
+    --prompt_tokenizer_name "google/flan-t5-large" \
+    --report_to "wandb" \
+    --wandb_key ${WANDB_KEY} \
+    --overwrite_output_dir true \
+    --train_dataset_name "OpenSound/CapSpeech" \
+    --train_split_name "train_PT" \
+    --eval_dataset_name "OpenSound/CapSpeech" \
+    --eval_split_name "validation_PT" \
+    --mls_dir ${MLS_WAV_DIR} \
+    --librittsrmix_dir ${LIBRITTSRMIX_WAV_DIR} \
+    --gigaspeech_dir ${GIGASPEECH_WAV_DIR} \
+    --commonvoice_dir ${COMMONVOICE_WAV_DIR} \
+    --emilia_dir ${EMILIA_WAV_DIR} \
+    --max_eval_samples 96 \
+    --per_device_eval_batch_size 32 \
+    --target_audio_column_name "audio_path" \
+    --description_column_name "caption" \
+    --source_column_name "source" \
+    --prompt_column_name "text" \
+    --max_duration_in_seconds 20 \
+    --min_duration_in_seconds 3 \
+    --max_text_length 600 \
+    --preprocessing_num_workers 32 \
+    --do_train true \
+    --num_train_epochs 10 \
+    --gradient_accumulation_steps 6 \
+    --gradient_checkpointing false \
+    --per_device_train_batch_size 4 \
+    --learning_rate 0.001 \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.99 \
+    --weight_decay 0.01 \
+    --lr_scheduler_type "constant_with_warmup" \
+    --warmup_steps 5000 \
+    --logging_steps 200 \
+    --freeze_text_encoder false \
+    --per_device_eval_batch_size 4 \
+    --audio_encoder_per_device_batch_size 24 \
+    --dtype "float16" \
+    --seed 456 \
+    --output_dir ${OUTPUT_DIR} \
+    --temporary_save_to_disk ${TEMPORY_SAVE_TO_DISK} \
+    --save_to_disk ${SAVE_TO_DISK} \
+    --dataloader_num_workers 32 \
+    --do_eval \
+    --evaluation_strategy steps \
+    --eval_steps 5000 \
+    --save_steps 5000 \
+    --group_by_length true

capspeech/ar/training/arguments.py

@@ -0,0 +1,403 @@
+from dataclasses import dataclass, field
+from typing import Optional, List
+
+from transformers import Seq2SeqTrainingArguments
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    feature_extractor_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained feature extractor name or path if not the same as model_name"}
+    )
+    description_tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained description tokenizer name or path if not the same as model_name"}
+    )
+    prompt_tokenizer_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "Pretrained prompt tokenizer name or path if not the same as description_tokenizer_name"},
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    pad_token_id: int = field(
+        default=None,
+        metadata={"help": "If specified, change the model pad token id."},
+    )
+    decoder_start_token_id: int = field(
+        default=None,
+        metadata={"help": "If specified, change the model decoder start token id."},
+    )
+    freeze_text_encoder: bool = field(
+        default=False,
+        metadata={"help": "Whether to freeze the text encoder."},
+    )
+    do_sample: bool = field(
+        default=True,
+        metadata={"help": "Whether to do sampling or greedy decoding."},
+    )
+    temperature: float = field(
+        default=1.0,
+        metadata={"help": "Temperature if sampling."},
+    )
+    max_length: int = field(
+        default=2580,
+        metadata={"help": "Generation max length."},
+    )
+    bandwidth: float = field(
+        default=6,
+        metadata={"help": "Audio encoder bandwidth."},
+    )
+    asr_model_name_or_path: str = field(
+        default="distil-whisper/distil-large-v2",
+        metadata={
+            "help": "Used to compute WER during evaluation. Path to pretrained model or model identifier from huggingface.co/models"
+        },
+    )
+    clap_model_name_or_path: str = field(
+        default="laion/larger_clap_music_and_speech",
+        metadata={
+            "help": "Used to compute audio similarity during evaluation. Path to pretrained model or model identifier from huggingface.co/models"
+        },
+    )
+    attn_implementation: str = field(
+        default="eager",
+        metadata={"help": "Attention implementation used. One of `eager`, `sdpa`, `flash_attention_2`"},
+    )
+    cross_attention_implementation_strategy: str = field(
+        default=None,
+        metadata={
+            "help": "If not specified, the cross-attention implementation will be the same as `_attn_implementation`. If `always_eager`, it will always be the eager implementation. If `always_sdpa`, it will always be the sdpa implementation."
+        },
+    )
+    prompt_padding_side: Optional[str] = field(
+        default="left",
+        metadata={
+            "help": "Prompt tokenizer padding side. Defaults to `left`. If the prompt is pre-pended to the codebooks hidden states, it should be padded on the left."
+        },
+    )
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+
+    Using `HfArgumentParser` we can turn this class
+    into argparse arguments to be able to specify them on
+    the command line.
+    """
+
+    train_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
+            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
+        },
+    )
+    train_dataset_config_name: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The configuration name of the training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset configs by a '+' symbol."
+        },
+    )
+    train_split_name: str = field(
+        default="train",
+        metadata={
+            "help": ("The name of the training data set split to use (via the datasets library). Defaults to 'train'")
+        },
+    )
+    train_dataset_samples: str = field(
+        default=None,
+        metadata={
+            "help": "Number of samples in the training data. Load and combine "
+            "multiple datasets by separating dataset samples by a '+' symbol."
+        },
+    )
+    train_metadata_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the metadata training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
+            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
+        },
+    )
+    eval_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the evaluation dataset to use (via the datasets library). Defaults to the training dataset name if unspecified."
+        },
+    )
+    eval_dataset_config_name: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The configuration name of the evaluation dataset to use (via the datasets library). Defaults to the training dataset config name if unspecified"
+        },
+    )
+    eval_split_name: str = field(
+        default="test",
+        metadata={
+            "help": "The name of the evaluation data set split to use (via the datasets library). Defaults to 'test'"
+        },
+    )
+    eval_metadata_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the metadata training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
+            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
+        },
+    )
+    target_audio_column_name: str = field(
+        default="audio",
+        metadata={"help": "The name of the dataset column containing the target audio data. Defaults to 'audio'"},
+    )
+    description_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the description text data. Defaults to 'None'."},
+    )
+    prompt_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the prompt text data. Defaults to 'None'."},
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."}
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of validation examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_duration_in_seconds: float = field(
+        default=35.0,
+        metadata={
+            "help": (
+                "Filter audio files that are longer than `max_duration_in_seconds` seconds to 'max_duration_in_seconds`."
+                "Also, used to set maximum audio length if `pad_to_max_length=True`."
+            )
+        },
+    )
+    min_duration_in_seconds: float = field(
+        default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}
+    )
+    max_text_length: int = field(
+        default=500, metadata={"help": "If set, max description lengths in number of characters."}
+    )
+    max_prompt_token_length: int = field(
+        default=None,
+        metadata={
+            "help": (
+                "If set, filter samples with prompts that are longer than `max_prompt_token_length` tokens."
+                "Also, used to set maximum prompt token length if `pad_to_max_length=True`."
+            )
+        },
+    )
+    max_description_token_length: int = field(
+        default=None,
+        metadata={
+            "help": (
+                "If set, filter samples with descriptions that are longer than `max_description_token_length` tokens."
+                "Also, used to set maximum description token length if `pad_to_max_length=True`."
+            )
+        },
+    )
+    pad_to_max_length: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "If `True`, pad audio, prompt and description to a maximum length set with respectively "
+                "`max_duration_in_seconds`, `max_prompt_token_length`, `max_description_token_length`."
+            )
+        },
+    )
+    preprocessing_only: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether to only do data preprocessing and skip training. This is especially useful when data"
+                " preprocessing errors out in distributed training due to timeout. In this case, one should run the"
+                " preprocessing in a non-distributed setup with `preprocessing_only=True` so that the cached datasets"
+                " can consequently be loaded in distributed training."
+                " In this training script, `save_to_disk` must be set to the path in which the dataset should be saved. "
+            )
+        },
+    )
+    token: str = field(
+        default=None,
+        metadata={
+            "help": (
+                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
+                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
+            )
+        },
+    )
+    use_auth_token: bool = field(
+        default=None,
+        metadata={
+            "help": "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead."
+        },
+    )
+    trust_remote_code: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option "
+                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
+                "execute code present on the Hub on your local machine."
+            )
+        },
+    )
+    add_audio_samples_to_wandb: bool = field(
+        default=False,
+        metadata={"help": "If set and if `wandb` in args.report_to, will add generated audio samples to wandb logs."},
+    )
+    id_column_name: str = field(default=None, metadata={"help": "id column name."})
+    wandb_project: str = field(
+        default="parler-speech",
+        metadata={"help": "The name of the wandb project."},
+    )
+    wandb_run_name: str = field(
+        default=None,
+        metadata={
+            "help": "If specified, the name of the run. If not specified, wandb will give a random name to this run."
+        },
+    )
+    save_to_disk: str = field(
+        default=None,
+        metadata={
+            "help": "If set, will save the dataset to this path if this is an empyt folder. If not empty, will load the datasets from it."
+        },
+    )
+    temporary_save_to_disk: str = field(default=None, metadata={"help": "Temporarily save audio labels here."})
+    save_codec_steps: Optional[int] = field(
+        default=500,
+        metadata={"help": "Temporarily save the audio labels every `save_steps`."},
+    )
+    pad_to_multiple_of: Optional[int] = field(
+        default=2,
+        metadata={"help": ("Pad to multiple of for tokenizers.")},
+    )
+    mls_dir: str = field(
+        default=None,
+        metadata={"help": "mls audio dir"},
+    )
+    librittsrmix_dir: str = field(
+        default=None,
+        metadata={"help": "librittsrmix audio dir"},
+    )
+    gigaspeech_dir: str = field(
+        default=None,
+        metadata={"help": "gigaspeech audio dir"},
+    )
+    commonvoice_dir: str = field(
+        default=None,
+        metadata={"help": "commonvoice audio dir"},
+    )
+    emilia_dir: str = field(
+        default=None,
+        metadata={"help": "emilia audio dir"},
+    )
+    source_column_name: str = field(
+        default="source",
+        metadata={"help": "The name of the source column."},
+    )
+    wandb_key: str = field(
+        default=None,
+        metadata={"help": "wandb key name"},
+    )
+
+
+@dataclass
+class ParlerTTSTrainingArguments(Seq2SeqTrainingArguments):
+    dtype: Optional[str] = field(
+        default="float32",
+        metadata={
+            "help": (
+                "The data type (dtype) in which to run training. One of `float32` (full-precision), "
+                "`float16` or `bfloat16` (both half-precision)."
+            )
+        },
+    )
+    audio_encoder_per_device_batch_size: int = field(
+        default=8,
+        metadata={"help": ("Specify the batch size of the audio encoding pre-processing steps.")},
+    )
+    eval_dataloader_num_workers: Optional[int] = field(
+        default=0,
+        metadata={
+            "help": (
+                "Number of subprocesses to use for evaluation data loading (PyTorch only). 0 means that the data will be loaded in the main process."
+            )
+        },
+    )
+    compute_clap_similarity_metric: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether or not to compute the clap similarity metric between the description and the generation during evalution."
+            )
+        },
+    )
+    compute_noise_level_metric: bool = field(
+        default=True,
+        metadata={"help": ("Whether or not to compute the squim si-sdr measure of the generations.")},
+    )
+    noise_level_to_compute_clean_wer: float = field(
+        default=25,
+        metadata={
+            "help": (
+                "if `compute_noise_level_metric=True`, will compute a 'clean' WER on samples with generated noise higher than `noise_level_to_compute_clean_wer`."
+                "This is a proxy measure to compute WER on clean audios, provided that the model learn to generate clean audios."
+            )
+        },
+    )
+    eval_generation_steps: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Number of update steps between two generation evaluation.  Will default to the same"
+                "value as `eval_steps` if not set. Should be an integer and a multiple of `eval_steps`."
+            )
+        },
+    )       
+    codebook_weights: Optional[List[float]] = field(
+        default=None,
+        metadata={"help": "Weights applied to each codebook."},
+    )

capspeech/ar/training/arguments_captts.py

@@ -0,0 +1,391 @@
+from dataclasses import dataclass, field
+from typing import Optional, List
+
+from transformers import Seq2SeqTrainingArguments
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    feature_extractor_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained feature extractor name or path if not the same as model_name"}
+    )
+    description_tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained description tokenizer name or path if not the same as model_name"}
+    )
+    prompt_tokenizer_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "Pretrained prompt tokenizer name or path if not the same as description_tokenizer_name"},
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    pad_token_id: int = field(
+        default=None,
+        metadata={"help": "If specified, change the model pad token id."},
+    )
+    decoder_start_token_id: int = field(
+        default=None,
+        metadata={"help": "If specified, change the model decoder start token id."},
+    )
+    freeze_text_encoder: bool = field(
+        default=False,
+        metadata={"help": "Whether to freeze the text encoder."},
+    )
+    do_sample: bool = field(
+        default=True,
+        metadata={"help": "Whether to do sampling or greedy decoding."},
+    )
+    temperature: float = field(
+        default=1.0,
+        metadata={"help": "Temperature if sampling."},
+    )
+    max_length: int = field(
+        default=2580,
+        metadata={"help": "Generation max length."},
+    )
+    bandwidth: float = field(
+        default=6,
+        metadata={"help": "Audio encoder bandwidth."},
+    )
+    asr_model_name_or_path: str = field(
+        default="distil-whisper/distil-large-v2",
+        metadata={
+            "help": "Used to compute WER during evaluation. Path to pretrained model or model identifier from huggingface.co/models"
+        },
+    )
+    clap_model_name_or_path: str = field(
+        default="laion/larger_clap_music_and_speech",
+        metadata={
+            "help": "Used to compute audio similarity during evaluation. Path to pretrained model or model identifier from huggingface.co/models"
+        },
+    )
+    attn_implementation: str = field(
+        default="eager",
+        metadata={"help": "Attention implementation used. One of `eager`, `sdpa`, `flash_attention_2`"},
+    )
+    cross_attention_implementation_strategy: str = field(
+        default=None,
+        metadata={
+            "help": "If not specified, the cross-attention implementation will be the same as `_attn_implementation`. If `always_eager`, it will always be the eager implementation. If `always_sdpa`, it will always be the sdpa implementation."
+        },
+    )
+    prompt_padding_side: Optional[str] = field(
+        default="left",
+        metadata={
+            "help": "Prompt tokenizer padding side. Defaults to `left`. If the prompt is pre-pended to the codebooks hidden states, it should be padded on the left."
+        },
+    )
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+
+    Using `HfArgumentParser` we can turn this class
+    into argparse arguments to be able to specify them on
+    the command line.
+    """
+
+    train_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
+            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
+        },
+    )
+    train_dataset_config_name: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The configuration name of the training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset configs by a '+' symbol."
+        },
+    )
+    train_split_name: str = field(
+        default="train",
+        metadata={
+            "help": ("The name of the training data set split to use (via the datasets library). Defaults to 'train'")
+        },
+    )
+    train_dataset_samples: str = field(
+        default=None,
+        metadata={
+            "help": "Number of samples in the training data. Load and combine "
+            "multiple datasets by separating dataset samples by a '+' symbol."
+        },
+    )
+    train_metadata_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the metadata training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
+            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
+        },
+    )
+    eval_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the evaluation dataset to use (via the datasets library). Defaults to the training dataset name if unspecified."
+        },
+    )
+    eval_dataset_config_name: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The configuration name of the evaluation dataset to use (via the datasets library). Defaults to the training dataset config name if unspecified"
+        },
+    )
+    eval_split_name: str = field(
+        default="test",
+        metadata={
+            "help": "The name of the evaluation data set split to use (via the datasets library). Defaults to 'test'"
+        },
+    )
+    eval_metadata_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the metadata training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
+            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
+        },
+    )
+    target_audio_column_name: str = field(
+        default="audio",
+        metadata={"help": "The name of the dataset column containing the target audio data. Defaults to 'audio'"},
+    )
+    description_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the description text data. Defaults to 'None'."},
+    )
+    prompt_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the prompt text data. Defaults to 'None'."},
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."}
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of validation examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_duration_in_seconds: float = field(
+        default=35.0,
+        metadata={
+            "help": (
+                "Filter audio files that are longer than `max_duration_in_seconds` seconds to 'max_duration_in_seconds`."
+                "Also, used to set maximum audio length if `pad_to_max_length=True`."
+            )
+        },
+    )
+    min_duration_in_seconds: float = field(
+        default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}
+    )
+    max_text_length: int = field(
+        default=500, metadata={"help": "If set, max description lengths in number of characters."}
+    )
+    max_prompt_token_length: int = field(
+        default=None,
+        metadata={
+            "help": (
+                "If set, filter samples with prompts that are longer than `max_prompt_token_length` tokens."
+                "Also, used to set maximum prompt token length if `pad_to_max_length=True`."
+            )
+        },
+    )
+    max_description_token_length: int = field(
+        default=None,
+        metadata={
+            "help": (
+                "If set, filter samples with descriptions that are longer than `max_description_token_length` tokens."
+                "Also, used to set maximum description token length if `pad_to_max_length=True`."
+            )
+        },
+    )
+    pad_to_max_length: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "If `True`, pad audio, prompt and description to a maximum length set with respectively "
+                "`max_duration_in_seconds`, `max_prompt_token_length`, `max_description_token_length`."
+            )
+        },
+    )
+    preprocessing_only: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether to only do data preprocessing and skip training. This is especially useful when data"
+                " preprocessing errors out in distributed training due to timeout. In this case, one should run the"
+                " preprocessing in a non-distributed setup with `preprocessing_only=True` so that the cached datasets"
+                " can consequently be loaded in distributed training."
+                " In this training script, `save_to_disk` must be set to the path in which the dataset should be saved. "
+            )
+        },
+    )
+    token: str = field(
+        default=None,
+        metadata={
+            "help": (
+                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
+                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
+            )
+        },
+    )
+    use_auth_token: bool = field(
+        default=None,
+        metadata={
+            "help": "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead."
+        },
+    )
+    trust_remote_code: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option "
+                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
+                "execute code present on the Hub on your local machine."
+            )
+        },
+    )
+    add_audio_samples_to_wandb: bool = field(
+        default=False,
+        metadata={"help": "If set and if `wandb` in args.report_to, will add generated audio samples to wandb logs."},
+    )
+    id_column_name: str = field(default=None, metadata={"help": "id column name."})
+    wandb_project: str = field(
+        default="parler-speech",
+        metadata={"help": "The name of the wandb project."},
+    )
+    wandb_run_name: str = field(
+        default=None,
+        metadata={
+            "help": "If specified, the name of the run. If not specified, wandb will give a random name to this run."
+        },
+    )
+    save_to_disk: str = field(
+        default=None,
+        metadata={
+            "help": "If set, will save the dataset to this path if this is an empyt folder. If not empty, will load the datasets from it."
+        },
+    )
+    temporary_save_to_disk: str = field(default=None, metadata={"help": "Temporarily save audio labels here."})
+    save_codec_steps: Optional[int] = field(
+        default=500,
+        metadata={"help": "Temporarily save the audio labels every `save_steps`."},
+    )
+    pad_to_multiple_of: Optional[int] = field(
+        default=2,
+        metadata={"help": ("Pad to multiple of for tokenizers.")},
+    )
+    librittsr_dir: str = field(
+        default=None,
+        metadata={"help": "librittsr audio dir"},
+    )
+    other_dir: str = field(
+        default=None,
+        metadata={"help": "other audio dir"},
+    )
+    source_column_name: str = field(
+        default="source",
+        metadata={"help": "The name of the source column."},
+    )
+    wandb_key: str = field(
+        default=None,
+        metadata={"help": "wandb key name"},
+    )
+
+
+@dataclass
+class ParlerTTSTrainingArguments(Seq2SeqTrainingArguments):
+    dtype: Optional[str] = field(
+        default="float32",
+        metadata={
+            "help": (
+                "The data type (dtype) in which to run training. One of `float32` (full-precision), "
+                "`float16` or `bfloat16` (both half-precision)."
+            )
+        },
+    )
+    audio_encoder_per_device_batch_size: int = field(
+        default=8,
+        metadata={"help": ("Specify the batch size of the audio encoding pre-processing steps.")},
+    )
+    eval_dataloader_num_workers: Optional[int] = field(
+        default=0,
+        metadata={
+            "help": (
+                "Number of subprocesses to use for evaluation data loading (PyTorch only). 0 means that the data will be loaded in the main process."
+            )
+        },
+    )
+    compute_clap_similarity_metric: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether or not to compute the clap similarity metric between the description and the generation during evalution."
+            )
+        },
+    )
+    compute_noise_level_metric: bool = field(
+        default=True,
+        metadata={"help": ("Whether or not to compute the squim si-sdr measure of the generations.")},
+    )
+    noise_level_to_compute_clean_wer: float = field(
+        default=25,
+        metadata={
+            "help": (
+                "if `compute_noise_level_metric=True`, will compute a 'clean' WER on samples with generated noise higher than `noise_level_to_compute_clean_wer`."
+                "This is a proxy measure to compute WER on clean audios, provided that the model learn to generate clean audios."
+            )
+        },
+    )
+    eval_generation_steps: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Number of update steps between two generation evaluation.  Will default to the same"
+                "value as `eval_steps` if not set. Should be an integer and a multiple of `eval_steps`."
+            )
+        },
+    )       
+    codebook_weights: Optional[List[float]] = field(
+        default=None,
+        metadata={"help": "Weights applied to each codebook."},
+    )

capspeech/ar/training/arguments_capttsse.py

@@ -0,0 +1,387 @@
+from dataclasses import dataclass, field
+from typing import Optional, List
+
+from transformers import Seq2SeqTrainingArguments
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    feature_extractor_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained feature extractor name or path if not the same as model_name"}
+    )
+    description_tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained description tokenizer name or path if not the same as model_name"}
+    )
+    prompt_tokenizer_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "Pretrained prompt tokenizer name or path if not the same as description_tokenizer_name"},
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    pad_token_id: int = field(
+        default=None,
+        metadata={"help": "If specified, change the model pad token id."},
+    )
+    decoder_start_token_id: int = field(
+        default=None,
+        metadata={"help": "If specified, change the model decoder start token id."},
+    )
+    freeze_text_encoder: bool = field(
+        default=False,
+        metadata={"help": "Whether to freeze the text encoder."},
+    )
+    do_sample: bool = field(
+        default=True,
+        metadata={"help": "Whether to do sampling or greedy decoding."},
+    )
+    temperature: float = field(
+        default=1.0,
+        metadata={"help": "Temperature if sampling."},
+    )
+    max_length: int = field(
+        default=2580,
+        metadata={"help": "Generation max length."},
+    )
+    bandwidth: float = field(
+        default=6,
+        metadata={"help": "Audio encoder bandwidth."},
+    )
+    asr_model_name_or_path: str = field(
+        default="distil-whisper/distil-large-v2",
+        metadata={
+            "help": "Used to compute WER during evaluation. Path to pretrained model or model identifier from huggingface.co/models"
+        },
+    )
+    clap_model_name_or_path: str = field(
+        default="laion/larger_clap_music_and_speech",
+        metadata={
+            "help": "Used to compute audio similarity during evaluation. Path to pretrained model or model identifier from huggingface.co/models"
+        },
+    )
+    attn_implementation: str = field(
+        default="eager",
+        metadata={"help": "Attention implementation used. One of `eager`, `sdpa`, `flash_attention_2`"},
+    )
+    cross_attention_implementation_strategy: str = field(
+        default=None,
+        metadata={
+            "help": "If not specified, the cross-attention implementation will be the same as `_attn_implementation`. If `always_eager`, it will always be the eager implementation. If `always_sdpa`, it will always be the sdpa implementation."
+        },
+    )
+    prompt_padding_side: Optional[str] = field(
+        default="left",
+        metadata={
+            "help": "Prompt tokenizer padding side. Defaults to `left`. If the prompt is pre-pended to the codebooks hidden states, it should be padded on the left."
+        },
+    )
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+
+    Using `HfArgumentParser` we can turn this class
+    into argparse arguments to be able to specify them on
+    the command line.
+    """
+
+    train_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
+            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
+        },
+    )
+    train_dataset_config_name: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The configuration name of the training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset configs by a '+' symbol."
+        },
+    )
+    train_split_name: str = field(
+        default="train",
+        metadata={
+            "help": ("The name of the training data set split to use (via the datasets library). Defaults to 'train'")
+        },
+    )
+    train_dataset_samples: str = field(
+        default=None,
+        metadata={
+            "help": "Number of samples in the training data. Load and combine "
+            "multiple datasets by separating dataset samples by a '+' symbol."
+        },
+    )
+    train_metadata_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the metadata training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
+            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
+        },
+    )
+    eval_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the evaluation dataset to use (via the datasets library). Defaults to the training dataset name if unspecified."
+        },
+    )
+    eval_dataset_config_name: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The configuration name of the evaluation dataset to use (via the datasets library). Defaults to the training dataset config name if unspecified"
+        },
+    )
+    eval_split_name: str = field(
+        default="test",
+        metadata={
+            "help": "The name of the evaluation data set split to use (via the datasets library). Defaults to 'test'"
+        },
+    )
+    eval_metadata_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the metadata training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
+            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
+        },
+    )
+    target_audio_column_name: str = field(
+        default="audio",
+        metadata={"help": "The name of the dataset column containing the target audio data. Defaults to 'audio'"},
+    )
+    description_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the description text data. Defaults to 'None'."},
+    )
+    prompt_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the prompt text data. Defaults to 'None'."},
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."}
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of validation examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_duration_in_seconds: float = field(
+        default=35.0,
+        metadata={
+            "help": (
+                "Filter audio files that are longer than `max_duration_in_seconds` seconds to 'max_duration_in_seconds`."
+                "Also, used to set maximum audio length if `pad_to_max_length=True`."
+            )
+        },
+    )
+    min_duration_in_seconds: float = field(
+        default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}
+    )
+    max_text_length: int = field(
+        default=500, metadata={"help": "If set, max description lengths in number of characters."}
+    )
+    max_prompt_token_length: int = field(
+        default=None,
+        metadata={
+            "help": (
+                "If set, filter samples with prompts that are longer than `max_prompt_token_length` tokens."
+                "Also, used to set maximum prompt token length if `pad_to_max_length=True`."
+            )
+        },
+    )
+    max_description_token_length: int = field(
+        default=None,
+        metadata={
+            "help": (
+                "If set, filter samples with descriptions that are longer than `max_description_token_length` tokens."
+                "Also, used to set maximum description token length if `pad_to_max_length=True`."
+            )
+        },
+    )
+    pad_to_max_length: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "If `True`, pad audio, prompt and description to a maximum length set with respectively "
+                "`max_duration_in_seconds`, `max_prompt_token_length`, `max_description_token_length`."
+            )
+        },
+    )
+    preprocessing_only: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether to only do data preprocessing and skip training. This is especially useful when data"
+                " preprocessing errors out in distributed training due to timeout. In this case, one should run the"
+                " preprocessing in a non-distributed setup with `preprocessing_only=True` so that the cached datasets"
+                " can consequently be loaded in distributed training."
+                " In this training script, `save_to_disk` must be set to the path in which the dataset should be saved. "
+            )
+        },
+    )
+    token: str = field(
+        default=None,
+        metadata={
+            "help": (
+                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
+                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
+            )
+        },
+    )
+    use_auth_token: bool = field(
+        default=None,
+        metadata={
+            "help": "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead."
+        },
+    )
+    trust_remote_code: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option "
+                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
+                "execute code present on the Hub on your local machine."
+            )
+        },
+    )
+    add_audio_samples_to_wandb: bool = field(
+        default=False,
+        metadata={"help": "If set and if `wandb` in args.report_to, will add generated audio samples to wandb logs."},
+    )
+    id_column_name: str = field(default=None, metadata={"help": "id column name."})
+    wandb_project: str = field(
+        default="parler-speech",
+        metadata={"help": "The name of the wandb project."},
+    )
+    wandb_run_name: str = field(
+        default=None,
+        metadata={
+            "help": "If specified, the name of the run. If not specified, wandb will give a random name to this run."
+        },
+    )
+    save_to_disk: str = field(
+        default=None,
+        metadata={
+            "help": "If set, will save the dataset to this path if this is an empyt folder. If not empty, will load the datasets from it."
+        },
+    )
+    temporary_save_to_disk: str = field(default=None, metadata={"help": "Temporarily save audio labels here."})
+    save_codec_steps: Optional[int] = field(
+        default=500,
+        metadata={"help": "Temporarily save the audio labels every `save_steps`."},
+    )
+    pad_to_multiple_of: Optional[int] = field(
+        default=2,
+        metadata={"help": ("Pad to multiple of for tokenizers.")},
+    )
+    librittsrmix_dir: str = field(
+        default=None,
+        metadata={"help": "librittsrmix audio dir"},
+    )
+    source_column_name: str = field(
+        default="source",
+        metadata={"help": "The name of the source column."},
+    )
+    wandb_key: str = field(
+        default=None,
+        metadata={"help": "wandb key name"},
+    )
+
+
+@dataclass
+class ParlerTTSTrainingArguments(Seq2SeqTrainingArguments):
+    dtype: Optional[str] = field(
+        default="float32",
+        metadata={
+            "help": (
+                "The data type (dtype) in which to run training. One of `float32` (full-precision), "
+                "`float16` or `bfloat16` (both half-precision)."
+            )
+        },
+    )
+    audio_encoder_per_device_batch_size: int = field(
+        default=8,
+        metadata={"help": ("Specify the batch size of the audio encoding pre-processing steps.")},
+    )
+    eval_dataloader_num_workers: Optional[int] = field(
+        default=0,
+        metadata={
+            "help": (
+                "Number of subprocesses to use for evaluation data loading (PyTorch only). 0 means that the data will be loaded in the main process."
+            )
+        },
+    )
+    compute_clap_similarity_metric: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether or not to compute the clap similarity metric between the description and the generation during evalution."
+            )
+        },
+    )
+    compute_noise_level_metric: bool = field(
+        default=True,
+        metadata={"help": ("Whether or not to compute the squim si-sdr measure of the generations.")},
+    )
+    noise_level_to_compute_clean_wer: float = field(
+        default=25,
+        metadata={
+            "help": (
+                "if `compute_noise_level_metric=True`, will compute a 'clean' WER on samples with generated noise higher than `noise_level_to_compute_clean_wer`."
+                "This is a proxy measure to compute WER on clean audios, provided that the model learn to generate clean audios."
+            )
+        },
+    )
+    eval_generation_steps: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Number of update steps between two generation evaluation.  Will default to the same"
+                "value as `eval_steps` if not set. Should be an integer and a multiple of `eval_steps`."
+            )
+        },
+    )       
+    codebook_weights: Optional[List[float]] = field(
+        default=None,
+        metadata={"help": "Weights applied to each codebook."},
+    )

capspeech/ar/training/data.py

@@ -0,0 +1,277 @@
+import logging
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Set, Union
+import os
+import datasets
+import numpy as np
+import torch
+from accelerate import Accelerator
+from datasets import Dataset, IterableDataset, concatenate_datasets, interleave_datasets, load_dataset
+from tqdm import tqdm
+from transformers import AutoFeatureExtractor, AutoTokenizer
+import torchaudio
+import torchaudio.transforms as T
+
+@dataclass
+class DataCollatorEncodecWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received to the longest sequence in the batch or
+    to `max_length` if `max_length` is set and `padding=max_length`.
+    """
+
+    feature_extractor: AutoFeatureExtractor
+    audio_column_name: str
+    mls_dir: Optional[str] = None
+    librittsrmix_dir: Optional[str] = None
+    gigaspeech_dir: Optional[str] = None
+    commonvoice_dir: Optional[str] = None
+    emilia_dir: Optional[str] = None
+    feature_extractor_input_name: Optional[str] = "input_values"
+    max_length: Optional[int] = None
+    padding: Optional[str] = "longest"
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need
+        # different padding methods
+        sampling_rate = self.feature_extractor.sampling_rate
+        # load audio
+        audios = []
+        for f in features:
+            path = f[self.audio_column_name]
+            source = f["source"]
+            if source == "libritts-r":
+                path = os.path.join(self.librittsrmix_dir, path)
+            elif source == "mls":
+                path = os.path.join(self.mls_dir, path)
+            elif source == "gigaspeech":
+                path = os.path.join(self.gigaspeech_dir, path)
+            elif source == "commonvoice":
+                path = os.path.join(self.commonvoice_dir, path)
+            elif source == "emilia":
+                path = os.path.join(self.emilia_dir, path)
+            else:
+                raise ValueError(source)
+
+            if os.path.exists(path):
+                waveform, sr = torchaudio.load(path)
+                if sr != sampling_rate:
+                    resampler = T.Resample(orig_freq=sr, new_freq=sampling_rate)
+                    waveform = resampler(waveform)
+                if waveform.shape[0] > 1:
+                    waveform = waveform.mean(dim=0, keepdim=True)
+                audios.append(waveform.squeeze())
+            else:
+                print(f"Read error: {path}")
+
+
+        len_audio = [len(audio) for audio in audios]
+        if self.max_length is not None:
+            audios = [audio[: min(l, self.max_length)] for audio, l in zip(audios, len_audio)]
+
+        # since resampling has already been performed in the 'load_multiple_datasets' function,
+        # a fixed sampling_rate(44100hz) is passed to the feature_extractor.
+        batch = self.feature_extractor(
+            [np.asarray(a, dtype=np.float32) for a in audios], sampling_rate=sampling_rate, return_tensors="pt", padding=self.padding, max_length=self.max_length
+        )
+        batch["len_audio"] = torch.tensor(len_audio).unsqueeze(1)
+        return batch
+
+
+@dataclass
+class DataCollatorParlerTTSWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received.
+    Args:
+        prompt_tokenizer (:class:`~transformers.AutoTokenizer`)
+            The prompt_tokenizer used for proccessing the data.
+        description_tokenizer (:class:`~transformers.AutoTokenizer`)
+            The description_tokenizer used for proccessing the data.
+        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
+            among:
+            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+              sequence if provided).
+            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
+              maximum acceptable input length for the model if that argument is not provided.
+            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
+              different lengths).
+        pad_to_multiple_of (:obj:`int`, `optional`):
+            If set will pad the sequence to a multiple of the provided value.
+            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
+            7.5 (Volta).
+    """
+
+    prompt_tokenizer: AutoTokenizer
+    description_tokenizer: AutoTokenizer
+    padding: Union[bool, str] = "longest"
+    pad_to_multiple_of: Optional[int] = None
+    prompt_max_length: Optional[int] = None
+    description_max_length: Optional[int] = None
+    audio_max_length: Optional[int] = None
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need
+        # different padding methods
+
+        labels = [torch.tensor(feature["labels"]).transpose(0, 1) for feature in features]
+        # (bsz, seq_len, num_codebooks)
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=-100)
+        if self.audio_max_length is not None and self.padding == "max_length":
+            labels = torch.nn.functional.pad(
+                labels, pad=(0, 0, 0, max(self.audio_max_length - labels.shape[1], 0)), value=-100
+            )
+
+        input_ids = [{"input_ids": feature["input_ids"]} for feature in features]
+
+        input_ids = self.description_tokenizer.pad(
+            input_ids,
+            return_tensors="pt",
+            padding=self.padding,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            max_length=self.description_max_length,
+        )
+
+        batch = {"labels": labels, **input_ids}
+
+        prompt_input_ids = [{"input_ids": feature["prompt_input_ids"]} for feature in features]
+        prompt_input_ids = self.prompt_tokenizer.pad(
+            prompt_input_ids,
+            return_tensors="pt",
+            padding=self.padding,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            max_length=self.prompt_max_length,
+        )
+
+        batch["prompt_input_ids"] = prompt_input_ids["input_ids"]
+        if "attention_mask" in prompt_input_ids:
+            batch["prompt_attention_mask"] = prompt_input_ids["attention_mask"]
+
+        return batch
+
+
+def convert_dataset_str_to_list(
+    dataset_names,
+    splits=None,
+    dataset_samples=None,
+    default_split="train",
+):
+    if isinstance(dataset_names, str):
+        dataset_names = dataset_names.split("+")
+        splits = splits.split("+") if splits is not None else None
+        dataset_samples = dataset_samples.split("+") if dataset_samples is not None else None
+
+    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 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)
+
+    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,
+                "split": splits[i],
+                "samples": dataset_samples[i],
+            }
+        )
+    return dataset_names_dict
+
+
+def load_multiple_datasets(
+    accelerator: Accelerator,
+    dataset_names: Union[List, str],
+    splits: Optional[Union[List, str]] = None,
+    label_column_names: Optional[List] = None,
+    stopping_strategy: Optional[str] = "first_exhausted",
+    dataset_samples: Optional[Union[List, np.array]] = None,
+    streaming: Optional[bool] = False,
+    seed: Optional[int] = None,
+    id_column_name: Optional[str] = None,
+    columns_to_keep: Optional[Set[str]] = None,
+    prompt_column_name: Optional[str] = None,
+    sampling_rate: Optional[int] = None,
+    audio_column_name: Optional[str] = None,
+    logger: Optional[logging.Logger] = None,
+    librittsrmix_dir: Optional[Union[List, str]] = None,
+    mls_dir: Optional[Union[List, str]] = None,
+    gigaspeech_dir: Optional[Union[List, str]] = None,
+    commonvoice_dir: Optional[Union[List, str]] = None,
+    emilia_dir: Optional[Union[List, str]] = None,
+    **kwargs,
+) -> Union[Dataset, IterableDataset]:
+    dataset_names_dict = convert_dataset_str_to_list(
+        dataset_names, splits, label_column_names, dataset_samples
+    )
+
+    if dataset_samples is not None:
+        dataset_samples = [ds_dict["samples"] for ds_dict in dataset_names_dict]
+        probabilities = np.array(dataset_samples) / np.sum(dataset_samples)
+    else:
+        probabilities = None
+
+    all_datasets = []
+    # iterate over the datasets we want to interleave
+    for dataset_dict in tqdm(dataset_names_dict, desc="Combining datasets..."):
+        with accelerator.local_main_process_first():
+            dataset = load_dataset(
+                dataset_dict["name"],
+                split=dataset_dict["split"],
+                streaming=streaming,
+                **kwargs,
+            )
+            dataset_features = dataset.features.keys()
+
+            if columns_to_keep is not None:
+                dataset = dataset.remove_columns(set(dataset_features - columns_to_keep))
+
+            def resolve_path(example):
+                path = example["audio_path"]
+                source = example["source"]
+
+                if source == "libritts-r":
+                    full_path = os.path.join(librittsrmix_dir, path)
+                elif source == "mls":
+                    full_path = os.path.join(mls_dir, path)
+                elif source == "gigaspeech":
+                    full_path = os.path.join(gigaspeech_dir, path)
+                elif source == "commonvoice":
+                    full_path = os.path.join(commonvoice_dir, path)
+                elif source == "emilia":
+                    full_path = os.path.join(emilia_dir, path)
+                else:
+                    return False  # unknown source
+
+                return os.path.exists(full_path)
+                
+            dataset = dataset.filter(resolve_path, num_proc=16)
+
+        all_datasets.append(dataset)
+
+    if len(all_datasets) == 1:
+        # we have a single dataset so just return it as is
+        return all_datasets[0]
+
+    if streaming:
+        interleaved_dataset = interleave_datasets(
+            all_datasets,
+            stopping_strategy=stopping_strategy,
+            probabilities=probabilities,
+            seed=seed,
+        )
+    else:
+        with accelerator.local_main_process_first():
+            interleaved_dataset = concatenate_datasets(all_datasets)
+
+    return interleaved_dataset

capspeech/ar/training/data_captts.py

@@ -0,0 +1,255 @@
+import logging
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Set, Union
+import os
+import datasets
+import numpy as np
+import torch
+from accelerate import Accelerator
+from datasets import Dataset, IterableDataset, concatenate_datasets, interleave_datasets, load_dataset
+from tqdm import tqdm
+from transformers import AutoFeatureExtractor, AutoTokenizer
+import torchaudio
+import torchaudio.transforms as T
+
+@dataclass
+class DataCollatorEncodecWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received to the longest sequence in the batch or
+    to `max_length` if `max_length` is set and `padding=max_length`.
+    """
+
+    feature_extractor: AutoFeatureExtractor
+    audio_column_name: str
+    librittsr_dir: Optional[str] = None
+    other_dir: Optional[str] = None
+    feature_extractor_input_name: Optional[str] = "input_values"
+    max_length: Optional[int] = None
+    padding: Optional[str] = "longest"
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need
+        # different padding methods
+        sampling_rate = self.feature_extractor.sampling_rate
+        # load audio
+        audios = []
+        for f in features:
+            path = f[self.audio_column_name]
+            source = f["source"]
+            if source == "libritts-r":
+                path = os.path.join(self.librittsr_dir, path)
+            else:
+                path = os.path.join(self.other_dir, path)
+
+            if os.path.exists(path):
+                waveform, sr = torchaudio.load(path)
+                if sr != sampling_rate:
+                    resampler = T.Resample(orig_freq=sr, new_freq=sampling_rate)
+                    waveform = resampler(waveform)
+                if waveform.shape[0] > 1:
+                    waveform = waveform.mean(dim=0, keepdim=True)
+                audios.append(waveform.squeeze())
+            else:
+                print(f"Read error: {path}")
+
+
+        len_audio = [len(audio) for audio in audios]
+        if self.max_length is not None:
+            audios = [audio[: min(l, self.max_length)] for audio, l in zip(audios, len_audio)]
+
+        # since resampling has already been performed in the 'load_multiple_datasets' function,
+        # a fixed sampling_rate(44100hz) is passed to the feature_extractor.
+        batch = self.feature_extractor(
+            [np.asarray(a, dtype=np.float32) for a in audios], sampling_rate=sampling_rate, return_tensors="pt", padding=self.padding, max_length=self.max_length
+        )
+        batch["len_audio"] = torch.tensor(len_audio).unsqueeze(1)
+        return batch
+
+
+@dataclass
+class DataCollatorParlerTTSWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received.
+    Args:
+        prompt_tokenizer (:class:`~transformers.AutoTokenizer`)
+            The prompt_tokenizer used for proccessing the data.
+        description_tokenizer (:class:`~transformers.AutoTokenizer`)
+            The description_tokenizer used for proccessing the data.
+        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
+            among:
+            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+              sequence if provided).
+            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
+              maximum acceptable input length for the model if that argument is not provided.
+            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
+              different lengths).
+        pad_to_multiple_of (:obj:`int`, `optional`):
+            If set will pad the sequence to a multiple of the provided value.
+            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
+            7.5 (Volta).
+    """
+
+    prompt_tokenizer: AutoTokenizer
+    description_tokenizer: AutoTokenizer
+    padding: Union[bool, str] = "longest"
+    pad_to_multiple_of: Optional[int] = None
+    prompt_max_length: Optional[int] = None
+    description_max_length: Optional[int] = None
+    audio_max_length: Optional[int] = None
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need
+        # different padding methods
+
+        labels = [torch.tensor(feature["labels"]).transpose(0, 1) for feature in features]
+        # (bsz, seq_len, num_codebooks)
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=-100)
+        if self.audio_max_length is not None and self.padding == "max_length":
+            labels = torch.nn.functional.pad(
+                labels, pad=(0, 0, 0, max(self.audio_max_length - labels.shape[1], 0)), value=-100
+            )
+
+        input_ids = [{"input_ids": feature["input_ids"]} for feature in features]
+
+        input_ids = self.description_tokenizer.pad(
+            input_ids,
+            return_tensors="pt",
+            padding=self.padding,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            max_length=self.description_max_length,
+        )
+
+        batch = {"labels": labels, **input_ids}
+
+        prompt_input_ids = [{"input_ids": feature["prompt_input_ids"]} for feature in features]
+        prompt_input_ids = self.prompt_tokenizer.pad(
+            prompt_input_ids,
+            return_tensors="pt",
+            padding=self.padding,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            max_length=self.prompt_max_length,
+        )
+
+        batch["prompt_input_ids"] = prompt_input_ids["input_ids"]
+        if "attention_mask" in prompt_input_ids:
+            batch["prompt_attention_mask"] = prompt_input_ids["attention_mask"]
+
+        return batch
+
+
+def convert_dataset_str_to_list(
+    dataset_names,
+    splits=None,
+    dataset_samples=None,
+    default_split="train",
+):
+    if isinstance(dataset_names, str):
+        dataset_names = dataset_names.split("+")
+        splits = splits.split("+") if splits is not None else None
+        dataset_samples = dataset_samples.split("+") if dataset_samples is not None else None
+
+    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 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)
+
+    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,
+                "split": splits[i],
+                "samples": dataset_samples[i],
+            }
+        )
+    return dataset_names_dict
+
+
+def load_multiple_datasets(
+    accelerator: Accelerator,
+    dataset_names: Union[List, str],
+    splits: Optional[Union[List, str]] = None,
+    label_column_names: Optional[List] = None,
+    stopping_strategy: Optional[str] = "first_exhausted",
+    dataset_samples: Optional[Union[List, np.array]] = None,
+    streaming: Optional[bool] = False,
+    seed: Optional[int] = None,
+    id_column_name: Optional[str] = None,
+    columns_to_keep: Optional[Set[str]] = None,
+    prompt_column_name: Optional[str] = None,
+    sampling_rate: Optional[int] = None,
+    audio_column_name: Optional[str] = None,
+    logger: Optional[logging.Logger] = None,
+    librittsr_dir: Optional[Union[List, str]] = None,
+    other_dir: Optional[Union[List, str]] = None,
+    **kwargs,
+) -> Union[Dataset, IterableDataset]:
+    dataset_names_dict = convert_dataset_str_to_list(
+        dataset_names, splits, label_column_names, dataset_samples
+    )
+
+    if dataset_samples is not None:
+        dataset_samples = [ds_dict["samples"] for ds_dict in dataset_names_dict]
+        probabilities = np.array(dataset_samples) / np.sum(dataset_samples)
+    else:
+        probabilities = None
+
+    all_datasets = []
+    # iterate over the datasets we want to interleave
+    for dataset_dict in tqdm(dataset_names_dict, desc="Combining datasets..."):
+        with accelerator.local_main_process_first():
+            dataset = load_dataset(
+                dataset_dict["name"],
+                split=dataset_dict["split"],
+                streaming=streaming,
+                **kwargs,
+            )
+            dataset_features = dataset.features.keys()
+
+            if columns_to_keep is not None:
+                dataset = dataset.remove_columns(set(dataset_features - columns_to_keep))
+
+            def resolve_path(example):
+                path = example["audio_path"]
+                source = example["source"]
+
+                if source == "libritts-r":
+                    full_path = os.path.join(librittsr_dir, path)
+                else:
+                    full_path = os.path.join(other_dir, path)
+
+                return os.path.exists(full_path)
+                
+            dataset = dataset.filter(resolve_path, num_proc=16)
+
+        all_datasets.append(dataset)
+
+    if len(all_datasets) == 1:
+        # we have a single dataset so just return it as is
+        return all_datasets[0]
+
+    if streaming:
+        interleaved_dataset = interleave_datasets(
+            all_datasets,
+            stopping_strategy=stopping_strategy,
+            probabilities=probabilities,
+            seed=seed,
+        )
+    else:
+        with accelerator.local_main_process_first():
+            interleaved_dataset = concatenate_datasets(all_datasets)
+
+    return interleaved_dataset

capspeech/ar/training/data_capttsse.py

@@ -0,0 +1,253 @@
+import logging
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Set, Union
+import os
+import datasets
+import numpy as np
+import torch
+from accelerate import Accelerator
+from datasets import Dataset, IterableDataset, concatenate_datasets, interleave_datasets, load_dataset
+from tqdm import tqdm
+from transformers import AutoFeatureExtractor, AutoTokenizer
+import torchaudio
+import torchaudio.transforms as T
+
+@dataclass
+class DataCollatorEncodecWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received to the longest sequence in the batch or
+    to `max_length` if `max_length` is set and `padding=max_length`.
+    """
+
+    feature_extractor: AutoFeatureExtractor
+    audio_column_name: str
+    librittsrmix_dir: Optional[str] = None
+    feature_extractor_input_name: Optional[str] = "input_values"
+    max_length: Optional[int] = None
+    padding: Optional[str] = "longest"
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need
+        # different padding methods
+        sampling_rate = self.feature_extractor.sampling_rate
+        # load audio
+        audios = []
+        for f in features:
+            path = f[self.audio_column_name]
+            source = f["source"]
+            if source == "libritts-r":
+                path = os.path.join(self.librittsrmix_dir, path)
+            else:
+                raise ValueError(source)
+
+            if os.path.exists(path):
+                waveform, sr = torchaudio.load(path)
+                if sr != sampling_rate:
+                    resampler = T.Resample(orig_freq=sr, new_freq=sampling_rate)
+                    waveform = resampler(waveform)
+                if waveform.shape[0] > 1:
+                    waveform = waveform.mean(dim=0, keepdim=True)
+                audios.append(waveform.squeeze())
+            else:
+                print(f"Read error: {path}")
+
+
+        len_audio = [len(audio) for audio in audios]
+        if self.max_length is not None:
+            audios = [audio[: min(l, self.max_length)] for audio, l in zip(audios, len_audio)]
+
+        # since resampling has already been performed in the 'load_multiple_datasets' function,
+        # a fixed sampling_rate(44100hz) is passed to the feature_extractor.
+        batch = self.feature_extractor(
+            [np.asarray(a, dtype=np.float32) for a in audios], sampling_rate=sampling_rate, return_tensors="pt", padding=self.padding, max_length=self.max_length
+        )
+        batch["len_audio"] = torch.tensor(len_audio).unsqueeze(1)
+        return batch
+
+
+@dataclass
+class DataCollatorParlerTTSWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received.
+    Args:
+        prompt_tokenizer (:class:`~transformers.AutoTokenizer`)
+            The prompt_tokenizer used for proccessing the data.
+        description_tokenizer (:class:`~transformers.AutoTokenizer`)
+            The description_tokenizer used for proccessing the data.
+        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
+            among:
+            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+              sequence if provided).
+            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
+              maximum acceptable input length for the model if that argument is not provided.
+            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
+              different lengths).
+        pad_to_multiple_of (:obj:`int`, `optional`):
+            If set will pad the sequence to a multiple of the provided value.
+            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
+            7.5 (Volta).
+    """
+
+    prompt_tokenizer: AutoTokenizer
+    description_tokenizer: AutoTokenizer
+    padding: Union[bool, str] = "longest"
+    pad_to_multiple_of: Optional[int] = None
+    prompt_max_length: Optional[int] = None
+    description_max_length: Optional[int] = None
+    audio_max_length: Optional[int] = None
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need
+        # different padding methods
+
+        labels = [torch.tensor(feature["labels"]).transpose(0, 1) for feature in features]
+        # (bsz, seq_len, num_codebooks)
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=-100)
+        if self.audio_max_length is not None and self.padding == "max_length":
+            labels = torch.nn.functional.pad(
+                labels, pad=(0, 0, 0, max(self.audio_max_length - labels.shape[1], 0)), value=-100
+            )
+
+        input_ids = [{"input_ids": feature["input_ids"]} for feature in features]
+
+        input_ids = self.description_tokenizer.pad(
+            input_ids,
+            return_tensors="pt",
+            padding=self.padding,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            max_length=self.description_max_length,
+        )
+
+        batch = {"labels": labels, **input_ids}
+
+        prompt_input_ids = [{"input_ids": feature["prompt_input_ids"]} for feature in features]
+        prompt_input_ids = self.prompt_tokenizer.pad(
+            prompt_input_ids,
+            return_tensors="pt",
+            padding=self.padding,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            max_length=self.prompt_max_length,
+        )
+
+        batch["prompt_input_ids"] = prompt_input_ids["input_ids"]
+        if "attention_mask" in prompt_input_ids:
+            batch["prompt_attention_mask"] = prompt_input_ids["attention_mask"]
+
+        return batch
+
+
+def convert_dataset_str_to_list(
+    dataset_names,
+    splits=None,
+    dataset_samples=None,
+    default_split="train",
+):
+    if isinstance(dataset_names, str):
+        dataset_names = dataset_names.split("+")
+        splits = splits.split("+") if splits is not None else None
+        dataset_samples = dataset_samples.split("+") if dataset_samples is not None else None
+
+    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 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)
+
+    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,
+                "split": splits[i],
+                "samples": dataset_samples[i],
+            }
+        )
+    return dataset_names_dict
+
+
+def load_multiple_datasets(
+    accelerator: Accelerator,
+    dataset_names: Union[List, str],
+    splits: Optional[Union[List, str]] = None,
+    label_column_names: Optional[List] = None,
+    stopping_strategy: Optional[str] = "first_exhausted",
+    dataset_samples: Optional[Union[List, np.array]] = None,
+    streaming: Optional[bool] = False,
+    seed: Optional[int] = None,
+    id_column_name: Optional[str] = None,
+    columns_to_keep: Optional[Set[str]] = None,
+    prompt_column_name: Optional[str] = None,
+    sampling_rate: Optional[int] = None,
+    audio_column_name: Optional[str] = None,
+    logger: Optional[logging.Logger] = None,
+    librittsrmix_dir: Optional[Union[List, str]] = None,
+    **kwargs,
+) -> Union[Dataset, IterableDataset]:
+    dataset_names_dict = convert_dataset_str_to_list(
+        dataset_names, splits, label_column_names, dataset_samples
+    )
+
+    if dataset_samples is not None:
+        dataset_samples = [ds_dict["samples"] for ds_dict in dataset_names_dict]
+        probabilities = np.array(dataset_samples) / np.sum(dataset_samples)
+    else:
+        probabilities = None
+
+    all_datasets = []
+    # iterate over the datasets we want to interleave
+    for dataset_dict in tqdm(dataset_names_dict, desc="Combining datasets..."):
+        with accelerator.local_main_process_first():
+            dataset = load_dataset(
+                dataset_dict["name"],
+                split=dataset_dict["split"],
+                streaming=streaming,
+                **kwargs,
+            )
+            dataset_features = dataset.features.keys()
+
+            if columns_to_keep is not None:
+                dataset = dataset.remove_columns(set(dataset_features - columns_to_keep))
+
+            def resolve_path(example):
+                path = example["audio_path"]
+                source = example["source"]
+
+                if source == "libritts-r":
+                    full_path = os.path.join(librittsrmix_dir, path)
+                else:
+                    return False  # unknown source
+
+                return os.path.exists(full_path)
+                
+            dataset = dataset.filter(resolve_path, num_proc=16)
+
+        all_datasets.append(dataset)
+
+    if len(all_datasets) == 1:
+        # we have a single dataset so just return it as is
+        return all_datasets[0]
+
+    if streaming:
+        interleaved_dataset = interleave_datasets(
+            all_datasets,
+            stopping_strategy=stopping_strategy,
+            probabilities=probabilities,
+            seed=seed,
+        )
+    else:
+        with accelerator.local_main_process_first():
+            interleaved_dataset = concatenate_datasets(all_datasets)
+
+    return interleaved_dataset

capspeech/ar/training/finetune_captts.py

@@ -0,0 +1,1270 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+""" Train Parler-TTS using 🤗 Accelerate"""
+
+import logging
+import os
+import re
+import sys
+import time
+import math
+import contextlib
+from multiprocess import set_start_method
+from datetime import timedelta
+import inspect
+from tqdm import tqdm
+from pathlib import Path
+import wandb
+
+import torch
+from torch.utils.data import DataLoader
+
+import datasets
+from datasets import DatasetDict, Dataset, IterableDataset, concatenate_datasets
+
+from huggingface_hub import HfApi
+
+import transformers
+from transformers import AutoFeatureExtractor, AutoTokenizer, HfArgumentParser
+from transformers.trainer_pt_utils import LengthGroupedSampler
+from transformers.optimization import get_scheduler
+from transformers.utils import send_example_telemetry
+
+
+from accelerate import Accelerator, skip_first_batches
+from accelerate.utils import set_seed, AutocastKwargs, InitProcessGroupKwargs, TorchDynamoPlugin, DistributedDataParallelKwargs
+from accelerate.utils.memory import release_memory
+
+from parler_tts import (
+    ParlerTTSConfig,
+    ParlerTTSForConditionalGeneration,
+    build_delay_pattern_mask,
+)
+
+from training.utils import (
+    get_last_checkpoint,
+    rotate_checkpoints,
+    log_pred,
+    log_metric,
+    load_all_codec_checkpoints,
+    save_codec_checkpoint,
+    get_last_codec_checkpoint_step,
+)
+from training.arguments_captts import ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments
+from training.data_captts import load_multiple_datasets, DataCollatorParlerTTSWithPadding, DataCollatorEncodecWithPadding
+from training.eval import clap_similarity, wer, si_sdr
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+    # See all possible arguments in src/transformers/training_args.py
+    # or by passing the --help flag to this script.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments))
+    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.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
+    # information sent is the one passed as arguments along with your Python/PyTorch versions.
+    send_example_telemetry("run_parler_tts", model_args, data_args)
+    
+    if data_args.wandb_key is not None:
+        wandb.login(key=data_args.wandb_key)
+
+    if training_args.dtype == "float16":
+        mixed_precision = "fp16"
+        torch_dtype = torch.float16
+    elif training_args.dtype == "bfloat16":
+        mixed_precision = "bf16"
+        torch_dtype = torch.bfloat16
+    else:
+        mixed_precision = "no"
+        torch_dtype = torch.float32
+
+    if data_args.pad_to_max_length and (
+        data_args.max_duration_in_seconds is None
+        or data_args.max_prompt_token_length is None
+        or data_args.max_description_token_length is None
+    ):
+        raise ValueError(
+            "`pad_to_max_length` is `True` but one of the following parameters has not been set: `max_duration_in_seconds`, `max_prompt_token_length`, `max_description_token_length`"
+        )
+
+    padding = "max_length" if data_args.pad_to_max_length else "longest"
+
+    ####### A. Preparation
+    kwargs_handlers = [InitProcessGroupKwargs(timeout=timedelta(minutes=120)), DistributedDataParallelKwargs(find_unused_parameters=False)]
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
+        mixed_precision=mixed_precision,
+        log_with=training_args.report_to,
+        project_dir=training_args.output_dir,
+        kwargs_handlers=kwargs_handlers,
+    )
+
+    accelerator.init_trackers(
+        project_name=data_args.wandb_project,
+        config={
+            "learning_rate": training_args.learning_rate,
+            "model_name_or_path": model_args.model_name_or_path,
+            "num_train_epochs": training_args.num_train_epochs,
+            "gradient_accumulation_steps": training_args.gradient_accumulation_steps,
+            "per_device_train_batch_size": training_args.per_device_train_batch_size,
+            "global_batch_size": training_args.per_device_train_batch_size * accelerator.num_processes,
+            "mixed_precision": mixed_precision,
+            "lr_scheduler_type": training_args.lr_scheduler_type,
+            "warmup_steps": training_args.warmup_steps,
+            "freeze_text_encoder": model_args.freeze_text_encoder,
+            "max_duration_in_seconds": data_args.max_duration_in_seconds,
+            "weight_decay": training_args.weight_decay,
+            "adam_beta1": training_args.adam_beta1,
+            "adam_beta2": training_args.adam_beta2,
+            "temperature": model_args.temperature,
+        },
+        init_kwargs={"wandb": {"name": data_args.wandb_run_name}} if data_args.wandb_run_name else {},
+    )
+
+    # Detecting last checkpoint and eventually continue from last checkpoint
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+
+    # Setup logging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+    logger.setLevel(logging.INFO if accelerator.is_main_process else logging.WARN)
+
+    # Log a small summary on each proces
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
+        f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
+    )
+
+    # Set the verbosity to info of the Transformers logger (on main process only)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+    num_workers = data_args.preprocessing_num_workers
+
+    # 1. First, lett's instantiate the feature extractor, tokenizers and model
+    # Note for distributed training, the .from_pretrained methods guarantee that only
+    # one local process can concurrently download model & vocab.
+
+    # load feature extractor
+    feature_extractor = AutoFeatureExtractor.from_pretrained(
+        model_args.feature_extractor_name or model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+    ) 
+    sampling_rate = feature_extractor.sampling_rate
+
+    # load prompt tokenizer
+    prompt_tokenizer = AutoTokenizer.from_pretrained(
+        model_args.prompt_tokenizer_name or model_args.description_tokenizer_name or model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+        use_fast=model_args.use_fast_tokenizer,
+        padding_side=model_args.prompt_padding_side,
+    )
+
+    # load description tokenizer
+    description_tokenizer = AutoTokenizer.from_pretrained(
+        model_args.description_tokenizer_name or model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+        use_fast=model_args.use_fast_tokenizer,
+    )
+
+    if model_args.use_fast_tokenizer:
+        logger.warning(
+            "Disabling fast tokenizer warning: https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3231-L3235"
+        )
+        prompt_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
+        description_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
+
+    # 2. Now, let's load the dataset
+
+    if data_args.save_to_disk is not None:
+        os.makedirs(data_args.save_to_disk, exist_ok=True)
+
+    # assume that the dataset has been saved to `save_to_disk` if the latter is not empty
+    dataset_was_precomputed = len(os.listdir(data_args.save_to_disk)) > 0
+    if dataset_was_precomputed:
+        with accelerator.local_main_process_first():
+            vectorized_datasets = datasets.load_from_disk(data_args.save_to_disk)
+    else:
+        raw_datasets = DatasetDict()
+
+        columns_to_keep = {
+            "target_audio_column_name": data_args.target_audio_column_name,
+            "prompt_column_name": data_args.prompt_column_name,
+            "source": data_args.source_column_name,
+        }
+        if data_args.description_column_name is not None:
+            columns_to_keep["description_column_name"] = data_args.description_column_name
+
+        if training_args.do_train:
+            raw_datasets["train"] = load_multiple_datasets(
+                accelerator,
+                data_args.train_dataset_name,
+                splits=data_args.train_split_name,
+                dataset_samples=data_args.train_dataset_samples,
+                seed=training_args.seed,
+                cache_dir=model_args.cache_dir,
+                num_proc=data_args.preprocessing_num_workers,
+                id_column_name=data_args.id_column_name,
+                columns_to_keep=columns_to_keep.values(),
+                prompt_column_name=data_args.prompt_column_name,
+                audio_column_name=data_args.target_audio_column_name,
+                sampling_rate=sampling_rate,
+                logger=logger,
+                librittsr_dir=data_args.librittsr_dir,
+                other_dir=data_args.other_dir,
+                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
+            )
+
+            for key in columns_to_keep:
+                if columns_to_keep[key] not in raw_datasets["train"].column_names:
+                    raise ValueError(
+                        f"--{key} '{columns_to_keep[key]}' not found in dataset '{data_args.train_dataset_name}'."
+                        f" Make sure to set `--{key}` to the correct audio column - one of"
+                        f" {', '.join(raw_datasets['train'].column_names)}."
+                    )
+
+            if data_args.max_train_samples is not None:
+                raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))
+
+        if training_args.do_eval:
+            raw_datasets["eval"] = load_multiple_datasets(
+                accelerator,
+                data_args.eval_dataset_name if data_args.eval_dataset_name else data_args.train_dataset_name,
+                splits=data_args.eval_split_name,
+                cache_dir=model_args.cache_dir,
+                num_proc=data_args.preprocessing_num_workers,
+                id_column_name=data_args.id_column_name,
+                columns_to_keep=columns_to_keep.values(),
+                prompt_column_name=data_args.prompt_column_name,
+                audio_column_name=data_args.target_audio_column_name,
+                sampling_rate=sampling_rate,
+                logger=logger,
+                librittsr_dir=data_args.librittsr_dir,
+                other_dir=data_args.other_dir,
+                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
+            )
+
+            if data_args.max_eval_samples is not None:
+                with accelerator.local_main_process_first():
+                    raw_datasets["eval"] = (
+                        raw_datasets["eval"].shuffle(seed=training_args.seed).select(range(data_args.max_eval_samples))
+                    )
+
+    # 3. Next, let's load the config.
+    config = ParlerTTSConfig.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+    )
+ 
+    if training_args.codebook_weights is not None and len(training_args.codebook_weights) != config.decoder.num_codebooks:
+        raise ValueError(f"`codebook_weights` has length {len(training_args.codebook_weights)} when it should be of length {config.decoder.num_codebooks}.")
+
+    # update pad token id and decoder_start_token_id
+    config.decoder.update(
+        {
+            "cross_attention_implementation_strategy": model_args.cross_attention_implementation_strategy
+            if model_args.cross_attention_implementation_strategy is not None
+            else None,
+            "codebook_weights": training_args.codebook_weights if training_args.codebook_weights is not None else config.decoder.codebook_weights
+        }
+    )
+    config.update(
+        {
+            "pad_token_id": model_args.pad_token_id if model_args.pad_token_id is not None else config.pad_token_id,
+            "decoder_start_token_id": model_args.decoder_start_token_id
+            if model_args.decoder_start_token_id is not None
+            else config.decoder_start_token_id,
+        }
+    )
+
+    with open("events.txt", "r") as f:
+        events = [line.strip() for line in f]
+    events = ["<"+event.lower().replace(" ", "_")+">" for event in events]
+    events.append("<B_start>")
+    events.append("<B_end>")
+    events.append("<I_start>")
+    events.append("<I_end>")
+
+    special_tokens = {"additional_special_tokens": events}
+    prompt_tokenizer.add_special_tokens(special_tokens)
+    description_tokenizer.add_special_tokens(special_tokens)
+    padded_vocab_size = ((len(prompt_tokenizer) + 127) // 128) * 128 
+    config.vocab_size = padded_vocab_size 
+
+    # create model
+    model = ParlerTTSForConditionalGeneration.from_pretrained(
+        model_args.model_name_or_path,
+        ignore_mismatched_sizes=True,
+        cache_dir=model_args.cache_dir,
+        config=config,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+        attn_implementation={"decoder": model_args.attn_implementation, "text_encoder": "eager"},
+    )
+    model.text_encoder.resize_token_embeddings(padded_vocab_size)
+
+    # enable gradient checkpointing if necessary
+    if training_args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+
+    # 4. Now we preprocess the datasets including loading the audio, resampling and normalization
+    # Thankfully, `datasets` takes care of automatically loading and resampling the audio,
+    # so that we just need to set the correct target sampling rate and normalize the input
+    # via the `feature_extractor`
+
+    # derive max & min input length for sample rate & max duration
+    sampling_rate = feature_extractor.sampling_rate
+    max_target_length = int(data_args.max_duration_in_seconds * sampling_rate)
+    min_target_length = int(data_args.min_duration_in_seconds * sampling_rate)
+    target_audio_column_name = data_args.target_audio_column_name
+    description_column_name = data_args.description_column_name
+    prompt_column_name = data_args.prompt_column_name
+    feature_extractor_input_name = feature_extractor.model_input_names[0]
+    audio_encoder_pad_token_id = config.decoder.pad_token_id
+    audio_encoder_eos_token_id = config.decoder.eos_token_id
+    audio_encoder_bos_token_id = model.generation_config.decoder_start_token_id
+    max_length = model.generation_config.max_length
+    num_codebooks = model.decoder.config.num_codebooks
+    bandwidth = model_args.bandwidth
+    attn_implementation = model_args.attn_implementation
+
+    # Freeze Encoders
+    model.freeze_encoders(model_args.freeze_text_encoder)
+
+    # Test all gather - used for warmout and avoiding timeout
+    logger.debug(str(accelerator.process_index), main_process_only=False, in_order=True)
+    test_tensor = torch.tensor([accelerator.process_index], device=accelerator.device)
+    gathered_tensor = accelerator.gather(test_tensor)
+    print("gathered_tensor", gathered_tensor)
+    accelerator.wait_for_everyone()
+
+    if not dataset_was_precomputed:
+        # Filter on text length
+        if description_column_name is not None and data_args.max_text_length is not None:
+            with accelerator.local_main_process_first():
+                # filter description that is shorter than max_text_length
+                raw_datasets = raw_datasets.filter(
+                    lambda x: len(x) < data_args.max_text_length,
+                    num_proc=num_workers,
+                    input_columns=[description_column_name],
+                )
+
+        # Preprocessing the dataset.
+        # We need to tokenize the texts.
+        def pass_through_processors(description, prompt):
+            batch = {}
+
+            batch["input_ids"] = description_tokenizer(description.strip())["input_ids"]
+            batch["prompt_input_ids"] = prompt_tokenizer(prompt.strip())["input_ids"]
+
+            return batch
+
+        with accelerator.local_main_process_first():
+            # this is a trick to avoid to rewrite the entire audio column which takes ages
+            vectorized_datasets = raw_datasets.map(
+                pass_through_processors,
+                remove_columns=next(iter(raw_datasets.values())).column_names,
+                input_columns=[description_column_name, prompt_column_name],
+                num_proc=num_workers,
+                desc="preprocess datasets",
+            )
+
+        # We use Accelerate to perform distributed inference
+        # T5 doesn't support fp16
+        autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))
+
+        # Now we encode the audio labels with encodec.
+        ####### B. Encode audio
+
+        logger.info("*** Encode target audio with encodec ***")
+
+        # no need to prepare audio_decoder because used for inference without mixed precision
+        # see: https://huggingface.co/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare
+        if training_args.torch_compile:
+            audio_decoder = accelerator.prepare_model(model.audio_encoder, evaluation_mode=True)
+        else:
+            audio_decoder = model.audio_encoder
+
+        encoder_data_collator = DataCollatorEncodecWithPadding(
+            feature_extractor,
+            audio_column_name=target_audio_column_name,
+            librittsr_dir=data_args.librittsr_dir,
+            other_dir=data_args.other_dir,
+            feature_extractor_input_name=feature_extractor_input_name,
+            max_length=max_target_length,
+            padding=padding,
+        )
+        encoder_signature = set(inspect.signature(audio_decoder.forward).parameters)
+
+        def apply_audio_decoder(batch):
+            len_audio = batch.pop("len_audio")
+            audio_decoder.to(batch["input_values"].device).eval()
+            if bandwidth is not None:
+                batch["bandwidth"] = bandwidth
+            elif "num_quantizers" in encoder_signature:
+                batch["num_quantizers"] = num_codebooks
+            elif "num_codebooks" in encoder_signature:
+                batch["num_codebooks"] = num_codebooks
+            elif "n_quantizers" in encoder_signature:
+                batch["n_quantizers"] = num_codebooks
+
+            with torch.no_grad():
+                labels = audio_decoder.encode(**batch)["audio_codes"]
+            output = {}
+            output["len_audio"] = len_audio
+            # (1, bsz, codebooks, seq_len) -> (bsz, seq_len, codebooks)
+            output["labels"] = labels.squeeze(0).transpose(1, 2)
+
+            # if `pad_to_max_length`, the maximum corresponding audio length of the current batch is max_duration*sampling_rate
+            max_length = len_audio.max() if padding != "max_length" else max_target_length
+            output["ratio"] = torch.ones_like(len_audio) * labels.shape[-1] / max_length
+            return output
+
+        # (1, codebooks, seq_len) where seq_len=1
+        bos_labels = torch.ones((1, num_codebooks, 1)) * audio_encoder_bos_token_id
+
+        def postprocess_dataset(labels):
+            # (1, codebooks, seq_len)
+            labels = torch.tensor(labels).unsqueeze(0)
+            # add bos
+            labels = torch.cat([bos_labels, labels], dim=-1)
+
+            labels, delay_pattern_mask = build_delay_pattern_mask(
+                labels,
+                bos_token_id=audio_encoder_bos_token_id,
+                pad_token_id=audio_encoder_eos_token_id,
+                max_length=labels.shape[-1] + num_codebooks,
+                num_codebooks=num_codebooks,
+            )
+
+            # the first ids of the delay pattern mask are precisely labels, we use the rest of the labels mask
+            # to take care of EOS
+            # we want labels to look like this:
+            #  - [B, a, b, E, E, E, E]
+            #  - [B, B, c, d, E, E, E]
+            #  - [B, B, B, e, f, E, E]
+            #  - [B, B, B, B, g, h, E]
+            labels = torch.where(delay_pattern_mask == -1, audio_encoder_eos_token_id, delay_pattern_mask)
+
+            # the first timestamp is associated to a row full of BOS, let's get rid of it
+            # we also remove the last timestampts (full of PAD)
+            output = {"labels": labels[:, 1:]}
+            return output
+
+        for split in vectorized_datasets:
+            data_loader = DataLoader(
+                raw_datasets[split],
+                batch_size=training_args.audio_encoder_per_device_batch_size,
+                collate_fn=encoder_data_collator,
+                num_workers=training_args.dataloader_num_workers,
+                pin_memory=True,
+            )
+            data_loader = accelerator.prepare(data_loader)
+            total_inference_steps = len(data_loader)
+
+            start_step = get_last_codec_checkpoint_step(os.path.join(data_args.temporary_save_to_disk, split))
+            accelerator.wait_for_everyone()
+            if start_step > 0:
+                logger.info(f"Resuming {split} from step {start_step}")
+                # efficiently skip the first n batches
+                start_step += 1
+                data_loader = skip_first_batches(data_loader, start_step)
+
+            all_generated_labels = []
+            all_lens = []
+            if start_step < total_inference_steps:
+                for i, batch in enumerate(tqdm(data_loader, disable=not accelerator.is_local_main_process)):
+                    cur_step = start_step + i
+                    generate_labels = apply_audio_decoder(batch)
+                    generate_labels = accelerator.pad_across_processes(generate_labels, dim=1, pad_index=0)
+                    generate_labels = accelerator.gather_for_metrics(generate_labels)
+
+                    if accelerator.is_main_process:
+                        lab = generate_labels["labels"].cpu().transpose(1, 2).to(torch.int16)
+                        rat = generate_labels["ratio"].cpu().squeeze(1)
+                        lens = generate_labels["len_audio"].cpu().squeeze(1)
+                        lab = [l[:, : int(ratio * length)] for (l, ratio, length) in zip(lab, rat, lens)]
+
+                        all_generated_labels.extend(lab)
+                        all_lens.extend(lens)
+
+                        if ((cur_step + 1) % data_args.save_codec_steps == 0) or (
+                            cur_step == total_inference_steps - 1
+                        ):
+                            tmp_labels = Dataset.from_dict({"labels": all_generated_labels, "target_length": all_lens})
+                            tmp_labels = tmp_labels.map(
+                                postprocess_dataset,
+                                num_proc=data_args.preprocessing_num_workers,  # this one is resource consuming if many processor.
+                                input_columns=["labels"],
+                                desc="Postprocessing labeling",
+                            )
+                            save_codec_checkpoint(
+                                os.path.join(data_args.temporary_save_to_disk, split), tmp_labels, cur_step
+                            )
+                            all_generated_labels = []
+                            all_lens = []
+
+                accelerator.wait_for_everyone()
+
+            if accelerator.is_main_process and len(all_generated_labels) > 0:
+                tmp_labels = Dataset.from_dict({"labels": all_generated_labels, "target_length": all_lens})
+                tmp_labels = tmp_labels.map(
+                    postprocess_dataset,
+                    num_proc=data_args.preprocessing_num_workers,  # this one is resource consuming if many processor.
+                    input_columns=["labels"],
+                    desc="Postprocessing labeling",
+                )
+                save_codec_checkpoint(os.path.join(data_args.temporary_save_to_disk, split), tmp_labels, cur_step)
+                all_generated_labels = []
+                all_lens = []
+            accelerator.wait_for_everyone()
+
+            del all_generated_labels
+            accelerator.wait_for_everyone()
+
+            with accelerator.local_main_process_first():
+                tmp_labels = load_all_codec_checkpoints(os.path.join(data_args.temporary_save_to_disk, split)).select(
+                    range(len(vectorized_datasets[split]))
+                )
+                logger.info(f"Concatenating {split}: {tmp_labels} with {vectorized_datasets[split]}")
+                vectorized_datasets[split] = concatenate_datasets([vectorized_datasets[split], tmp_labels], axis=1)
+
+        accelerator.free_memory()
+        del generate_labels, all_lens
+
+        with accelerator.local_main_process_first():
+            # NOTE: filtering is done at the end because in the `datasets` library, caching audio files is done after most operations
+            # caching audio files is time and disk-space consuming, so we want to avoid it at all costs, especially for large (>1Kh) audio datasets.
+            # That's also why we avoid to concat the processed datasets (vectorized_datasets) with the audio column present in raw_datasets.
+
+            def is_audio_in_length_range(length):
+                return length > min_target_length and length < max_target_length
+
+            # filter data that is shorter than min_target_length
+            vectorized_datasets = vectorized_datasets.filter(
+                is_audio_in_length_range,
+                num_proc=num_workers,
+                input_columns=["target_length"],
+            )
+
+        if description_column_name is not None and data_args.max_description_token_length is not None:
+            with accelerator.local_main_process_first():
+                # filter description that is shorter than max_text_length
+                vectorized_datasets = vectorized_datasets.filter(
+                    lambda x: len(x) < data_args.max_description_token_length,
+                    num_proc=num_workers,
+                    input_columns=["input_ids"],
+                )
+
+        if data_args.max_prompt_token_length is not None:
+            with accelerator.local_main_process_first():
+                # filter description that is shorter than max_text_length
+                vectorized_datasets = vectorized_datasets.filter(
+                    lambda x: len(x) < data_args.max_prompt_token_length,
+                    num_proc=num_workers,
+                    input_columns=["prompt_input_ids"],
+                )
+
+    if data_args.save_to_disk is not None and not dataset_was_precomputed:
+        if accelerator.is_main_process:
+            vectorized_datasets.save_to_disk(
+                data_args.save_to_disk,
+                num_proc=min(data_args.preprocessing_num_workers, len(vectorized_datasets["eval"]) - 1),
+            )
+        accelerator.wait_for_everyone()
+        logger.info(f"Dataset saved at {data_args.save_to_disk}")
+
+    audio_max_length = None
+    if padding == "max_length":
+        audio_max_length = max(vectorized_datasets["train"]["target_length"])
+        with accelerator.local_main_process_first():
+            max_sample = vectorized_datasets["train"].filter(
+                lambda x: x == audio_max_length,
+                num_proc=num_workers,
+                input_columns=["target_length"],
+            )
+        audio_max_length = max([len(l[0]) for l in max_sample["labels"]])
+
+    if description_column_name is not None and data_args.max_description_token_length is not None:
+        with accelerator.local_main_process_first():
+            # filter description that is shorter than max_text_length
+            vectorized_datasets = vectorized_datasets.filter(
+                lambda x: len(x) < data_args.max_description_token_length,
+                num_proc=num_workers,
+                input_columns=["input_ids"],
+            )
+
+    if data_args.max_prompt_token_length is not None:
+        with accelerator.local_main_process_first():
+            # filter description that is shorter than max_text_length
+            vectorized_datasets = vectorized_datasets.filter(
+                lambda x: len(x) < data_args.max_prompt_token_length,
+                num_proc=num_workers,
+                input_columns=["prompt_input_ids"],
+            )
+
+    if training_args.group_by_length:
+        # apply a simple heuristic to take into account audio and text lengths
+        def add_target_lengths(target_length, prompt, description):
+            return {"target_length": target_length + len(prompt) + len(description)}
+
+        with accelerator.local_main_process_first():
+            vectorized_datasets = vectorized_datasets.map(
+                add_target_lengths,
+                num_proc=num_workers,
+                input_columns=["target_length", "prompt_input_ids", "input_ids"],
+            )
+
+    # for large datasets it is advised to run the preprocessing on a
+    # single machine first with ``args.preprocessing_only`` since there will mostly likely
+    # be a timeout when running the script in distributed mode.
+    # In a second step ``args.preprocessing_only`` can then be set to `False` to load the
+    # cached dataset
+    if data_args.preprocessing_only and data_args.save_to_disk is None:
+        raise ValueError(
+            "`preprocessing_only=True` but `save_to_disk` is not set. The latter should indicates where to save the dataset locally."
+        )
+    elif data_args.preprocessing_only:
+        logger.info(f"Data preprocessing finished. Files save at {data_args.save_to_disk}")
+        return
+
+    # 6. Next, we can prepare the training.
+
+    # Let's use word CLAP similary and WER metrics as our evaluation metrics,
+    def compute_metrics(
+        audios,
+        descriptions,
+        prompts,
+        device="cpu",
+        compute_clap_similarity_metric=False,
+        compute_noise_level_metric=False,
+        noise_level_to_compute_clean_wer=None,
+    ):
+        results = {}
+        input_ids = descriptions
+        texts = description_tokenizer.batch_decode(input_ids, skip_special_tokens=True)
+        prompts = prompt_tokenizer.batch_decode(prompts, skip_special_tokens=True)
+        audios = [a.float().cpu().numpy() for a in audios]
+
+        if compute_clap_similarity_metric:
+            clap_score = clap_similarity(
+                model_args.clap_model_name_or_path, texts, audios, device, input_sampling_rate=sampling_rate
+            )
+            results["clap"] = clap_score
+
+        si_sdr_measures = None
+        if compute_noise_level_metric:
+            si_sdr_measures = si_sdr(audios, device, input_sampling_rate=sampling_rate)
+
+        word_error, transcriptions, clean_word_error, noisy_word_error, percent_clean_samples = wer(
+            model_args.asr_model_name_or_path,
+            prompts,
+            audios,
+            device,
+            training_args.per_device_eval_batch_size,
+            sampling_rate,
+            noise_level_to_compute_clean_wer,
+            si_sdr_measures,
+        )
+        results["wer"] = word_error
+        if clean_word_error is not None:
+            results["clean_wer"] = clean_word_error
+            results["noisy_word_error"] = noisy_word_error
+            results["percent_clean_samples"] = percent_clean_samples
+
+        return results, texts, prompts, audios, transcriptions, si_sdr_measures
+
+    # Define Training Schedule
+    # Store some constants
+    per_device_train_batch_size = int(training_args.per_device_train_batch_size)
+    train_batch_size = per_device_train_batch_size * accelerator.num_processes
+    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
+    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
+
+    if training_args.max_steps < 0:
+        num_epochs = int(training_args.num_train_epochs)
+        steps_per_epoch = len(vectorized_datasets["train"]) // (train_batch_size * gradient_accumulation_steps)
+        total_train_steps = steps_per_epoch * num_epochs
+    elif training_args.max_steps > 0:
+        logger.info("max_steps is given, it will override any value given in num_train_epochs")
+        total_train_steps = int(training_args.max_steps)
+        # Setting a very large number of epochs so we go as many times as necessary over the iterator.
+        num_epochs = sys.maxsize
+        steps_per_epoch = total_train_steps
+
+    if training_args.eval_steps is None:
+        logger.info(f"eval_steps is not set, evaluating at the end of each epoch")
+        eval_steps = steps_per_epoch
+    else:
+        eval_steps = training_args.eval_steps
+        
+    if training_args.eval_generation_steps is None:
+        eval_generation_steps = eval_steps
+    else:
+        eval_generation_steps = training_args.eval_generation_steps
+
+    # T5 doesn't support fp16
+    autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))
+
+    # Define optimizer, LR scheduler, collator
+    optimizer = torch.optim.AdamW(
+        params=model.parameters(),
+        lr=training_args.learning_rate,
+        betas=(training_args.adam_beta1, training_args.adam_beta2),
+        eps=training_args.adam_epsilon,
+        weight_decay=training_args.weight_decay,
+    )
+
+    # LR scheduler gets stepped by `num_processes` each time -> account for this in warmup / total steps
+    lr_scheduler = get_scheduler(
+        name=training_args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=training_args.get_warmup_steps(total_train_steps) * accelerator.num_processes,
+        num_training_steps=total_train_steps * accelerator.num_processes,
+    )
+
+    # Instantiate custom data collator
+    data_collator = DataCollatorParlerTTSWithPadding(
+        prompt_tokenizer=prompt_tokenizer,
+        description_tokenizer=description_tokenizer,
+        pad_to_multiple_of=data_args.pad_to_multiple_of,
+        padding=padding,
+        prompt_max_length=data_args.max_prompt_token_length,
+        description_max_length=data_args.max_description_token_length,
+        audio_max_length=audio_max_length,
+    )
+
+    # Prepare everything with accelerate
+    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
+
+    num_examples = total_train_steps * train_batch_size * gradient_accumulation_steps
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {num_examples}")
+    logger.info("  Instantaneous batch size per device =" f" {per_device_train_batch_size}")
+    logger.info("  Gradient accumulation steps =" f" {gradient_accumulation_steps}")
+    logger.info(
+        f"  Total train batch size (w. parallel & distributed) = {train_batch_size * gradient_accumulation_steps}"
+    )
+    logger.info(f"  Total optimization steps = {total_train_steps}")
+
+    # ======================== Training ================================
+    train_time = 0
+    train_start = time.time()
+    steps_trained_progress_bar = tqdm(
+        range(total_train_steps), desc="Train steps ... ", position=0, disable=not accelerator.is_local_main_process
+    )
+    continue_training = True
+    epochs_trained = 0
+    cur_step = 0
+
+    checkpoint = None
+    if training_args.resume_from_checkpoint is not None:
+        checkpoint = training_args.resume_from_checkpoint
+    elif last_checkpoint is not None:
+        checkpoint = last_checkpoint
+
+    if accelerator.is_main_process:
+        if training_args.push_to_hub:
+            api = HfApi(token=training_args.hub_token)
+
+            # Create repo (repo_name from args or inferred)
+            repo_name = training_args.hub_model_id
+            if repo_name is None:
+                repo_name = Path(training_args.output_dir).absolute().name
+            repo_id = api.create_repo(repo_name, exist_ok=True).repo_id
+
+            with open(os.path.join(training_args.output_dir, ".gitignore"), "w+") as gitignore:
+                if "wandb" not in gitignore:
+                    gitignore.write("wandb\n")
+        elif training_args.output_dir is not None:
+            os.makedirs(training_args.output_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    # Now save everything to be able to create a single processor later
+    # make sure all processes wait until data is saved
+    # only the main process saves them
+    if accelerator.is_main_process:
+        # save feature extractor, tokenizer and config
+        if (
+            model_args.prompt_tokenizer_name is None
+            and model_args.description_tokenizer_name
+            or (model_args.prompt_tokenizer_name == model_args.description_tokenizer_name)
+        ):
+            prompt_tokenizer.save_pretrained(training_args.output_dir)
+        else:
+            logger.warning(
+                f"Prompt tokenizer ('{model_args.prompt_tokenizer_name}') and description tokenizer ('{model_args.description_tokenizer_name}') are not the same. Saving only the prompt tokenizer."
+            )
+            prompt_tokenizer.save_pretrained(training_args.output_dir)
+
+        feature_extractor.save_pretrained(training_args.output_dir)
+        config.save_pretrained(training_args.output_dir)
+    accelerator.wait_for_everyone()
+
+    if checkpoint is not None:
+        accelerator.load_state(checkpoint)
+        # Find num steps and epoch from saved state string pattern
+        pattern = r"checkpoint-(\d+)-epoch-(\d+)"
+        match = re.search(pattern, checkpoint)
+        cur_step = int(match.group(1))
+        epochs_trained = int(match.group(2))
+
+        logger.info("  Continuing training from checkpoint, will skip to saved global_step")
+        logger.info(f"  Continuing training from epoch {epochs_trained}")
+        logger.info(f"  Continuing training from global step {cur_step}")
+
+        steps_trained_progress_bar.update(cur_step)
+
+        for epoch in range(0, epochs_trained):
+            with accelerator.local_main_process_first():
+                vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+
+        if training_args.max_steps < 0:
+            # we know exactly the number of steps per epoch, so can skip through the required number of batches
+            resume_step = (cur_step - epochs_trained * steps_per_epoch) * gradient_accumulation_steps
+        else:
+            # Currently we don't know how many steps we've taken in the current epoch
+            # So we just shuffle the dataset one extra time and start from a fresh epoch
+            # This is "good enough" for our purposes but not fully correct
+            resume_step = None
+            with accelerator.local_main_process_first():
+                vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+    else:
+        resume_step = None
+
+    gen_kwargs = {
+        "do_sample": model_args.do_sample,
+        "temperature": model_args.temperature,
+        "max_length": model_args.max_length,
+        # Because of the delayed pattern mask, generation might stop earlier because of unexpected behaviour
+        # on the first tokens of the codebooks that are delayed.
+        # This fix the issue.
+        "min_new_tokens": num_codebooks + 1,
+    }
+
+    # Define gradient update step fn
+    def train_step(
+        batch,
+        accelerator,
+        autocast_kwargs,
+        num_items_in_batch,
+        gradient_accumulation_steps,
+    ):
+        if mixed_precision == "fp16":
+            # fp16 doesn't work with T5-like models
+            with accelerator.autocast(autocast_handler=autocast_kwargs):
+                if training_args.parallel_mode.value != "distributed":
+                    encoder_outputs = model.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                else:
+                    encoder_outputs = model.module.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                # we optionnally project last_hidden_state to avoid recomputing every time
+                encoder_hidden_states = encoder_outputs.last_hidden_state
+                if (
+                    config.text_encoder.hidden_size != config.decoder.hidden_size
+                    and config.decoder.cross_attention_hidden_size is None
+                ):
+                    encoder_hidden_states = (
+                        model.enc_to_dec_proj(encoder_hidden_states)
+                        if training_args.parallel_mode.value != "distributed"
+                        else model.module.enc_to_dec_proj(encoder_hidden_states)
+                    )
+
+                if batch.get("attention_mask", None) is not None:
+                    encoder_hidden_states = encoder_hidden_states * batch.get("attention_mask", None)[..., None]
+
+                encoder_outputs.last_hidden_state = encoder_hidden_states
+                batch["encoder_outputs"] = encoder_outputs
+
+        outputs = model(**batch, loss_reduction="sum")
+        # CE (data) loss
+        ce_loss = (outputs.loss * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch
+
+        metrics = {"loss": ce_loss}
+        
+        # per CE loss
+        per_codebook_losses = outputs.per_codebook_losses
+        metrics.update({f"codebook_{i}_loss": ((l  * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch) for (i,l) in enumerate(per_codebook_losses)})
+        return ce_loss, metrics
+
+    # Define eval fn
+    def eval_step(
+        batch,
+        accelerator,
+        autocast_kwargs,
+    ):
+        eval_model = model if not training_args.torch_compile else model._orig_mod
+
+        if mixed_precision == "fp16":
+            # fp16 doesn't work with T5-like models
+            with accelerator.autocast(autocast_handler=autocast_kwargs):
+                if training_args.parallel_mode.value != "distributed":
+                    encoder_outputs = model.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                else:
+                    encoder_outputs = model.module.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                # we optionnally project last_hidden_state to avoid recomputing every time
+                encoder_hidden_states = encoder_outputs.last_hidden_state
+                if (
+                    config.text_encoder.hidden_size != config.decoder.hidden_size
+                    and config.decoder.cross_attention_hidden_size is None
+                ):
+                    encoder_hidden_states = (
+                        model.enc_to_dec_proj(encoder_hidden_states)
+                        if training_args.parallel_mode.value != "distributed"
+                        else model.module.enc_to_dec_proj(encoder_hidden_states)
+                    )
+
+                if batch.get("attention_mask", None) is not None:
+                    encoder_hidden_states = encoder_hidden_states * batch.get("attention_mask", None)[..., None]
+
+                encoder_outputs.last_hidden_state = encoder_hidden_states
+                batch["encoder_outputs"] = encoder_outputs
+
+        with torch.no_grad():
+            outputs = eval_model(**batch)
+        # CE (data) loss
+        ce_loss = outputs.loss
+        metrics = {"loss": ce_loss}
+        
+        # per CE loss
+        per_codebook_losses = outputs.per_codebook_losses
+        metrics.update({f"codebook_{i}_loss": l for (i,l) in enumerate(per_codebook_losses)})
+        return metrics
+
+    def generate_step(batch, accelerator):
+        batch.pop("decoder_attention_mask", None)
+        eval_model = accelerator.unwrap_model(model, keep_fp32_wrapper=True)
+        if training_args.torch_compile:
+            # if the model is compiled, we use the original model bc compile is not compatible with .generate
+            eval_model = model._orig_mod
+
+        # since we've might have loaded the weights in fp32, we have to autocast to ensure FA2 weights are in half-precision.
+        # with accelerator.autocast(autocast_handler=AutocastKwargs(enabled=(attn_implementation=="flash_attention_2"))):
+        output_audios = eval_model.generate(**batch, **gen_kwargs)
+        output_audios = accelerator.pad_across_processes(output_audios, dim=1, pad_index=0)
+        return output_audios
+
+    model.train()
+
+    total_batched_samples = resume_step if resume_step is not None else 0
+    for epoch in range(epochs_trained, num_epochs):
+        with accelerator.local_main_process_first():
+            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+        sampler = None
+        if training_args.group_by_length:
+            sampler = LengthGroupedSampler(train_batch_size, lengths=vectorized_datasets["train"]["target_length"])
+        train_dataloader = DataLoader(
+            vectorized_datasets["train"],
+            collate_fn=data_collator,
+            batch_size=per_device_train_batch_size,
+            sampler=sampler,
+            shuffle=not training_args.group_by_length,
+            num_workers=training_args.dataloader_num_workers,
+            pin_memory=training_args.dataloader_pin_memory,
+        )
+        train_dataloader = accelerator.prepare(train_dataloader)
+        if hasattr(train_dataloader, "dataset") and isinstance(train_dataloader.dataset, IterableDataset):
+            train_dataloader.dataset.set_epoch(epoch)
+
+        if resume_step is not None:
+            # Skip the first N batches in the dataloader when resuming from a checkpoint
+            logger.info(f"  Skip first {resume_step} batches")
+            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
+            resume_step = None
+            accelerator.wait_for_everyone()
+
+        # We chunkify the epoch iterator into gradient accumulation steps `n` batches
+        train_iterator = iter(train_dataloader)
+        num_steps_in_epoch = len(train_dataloader)
+        remainder = num_steps_in_epoch % gradient_accumulation_steps
+        remainder = remainder if remainder != 0 else gradient_accumulation_steps
+        total_updates = math.ceil(num_steps_in_epoch / gradient_accumulation_steps)
+        
+        update_step = -1
+        for _ in range(total_updates):
+            update_step += 1
+            
+            # preload the total batch per step
+            batch_samples = []
+            num_batches_in_step = gradient_accumulation_steps if update_step != (total_updates - 1) else remainder
+            for _ in range(num_batches_in_step):
+                batch_samples += [next(train_iterator)]
+                
+            # get num items in batch - if different than BOS and than -100
+            num_items_in_batch = sum([(batch["labels"].ne(audio_encoder_bos_token_id) | batch["labels"].ne(-100) | batch["labels"].ne(audio_encoder_eos_token_id)).sum((0,1))[0] for batch in batch_samples])
+            num_items_in_batch = accelerator.gather(num_items_in_batch).sum().item()
+            
+            # losses = []
+            for i,batch in enumerate(batch_samples):
+                total_batched_samples += 1
+                ctx = model.no_sync if (i < len(batch_samples) - 1 and accelerator.num_processes > 1) else contextlib.nullcontext
+                
+                with ctx():
+                    loss, train_metric = train_step(batch, accelerator, autocast_kwargs, num_items_in_batch, gradient_accumulation_steps)
+                    accelerator.backward(loss)
+                    # losses.append(loss.detach())
+            
+            grad_norm = accelerator.clip_grad_norm_(model.parameters(), training_args.max_grad_norm)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+            # The accelerator has performed an optimization step behind the scenes
+            steps_trained_progress_bar.update(1)
+            cur_step += 1
+
+            # losses = accelerator.gather(sum(losses)).sum().item() / (accelerator.num_processes * gradient_accumulation_steps)
+            
+            if cur_step % training_args.logging_steps == 0:
+                steps_trained_progress_bar.write(
+                    f"Step... ({cur_step} / {total_train_steps} | Loss:"
+                    f" {train_metric['loss']}, Learning Rate:"
+                    f" {lr_scheduler.get_last_lr()[0]})"
+                )
+                train_metric["grad_norm"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm
+                log_metric(
+                    accelerator,
+                    metrics=train_metric,
+                    learning_rate=lr_scheduler.get_last_lr()[0],
+                    train_time=train_time + time.time() - train_start,
+                    step=cur_step,
+                    epoch=epoch,
+                    prefix="train",
+                )
+
+            # save checkpoint and weights after each save_steps and at the end of training
+            if (cur_step % training_args.save_steps == 0) or cur_step == total_train_steps:
+                intermediate_dir = os.path.join(training_args.output_dir, f"checkpoint-{cur_step}-epoch-{epoch}")
+                # safe_serialization=False to avoid shared tensors saving issue (TODO(YL): it's a temporary fix)
+                # https://github.com/huggingface/transformers/issues/27293#issuecomment-1872560074
+                accelerator.save_state(output_dir=intermediate_dir, safe_serialization=False)
+                accelerator.wait_for_everyone()
+                if accelerator.is_main_process:
+                    rotate_checkpoints(
+                        training_args.save_total_limit, output_dir=training_args.output_dir, logger=logger
+                    )
+
+                    if cur_step == total_train_steps:
+                        # un-wrap student model for save
+                        unwrapped_model = accelerator.unwrap_model(model)
+                        unwrapped_model.save_pretrained(training_args.output_dir)
+
+                    if training_args.push_to_hub:
+                        api.upload_folder(
+                            repo_id=repo_id,
+                            folder_path=training_args.output_dir,
+                            commit_message=f"Saving train state of step {cur_step}",
+                            run_as_future=True,
+                        )
+                accelerator.wait_for_everyone()
+
+            if training_args.do_eval and (cur_step % eval_steps == 0 or cur_step == total_train_steps):
+                train_time += time.time() - train_start
+                # ======================== Evaluating ==============================
+                model.eval()
+                eval_metrics = []
+                eval_preds = []
+                eval_descriptions = []
+                eval_prompts = []
+                eval_start = time.time()
+
+                # release training input batch
+                batch = release_memory(batch)
+
+                validation_dataloader = DataLoader(
+                    vectorized_datasets["eval"],
+                    collate_fn=data_collator,
+                    batch_size=per_device_eval_batch_size,
+                    drop_last=False,
+                    num_workers=training_args.eval_dataloader_num_workers,
+                    pin_memory=training_args.dataloader_pin_memory,
+                )
+                validation_dataloader = accelerator.prepare(validation_dataloader)
+
+                for batch in tqdm(
+                    validation_dataloader,
+                    desc=f"Evaluating - Inference ...",
+                    position=2,
+                    disable=not accelerator.is_local_main_process,
+                ):
+                    # Model forward
+                    eval_metric = eval_step(batch, accelerator, autocast_kwargs)
+                    eval_metric = accelerator.gather_for_metrics(eval_metric)
+                    eval_metric = {key: val.unsqueeze(0) if val.ndim == 0 else val for (key,val) in eval_metric.items()}
+                    eval_metrics.append(eval_metric)
+
+                if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
+                    validation_dataloader = DataLoader(
+                        vectorized_datasets["eval"],
+                        collate_fn=data_collator,
+                        batch_size=per_device_eval_batch_size,
+                        drop_last=False,
+                        num_workers=training_args.eval_dataloader_num_workers,
+                        pin_memory=training_args.dataloader_pin_memory,
+                    )
+                    validation_dataloader = accelerator.prepare(validation_dataloader)
+                    # generation
+                    for batch in tqdm(
+                        validation_dataloader,
+                        desc=f"Evaluating - Generation ...",
+                        position=2,
+                        disable=not accelerator.is_local_main_process,
+                    ):
+                        generated_audios = generate_step(batch, accelerator)
+                        # Gather all predictions and targets
+                        generated_audios, input_ids, prompts = accelerator.pad_across_processes(
+                            (generated_audios, batch["input_ids"], batch["prompt_input_ids"]), dim=1, pad_index=0
+                        )
+                        generated_audios, input_ids, prompts = accelerator.gather_for_metrics(
+                            (generated_audios, input_ids, prompts)
+                        )
+                        eval_preds.extend(generated_audios.to("cpu"))
+                        eval_descriptions.extend(input_ids.to("cpu"))
+                        eval_prompts.extend(prompts.to("cpu"))
+
+                eval_time = time.time() - eval_start
+                # normalize eval metrics
+                eval_metrics = {
+                    key: torch.mean(torch.cat([d[key] for d in eval_metrics])).to("cpu") for key in eval_metrics[0]
+                }
+
+                # compute metrics
+                metrics_desc = ""
+                if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
+                    if accelerator.is_local_main_process:
+                        (
+                            metric_values,
+                            pred_descriptions,
+                            pred_prompts,
+                            audios,
+                            transcriptions,
+                            si_sdr_measures,
+                        ) = compute_metrics(
+                            eval_preds,
+                            eval_descriptions,
+                            eval_prompts,
+                            accelerator.device,
+                            training_args.compute_clap_similarity_metric,
+                            training_args.compute_noise_level_metric,
+                            training_args.noise_level_to_compute_clean_wer,
+                        )
+                        eval_metrics.update(metric_values)
+                        metrics_desc = " ".join([f"Eval {key}: {value} |" for key, value in metric_values.items()])
+                        if "wandb" in training_args.report_to:
+                            log_pred(
+                                accelerator,
+                                pred_descriptions,
+                                pred_prompts,
+                                transcriptions,
+                                audios,
+                                si_sdr_measures,
+                                sampling_rate=sampling_rate,
+                                step=cur_step,
+                                prefix="eval",
+                            )
+                    accelerator.wait_for_everyone()
+
+                # Print metrics and update progress bar
+                if accelerator.is_local_main_process:
+                    steps_trained_progress_bar.write(
+                        f"Eval results for step ({cur_step} / {total_train_steps} | Eval Loss: {eval_metrics['loss']} |"
+                        f" {metrics_desc})"
+                    )
+
+                log_metric(
+                    accelerator,
+                    metrics=eval_metrics,
+                    train_time=eval_time,
+                    step=cur_step,
+                    epoch=epoch,
+                    prefix="eval",
+                )
+
+                # release eval batch and relax metrics
+                eval_metrics, eval_preds, eval_descriptions, eval_prompts, batch, eval_metric = release_memory(
+                    eval_metrics, eval_preds, eval_descriptions, eval_prompts, batch, eval_metric
+                )
+                if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
+                    generated_audios, input_ids, prompts = release_memory(generated_audios, input_ids, prompts)
+
+                # train mode
+                model.train()
+
+                # flush the train metrics
+                train_start = time.time()
+
+            # break condition
+            if cur_step == total_train_steps:
+                continue_training = False
+                break
+
+        if not continue_training:
+            break
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main()

capspeech/ar/training/finetune_capttsse.py

@@ -0,0 +1,1267 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+""" Train Parler-TTS using 🤗 Accelerate"""
+
+import logging
+import os
+import re
+import sys
+import time
+import math
+import contextlib
+from multiprocess import set_start_method
+from datetime import timedelta
+import inspect
+from tqdm import tqdm
+from pathlib import Path
+import wandb
+
+import torch
+from torch.utils.data import DataLoader
+
+import datasets
+from datasets import DatasetDict, Dataset, IterableDataset, concatenate_datasets
+
+from huggingface_hub import HfApi
+
+import transformers
+from transformers import AutoFeatureExtractor, AutoTokenizer, HfArgumentParser
+from transformers.trainer_pt_utils import LengthGroupedSampler
+from transformers.optimization import get_scheduler
+from transformers.utils import send_example_telemetry
+
+
+from accelerate import Accelerator, skip_first_batches
+from accelerate.utils import set_seed, AutocastKwargs, InitProcessGroupKwargs, TorchDynamoPlugin, DistributedDataParallelKwargs
+from accelerate.utils.memory import release_memory
+
+from parler_tts import (
+    ParlerTTSConfig,
+    ParlerTTSForConditionalGeneration,
+    build_delay_pattern_mask,
+)
+
+from training.utils import (
+    get_last_checkpoint,
+    rotate_checkpoints,
+    log_pred,
+    log_metric,
+    load_all_codec_checkpoints,
+    save_codec_checkpoint,
+    get_last_codec_checkpoint_step,
+)
+from training.arguments_capttsse import ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments
+from training.data_capttsse import load_multiple_datasets, DataCollatorParlerTTSWithPadding, DataCollatorEncodecWithPadding
+from training.eval import clap_similarity, wer, si_sdr
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+    # See all possible arguments in src/transformers/training_args.py
+    # or by passing the --help flag to this script.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments))
+    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.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
+    # information sent is the one passed as arguments along with your Python/PyTorch versions.
+    send_example_telemetry("run_parler_tts", model_args, data_args)
+    
+    if data_args.wandb_key is not None:
+        wandb.login(key=data_args.wandb_key)
+
+    if training_args.dtype == "float16":
+        mixed_precision = "fp16"
+        torch_dtype = torch.float16
+    elif training_args.dtype == "bfloat16":
+        mixed_precision = "bf16"
+        torch_dtype = torch.bfloat16
+    else:
+        mixed_precision = "no"
+        torch_dtype = torch.float32
+
+    if data_args.pad_to_max_length and (
+        data_args.max_duration_in_seconds is None
+        or data_args.max_prompt_token_length is None
+        or data_args.max_description_token_length is None
+    ):
+        raise ValueError(
+            "`pad_to_max_length` is `True` but one of the following parameters has not been set: `max_duration_in_seconds`, `max_prompt_token_length`, `max_description_token_length`"
+        )
+
+    padding = "max_length" if data_args.pad_to_max_length else "longest"
+
+    ####### A. Preparation
+    kwargs_handlers = [InitProcessGroupKwargs(timeout=timedelta(minutes=120)), DistributedDataParallelKwargs(find_unused_parameters=False)]
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
+        mixed_precision=mixed_precision,
+        log_with=training_args.report_to,
+        project_dir=training_args.output_dir,
+        kwargs_handlers=kwargs_handlers,
+    )
+
+    accelerator.init_trackers(
+        project_name=data_args.wandb_project,
+        config={
+            "learning_rate": training_args.learning_rate,
+            "model_name_or_path": model_args.model_name_or_path,
+            "num_train_epochs": training_args.num_train_epochs,
+            "gradient_accumulation_steps": training_args.gradient_accumulation_steps,
+            "per_device_train_batch_size": training_args.per_device_train_batch_size,
+            "global_batch_size": training_args.per_device_train_batch_size * accelerator.num_processes,
+            "mixed_precision": mixed_precision,
+            "lr_scheduler_type": training_args.lr_scheduler_type,
+            "warmup_steps": training_args.warmup_steps,
+            "freeze_text_encoder": model_args.freeze_text_encoder,
+            "max_duration_in_seconds": data_args.max_duration_in_seconds,
+            "weight_decay": training_args.weight_decay,
+            "adam_beta1": training_args.adam_beta1,
+            "adam_beta2": training_args.adam_beta2,
+            "temperature": model_args.temperature,
+        },
+        init_kwargs={"wandb": {"name": data_args.wandb_run_name}} if data_args.wandb_run_name else {},
+    )
+
+    # Detecting last checkpoint and eventually continue from last checkpoint
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+
+    # Setup logging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+    logger.setLevel(logging.INFO if accelerator.is_main_process else logging.WARN)
+
+    # Log a small summary on each proces
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
+        f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
+    )
+
+    # Set the verbosity to info of the Transformers logger (on main process only)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+    num_workers = data_args.preprocessing_num_workers
+
+    # 1. First, lett's instantiate the feature extractor, tokenizers and model
+    # Note for distributed training, the .from_pretrained methods guarantee that only
+    # one local process can concurrently download model & vocab.
+
+    # load feature extractor
+    feature_extractor = AutoFeatureExtractor.from_pretrained(
+        model_args.feature_extractor_name or model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+    ) 
+    sampling_rate = feature_extractor.sampling_rate
+
+    # load prompt tokenizer
+    prompt_tokenizer = AutoTokenizer.from_pretrained(
+        model_args.prompt_tokenizer_name or model_args.description_tokenizer_name or model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+        use_fast=model_args.use_fast_tokenizer,
+        padding_side=model_args.prompt_padding_side,
+    )
+
+    # load description tokenizer
+    description_tokenizer = AutoTokenizer.from_pretrained(
+        model_args.description_tokenizer_name or model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+        use_fast=model_args.use_fast_tokenizer,
+    )
+
+    if model_args.use_fast_tokenizer:
+        logger.warning(
+            "Disabling fast tokenizer warning: https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3231-L3235"
+        )
+        prompt_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
+        description_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
+
+    # 2. Now, let's load the dataset
+
+    if data_args.save_to_disk is not None:
+        os.makedirs(data_args.save_to_disk, exist_ok=True)
+
+    # assume that the dataset has been saved to `save_to_disk` if the latter is not empty
+    dataset_was_precomputed = len(os.listdir(data_args.save_to_disk)) > 0
+    if dataset_was_precomputed:
+        with accelerator.local_main_process_first():
+            vectorized_datasets = datasets.load_from_disk(data_args.save_to_disk)
+    else:
+        raw_datasets = DatasetDict()
+
+        columns_to_keep = {
+            "target_audio_column_name": data_args.target_audio_column_name,
+            "prompt_column_name": data_args.prompt_column_name,
+            "source": data_args.source_column_name,
+        }
+        if data_args.description_column_name is not None:
+            columns_to_keep["description_column_name"] = data_args.description_column_name
+
+        if training_args.do_train:
+            raw_datasets["train"] = load_multiple_datasets(
+                accelerator,
+                data_args.train_dataset_name,
+                splits=data_args.train_split_name,
+                dataset_samples=data_args.train_dataset_samples,
+                seed=training_args.seed,
+                cache_dir=model_args.cache_dir,
+                num_proc=data_args.preprocessing_num_workers,
+                id_column_name=data_args.id_column_name,
+                columns_to_keep=columns_to_keep.values(),
+                prompt_column_name=data_args.prompt_column_name,
+                audio_column_name=data_args.target_audio_column_name,
+                sampling_rate=sampling_rate,
+                logger=logger,
+                librittsrmix_dir=data_args.librittsrmix_dir,
+                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
+            )
+
+            for key in columns_to_keep:
+                if columns_to_keep[key] not in raw_datasets["train"].column_names:
+                    raise ValueError(
+                        f"--{key} '{columns_to_keep[key]}' not found in dataset '{data_args.train_dataset_name}'."
+                        f" Make sure to set `--{key}` to the correct audio column - one of"
+                        f" {', '.join(raw_datasets['train'].column_names)}."
+                    )
+
+            if data_args.max_train_samples is not None:
+                raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))
+
+        if training_args.do_eval:
+            raw_datasets["eval"] = load_multiple_datasets(
+                accelerator,
+                data_args.eval_dataset_name if data_args.eval_dataset_name else data_args.train_dataset_name,
+                splits=data_args.eval_split_name,
+                cache_dir=model_args.cache_dir,
+                num_proc=data_args.preprocessing_num_workers,
+                id_column_name=data_args.id_column_name,
+                columns_to_keep=columns_to_keep.values(),
+                prompt_column_name=data_args.prompt_column_name,
+                audio_column_name=data_args.target_audio_column_name,
+                sampling_rate=sampling_rate,
+                logger=logger,
+                librittsrmix_dir=data_args.librittsrmix_dir,
+                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
+            )
+
+            if data_args.max_eval_samples is not None:
+                with accelerator.local_main_process_first():
+                    raw_datasets["eval"] = (
+                        raw_datasets["eval"].shuffle(seed=training_args.seed).select(range(data_args.max_eval_samples))
+                    )
+
+    # 3. Next, let's load the config.
+    config = ParlerTTSConfig.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+    )
+ 
+    if training_args.codebook_weights is not None and len(training_args.codebook_weights) != config.decoder.num_codebooks:
+        raise ValueError(f"`codebook_weights` has length {len(training_args.codebook_weights)} when it should be of length {config.decoder.num_codebooks}.")
+
+    # update pad token id and decoder_start_token_id
+    config.decoder.update(
+        {
+            "cross_attention_implementation_strategy": model_args.cross_attention_implementation_strategy
+            if model_args.cross_attention_implementation_strategy is not None
+            else None,
+            "codebook_weights": training_args.codebook_weights if training_args.codebook_weights is not None else config.decoder.codebook_weights
+        }
+    )
+    config.update(
+        {
+            "pad_token_id": model_args.pad_token_id if model_args.pad_token_id is not None else config.pad_token_id,
+            "decoder_start_token_id": model_args.decoder_start_token_id
+            if model_args.decoder_start_token_id is not None
+            else config.decoder_start_token_id,
+        }
+    )
+
+    with open("events.txt", "r") as f:
+        events = [line.strip() for line in f]
+    events = ["<"+event.lower().replace(" ", "_")+">" for event in events]
+    events.append("<B_start>")
+    events.append("<B_end>")
+    events.append("<I_start>")
+    events.append("<I_end>")
+
+    special_tokens = {"additional_special_tokens": events}
+    prompt_tokenizer.add_special_tokens(special_tokens)
+    description_tokenizer.add_special_tokens(special_tokens)
+    padded_vocab_size = ((len(prompt_tokenizer) + 127) // 128) * 128 
+    config.vocab_size = padded_vocab_size 
+
+    # create model
+    model = ParlerTTSForConditionalGeneration.from_pretrained(
+        model_args.model_name_or_path,
+        ignore_mismatched_sizes=True,
+        cache_dir=model_args.cache_dir,
+        config=config,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+        attn_implementation={"decoder": model_args.attn_implementation, "text_encoder": "eager"},
+    )
+    model.text_encoder.resize_token_embeddings(padded_vocab_size)
+
+    # enable gradient checkpointing if necessary
+    if training_args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+
+    # 4. Now we preprocess the datasets including loading the audio, resampling and normalization
+    # Thankfully, `datasets` takes care of automatically loading and resampling the audio,
+    # so that we just need to set the correct target sampling rate and normalize the input
+    # via the `feature_extractor`
+
+    # derive max & min input length for sample rate & max duration
+    sampling_rate = feature_extractor.sampling_rate
+    max_target_length = int(data_args.max_duration_in_seconds * sampling_rate)
+    min_target_length = int(data_args.min_duration_in_seconds * sampling_rate)
+    target_audio_column_name = data_args.target_audio_column_name
+    description_column_name = data_args.description_column_name
+    prompt_column_name = data_args.prompt_column_name
+    feature_extractor_input_name = feature_extractor.model_input_names[0]
+    audio_encoder_pad_token_id = config.decoder.pad_token_id
+    audio_encoder_eos_token_id = config.decoder.eos_token_id
+    audio_encoder_bos_token_id = model.generation_config.decoder_start_token_id
+    max_length = model.generation_config.max_length
+    num_codebooks = model.decoder.config.num_codebooks
+    bandwidth = model_args.bandwidth
+    attn_implementation = model_args.attn_implementation
+
+    # Freeze Encoders
+    model.freeze_encoders(model_args.freeze_text_encoder)
+
+    # Test all gather - used for warmout and avoiding timeout
+    logger.debug(str(accelerator.process_index), main_process_only=False, in_order=True)
+    test_tensor = torch.tensor([accelerator.process_index], device=accelerator.device)
+    gathered_tensor = accelerator.gather(test_tensor)
+    print("gathered_tensor", gathered_tensor)
+    accelerator.wait_for_everyone()
+
+    if not dataset_was_precomputed:
+        # Filter on text length
+        if description_column_name is not None and data_args.max_text_length is not None:
+            with accelerator.local_main_process_first():
+                # filter description that is shorter than max_text_length
+                raw_datasets = raw_datasets.filter(
+                    lambda x: len(x) < data_args.max_text_length,
+                    num_proc=num_workers,
+                    input_columns=[description_column_name],
+                )
+
+        # Preprocessing the dataset.
+        # We need to tokenize the texts.
+        def pass_through_processors(description, prompt):
+            batch = {}
+
+            batch["input_ids"] = description_tokenizer(description.strip())["input_ids"]
+            batch["prompt_input_ids"] = prompt_tokenizer(prompt.strip())["input_ids"]
+
+            return batch
+
+        with accelerator.local_main_process_first():
+            # this is a trick to avoid to rewrite the entire audio column which takes ages
+            vectorized_datasets = raw_datasets.map(
+                pass_through_processors,
+                remove_columns=next(iter(raw_datasets.values())).column_names,
+                input_columns=[description_column_name, prompt_column_name],
+                num_proc=num_workers,
+                desc="preprocess datasets",
+            )
+
+        # We use Accelerate to perform distributed inference
+        # T5 doesn't support fp16
+        autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))
+
+        # Now we encode the audio labels with encodec.
+        ####### B. Encode audio
+
+        logger.info("*** Encode target audio with encodec ***")
+
+        # no need to prepare audio_decoder because used for inference without mixed precision
+        # see: https://huggingface.co/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare
+        if training_args.torch_compile:
+            audio_decoder = accelerator.prepare_model(model.audio_encoder, evaluation_mode=True)
+        else:
+            audio_decoder = model.audio_encoder
+
+        encoder_data_collator = DataCollatorEncodecWithPadding(
+            feature_extractor,
+            audio_column_name=target_audio_column_name,
+            librittsrmix_dir=data_args.librittsrmix_dir,
+            feature_extractor_input_name=feature_extractor_input_name,
+            max_length=max_target_length,
+            padding=padding,
+        )
+        encoder_signature = set(inspect.signature(audio_decoder.forward).parameters)
+
+        def apply_audio_decoder(batch):
+            len_audio = batch.pop("len_audio")
+            audio_decoder.to(batch["input_values"].device).eval()
+            if bandwidth is not None:
+                batch["bandwidth"] = bandwidth
+            elif "num_quantizers" in encoder_signature:
+                batch["num_quantizers"] = num_codebooks
+            elif "num_codebooks" in encoder_signature:
+                batch["num_codebooks"] = num_codebooks
+            elif "n_quantizers" in encoder_signature:
+                batch["n_quantizers"] = num_codebooks
+
+            with torch.no_grad():
+                labels = audio_decoder.encode(**batch)["audio_codes"]
+            output = {}
+            output["len_audio"] = len_audio
+            # (1, bsz, codebooks, seq_len) -> (bsz, seq_len, codebooks)
+            output["labels"] = labels.squeeze(0).transpose(1, 2)
+
+            # if `pad_to_max_length`, the maximum corresponding audio length of the current batch is max_duration*sampling_rate
+            max_length = len_audio.max() if padding != "max_length" else max_target_length
+            output["ratio"] = torch.ones_like(len_audio) * labels.shape[-1] / max_length
+            return output
+
+        # (1, codebooks, seq_len) where seq_len=1
+        bos_labels = torch.ones((1, num_codebooks, 1)) * audio_encoder_bos_token_id
+
+        def postprocess_dataset(labels):
+            # (1, codebooks, seq_len)
+            labels = torch.tensor(labels).unsqueeze(0)
+            # add bos
+            labels = torch.cat([bos_labels, labels], dim=-1)
+
+            labels, delay_pattern_mask = build_delay_pattern_mask(
+                labels,
+                bos_token_id=audio_encoder_bos_token_id,
+                pad_token_id=audio_encoder_eos_token_id,
+                max_length=labels.shape[-1] + num_codebooks,
+                num_codebooks=num_codebooks,
+            )
+
+            # the first ids of the delay pattern mask are precisely labels, we use the rest of the labels mask
+            # to take care of EOS
+            # we want labels to look like this:
+            #  - [B, a, b, E, E, E, E]
+            #  - [B, B, c, d, E, E, E]
+            #  - [B, B, B, e, f, E, E]
+            #  - [B, B, B, B, g, h, E]
+            labels = torch.where(delay_pattern_mask == -1, audio_encoder_eos_token_id, delay_pattern_mask)
+
+            # the first timestamp is associated to a row full of BOS, let's get rid of it
+            # we also remove the last timestampts (full of PAD)
+            output = {"labels": labels[:, 1:]}
+            return output
+
+        for split in vectorized_datasets:
+            data_loader = DataLoader(
+                raw_datasets[split],
+                batch_size=training_args.audio_encoder_per_device_batch_size,
+                collate_fn=encoder_data_collator,
+                num_workers=training_args.dataloader_num_workers,
+                pin_memory=True,
+            )
+            data_loader = accelerator.prepare(data_loader)
+            total_inference_steps = len(data_loader)
+
+            start_step = get_last_codec_checkpoint_step(os.path.join(data_args.temporary_save_to_disk, split))
+            accelerator.wait_for_everyone()
+            if start_step > 0:
+                logger.info(f"Resuming {split} from step {start_step}")
+                # efficiently skip the first n batches
+                start_step += 1
+                data_loader = skip_first_batches(data_loader, start_step)
+
+            all_generated_labels = []
+            all_lens = []
+            if start_step < total_inference_steps:
+                for i, batch in enumerate(tqdm(data_loader, disable=not accelerator.is_local_main_process)):
+                    cur_step = start_step + i
+                    generate_labels = apply_audio_decoder(batch)
+                    generate_labels = accelerator.pad_across_processes(generate_labels, dim=1, pad_index=0)
+                    generate_labels = accelerator.gather_for_metrics(generate_labels)
+
+                    if accelerator.is_main_process:
+                        lab = generate_labels["labels"].cpu().transpose(1, 2).to(torch.int16)
+                        rat = generate_labels["ratio"].cpu().squeeze(1)
+                        lens = generate_labels["len_audio"].cpu().squeeze(1)
+                        lab = [l[:, : int(ratio * length)] for (l, ratio, length) in zip(lab, rat, lens)]
+
+                        all_generated_labels.extend(lab)
+                        all_lens.extend(lens)
+
+                        if ((cur_step + 1) % data_args.save_codec_steps == 0) or (
+                            cur_step == total_inference_steps - 1
+                        ):
+                            tmp_labels = Dataset.from_dict({"labels": all_generated_labels, "target_length": all_lens})
+                            tmp_labels = tmp_labels.map(
+                                postprocess_dataset,
+                                num_proc=data_args.preprocessing_num_workers,  # this one is resource consuming if many processor.
+                                input_columns=["labels"],
+                                desc="Postprocessing labeling",
+                            )
+                            save_codec_checkpoint(
+                                os.path.join(data_args.temporary_save_to_disk, split), tmp_labels, cur_step
+                            )
+                            all_generated_labels = []
+                            all_lens = []
+
+                accelerator.wait_for_everyone()
+
+            if accelerator.is_main_process and len(all_generated_labels) > 0:
+                tmp_labels = Dataset.from_dict({"labels": all_generated_labels, "target_length": all_lens})
+                tmp_labels = tmp_labels.map(
+                    postprocess_dataset,
+                    num_proc=data_args.preprocessing_num_workers,  # this one is resource consuming if many processor.
+                    input_columns=["labels"],
+                    desc="Postprocessing labeling",
+                )
+                save_codec_checkpoint(os.path.join(data_args.temporary_save_to_disk, split), tmp_labels, cur_step)
+                all_generated_labels = []
+                all_lens = []
+            accelerator.wait_for_everyone()
+
+            del all_generated_labels
+            accelerator.wait_for_everyone()
+
+            with accelerator.local_main_process_first():
+                tmp_labels = load_all_codec_checkpoints(os.path.join(data_args.temporary_save_to_disk, split)).select(
+                    range(len(vectorized_datasets[split]))
+                )
+                logger.info(f"Concatenating {split}: {tmp_labels} with {vectorized_datasets[split]}")
+                vectorized_datasets[split] = concatenate_datasets([vectorized_datasets[split], tmp_labels], axis=1)
+
+        accelerator.free_memory()
+        del generate_labels, all_lens
+
+        with accelerator.local_main_process_first():
+            # NOTE: filtering is done at the end because in the `datasets` library, caching audio files is done after most operations
+            # caching audio files is time and disk-space consuming, so we want to avoid it at all costs, especially for large (>1Kh) audio datasets.
+            # That's also why we avoid to concat the processed datasets (vectorized_datasets) with the audio column present in raw_datasets.
+
+            def is_audio_in_length_range(length):
+                return length > min_target_length and length < max_target_length
+
+            # filter data that is shorter than min_target_length
+            vectorized_datasets = vectorized_datasets.filter(
+                is_audio_in_length_range,
+                num_proc=num_workers,
+                input_columns=["target_length"],
+            )
+
+        if description_column_name is not None and data_args.max_description_token_length is not None:
+            with accelerator.local_main_process_first():
+                # filter description that is shorter than max_text_length
+                vectorized_datasets = vectorized_datasets.filter(
+                    lambda x: len(x) < data_args.max_description_token_length,
+                    num_proc=num_workers,
+                    input_columns=["input_ids"],
+                )
+
+        if data_args.max_prompt_token_length is not None:
+            with accelerator.local_main_process_first():
+                # filter description that is shorter than max_text_length
+                vectorized_datasets = vectorized_datasets.filter(
+                    lambda x: len(x) < data_args.max_prompt_token_length,
+                    num_proc=num_workers,
+                    input_columns=["prompt_input_ids"],
+                )
+
+    if data_args.save_to_disk is not None and not dataset_was_precomputed:
+        if accelerator.is_main_process:
+            vectorized_datasets.save_to_disk(
+                data_args.save_to_disk,
+                num_proc=min(data_args.preprocessing_num_workers, len(vectorized_datasets["eval"]) - 1),
+            )
+        accelerator.wait_for_everyone()
+        logger.info(f"Dataset saved at {data_args.save_to_disk}")
+
+    audio_max_length = None
+    if padding == "max_length":
+        audio_max_length = max(vectorized_datasets["train"]["target_length"])
+        with accelerator.local_main_process_first():
+            max_sample = vectorized_datasets["train"].filter(
+                lambda x: x == audio_max_length,
+                num_proc=num_workers,
+                input_columns=["target_length"],
+            )
+        audio_max_length = max([len(l[0]) for l in max_sample["labels"]])
+
+    if description_column_name is not None and data_args.max_description_token_length is not None:
+        with accelerator.local_main_process_first():
+            # filter description that is shorter than max_text_length
+            vectorized_datasets = vectorized_datasets.filter(
+                lambda x: len(x) < data_args.max_description_token_length,
+                num_proc=num_workers,
+                input_columns=["input_ids"],
+            )
+
+    if data_args.max_prompt_token_length is not None:
+        with accelerator.local_main_process_first():
+            # filter description that is shorter than max_text_length
+            vectorized_datasets = vectorized_datasets.filter(
+                lambda x: len(x) < data_args.max_prompt_token_length,
+                num_proc=num_workers,
+                input_columns=["prompt_input_ids"],
+            )
+
+    if training_args.group_by_length:
+        # apply a simple heuristic to take into account audio and text lengths
+        def add_target_lengths(target_length, prompt, description):
+            return {"target_length": target_length + len(prompt) + len(description)}
+
+        with accelerator.local_main_process_first():
+            vectorized_datasets = vectorized_datasets.map(
+                add_target_lengths,
+                num_proc=num_workers,
+                input_columns=["target_length", "prompt_input_ids", "input_ids"],
+            )
+
+    # for large datasets it is advised to run the preprocessing on a
+    # single machine first with ``args.preprocessing_only`` since there will mostly likely
+    # be a timeout when running the script in distributed mode.
+    # In a second step ``args.preprocessing_only`` can then be set to `False` to load the
+    # cached dataset
+    if data_args.preprocessing_only and data_args.save_to_disk is None:
+        raise ValueError(
+            "`preprocessing_only=True` but `save_to_disk` is not set. The latter should indicates where to save the dataset locally."
+        )
+    elif data_args.preprocessing_only:
+        logger.info(f"Data preprocessing finished. Files save at {data_args.save_to_disk}")
+        return
+
+    # 6. Next, we can prepare the training.
+
+    # Let's use word CLAP similary and WER metrics as our evaluation metrics,
+    def compute_metrics(
+        audios,
+        descriptions,
+        prompts,
+        device="cpu",
+        compute_clap_similarity_metric=False,
+        compute_noise_level_metric=False,
+        noise_level_to_compute_clean_wer=None,
+    ):
+        results = {}
+        input_ids = descriptions
+        texts = description_tokenizer.batch_decode(input_ids, skip_special_tokens=True)
+        prompts = prompt_tokenizer.batch_decode(prompts, skip_special_tokens=True)
+        audios = [a.float().cpu().numpy() for a in audios]
+
+        if compute_clap_similarity_metric:
+            clap_score = clap_similarity(
+                model_args.clap_model_name_or_path, texts, audios, device, input_sampling_rate=sampling_rate
+            )
+            results["clap"] = clap_score
+
+        si_sdr_measures = None
+        if compute_noise_level_metric:
+            si_sdr_measures = si_sdr(audios, device, input_sampling_rate=sampling_rate)
+
+        word_error, transcriptions, clean_word_error, noisy_word_error, percent_clean_samples = wer(
+            model_args.asr_model_name_or_path,
+            prompts,
+            audios,
+            device,
+            training_args.per_device_eval_batch_size,
+            sampling_rate,
+            noise_level_to_compute_clean_wer,
+            si_sdr_measures,
+        )
+        results["wer"] = word_error
+        if clean_word_error is not None:
+            results["clean_wer"] = clean_word_error
+            results["noisy_word_error"] = noisy_word_error
+            results["percent_clean_samples"] = percent_clean_samples
+
+        return results, texts, prompts, audios, transcriptions, si_sdr_measures
+
+    # Define Training Schedule
+    # Store some constants
+    per_device_train_batch_size = int(training_args.per_device_train_batch_size)
+    train_batch_size = per_device_train_batch_size * accelerator.num_processes
+    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
+    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
+
+    if training_args.max_steps < 0:
+        num_epochs = int(training_args.num_train_epochs)
+        steps_per_epoch = len(vectorized_datasets["train"]) // (train_batch_size * gradient_accumulation_steps)
+        total_train_steps = steps_per_epoch * num_epochs
+    elif training_args.max_steps > 0:
+        logger.info("max_steps is given, it will override any value given in num_train_epochs")
+        total_train_steps = int(training_args.max_steps)
+        # Setting a very large number of epochs so we go as many times as necessary over the iterator.
+        num_epochs = sys.maxsize
+        steps_per_epoch = total_train_steps
+
+    if training_args.eval_steps is None:
+        logger.info(f"eval_steps is not set, evaluating at the end of each epoch")
+        eval_steps = steps_per_epoch
+    else:
+        eval_steps = training_args.eval_steps
+        
+    if training_args.eval_generation_steps is None:
+        eval_generation_steps = eval_steps
+    else:
+        eval_generation_steps = training_args.eval_generation_steps
+
+    # T5 doesn't support fp16
+    autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))
+
+    # Define optimizer, LR scheduler, collator
+    optimizer = torch.optim.AdamW(
+        params=model.parameters(),
+        lr=training_args.learning_rate,
+        betas=(training_args.adam_beta1, training_args.adam_beta2),
+        eps=training_args.adam_epsilon,
+        weight_decay=training_args.weight_decay,
+    )
+
+    # LR scheduler gets stepped by `num_processes` each time -> account for this in warmup / total steps
+    lr_scheduler = get_scheduler(
+        name=training_args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=training_args.get_warmup_steps(total_train_steps) * accelerator.num_processes,
+        num_training_steps=total_train_steps * accelerator.num_processes,
+    )
+
+    # Instantiate custom data collator
+    data_collator = DataCollatorParlerTTSWithPadding(
+        prompt_tokenizer=prompt_tokenizer,
+        description_tokenizer=description_tokenizer,
+        pad_to_multiple_of=data_args.pad_to_multiple_of,
+        padding=padding,
+        prompt_max_length=data_args.max_prompt_token_length,
+        description_max_length=data_args.max_description_token_length,
+        audio_max_length=audio_max_length,
+    )
+
+    # Prepare everything with accelerate
+    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
+
+    num_examples = total_train_steps * train_batch_size * gradient_accumulation_steps
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {num_examples}")
+    logger.info("  Instantaneous batch size per device =" f" {per_device_train_batch_size}")
+    logger.info("  Gradient accumulation steps =" f" {gradient_accumulation_steps}")
+    logger.info(
+        f"  Total train batch size (w. parallel & distributed) = {train_batch_size * gradient_accumulation_steps}"
+    )
+    logger.info(f"  Total optimization steps = {total_train_steps}")
+
+    # ======================== Training ================================
+    train_time = 0
+    train_start = time.time()
+    steps_trained_progress_bar = tqdm(
+        range(total_train_steps), desc="Train steps ... ", position=0, disable=not accelerator.is_local_main_process
+    )
+    continue_training = True
+    epochs_trained = 0
+    cur_step = 0
+
+    checkpoint = None
+    if training_args.resume_from_checkpoint is not None:
+        checkpoint = training_args.resume_from_checkpoint
+    elif last_checkpoint is not None:
+        checkpoint = last_checkpoint
+
+    if accelerator.is_main_process:
+        if training_args.push_to_hub:
+            api = HfApi(token=training_args.hub_token)
+
+            # Create repo (repo_name from args or inferred)
+            repo_name = training_args.hub_model_id
+            if repo_name is None:
+                repo_name = Path(training_args.output_dir).absolute().name
+            repo_id = api.create_repo(repo_name, exist_ok=True).repo_id
+
+            with open(os.path.join(training_args.output_dir, ".gitignore"), "w+") as gitignore:
+                if "wandb" not in gitignore:
+                    gitignore.write("wandb\n")
+        elif training_args.output_dir is not None:
+            os.makedirs(training_args.output_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    # Now save everything to be able to create a single processor later
+    # make sure all processes wait until data is saved
+    # only the main process saves them
+    if accelerator.is_main_process:
+        # save feature extractor, tokenizer and config
+        if (
+            model_args.prompt_tokenizer_name is None
+            and model_args.description_tokenizer_name
+            or (model_args.prompt_tokenizer_name == model_args.description_tokenizer_name)
+        ):
+            prompt_tokenizer.save_pretrained(training_args.output_dir)
+        else:
+            logger.warning(
+                f"Prompt tokenizer ('{model_args.prompt_tokenizer_name}') and description tokenizer ('{model_args.description_tokenizer_name}') are not the same. Saving only the prompt tokenizer."
+            )
+            prompt_tokenizer.save_pretrained(training_args.output_dir)
+
+        feature_extractor.save_pretrained(training_args.output_dir)
+        config.save_pretrained(training_args.output_dir)
+    accelerator.wait_for_everyone()
+
+    if checkpoint is not None:
+        accelerator.load_state(checkpoint)
+        # Find num steps and epoch from saved state string pattern
+        pattern = r"checkpoint-(\d+)-epoch-(\d+)"
+        match = re.search(pattern, checkpoint)
+        cur_step = int(match.group(1))
+        epochs_trained = int(match.group(2))
+
+        logger.info("  Continuing training from checkpoint, will skip to saved global_step")
+        logger.info(f"  Continuing training from epoch {epochs_trained}")
+        logger.info(f"  Continuing training from global step {cur_step}")
+
+        steps_trained_progress_bar.update(cur_step)
+
+        for epoch in range(0, epochs_trained):
+            with accelerator.local_main_process_first():
+                vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+
+        if training_args.max_steps < 0:
+            # we know exactly the number of steps per epoch, so can skip through the required number of batches
+            resume_step = (cur_step - epochs_trained * steps_per_epoch) * gradient_accumulation_steps
+        else:
+            # Currently we don't know how many steps we've taken in the current epoch
+            # So we just shuffle the dataset one extra time and start from a fresh epoch
+            # This is "good enough" for our purposes but not fully correct
+            resume_step = None
+            with accelerator.local_main_process_first():
+                vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+    else:
+        resume_step = None
+
+    gen_kwargs = {
+        "do_sample": model_args.do_sample,
+        "temperature": model_args.temperature,
+        "max_length": model_args.max_length,
+        # Because of the delayed pattern mask, generation might stop earlier because of unexpected behaviour
+        # on the first tokens of the codebooks that are delayed.
+        # This fix the issue.
+        "min_new_tokens": num_codebooks + 1,
+    }
+
+    # Define gradient update step fn
+    def train_step(
+        batch,
+        accelerator,
+        autocast_kwargs,
+        num_items_in_batch,
+        gradient_accumulation_steps,
+    ):
+        if mixed_precision == "fp16":
+            # fp16 doesn't work with T5-like models
+            with accelerator.autocast(autocast_handler=autocast_kwargs):
+                if training_args.parallel_mode.value != "distributed":
+                    encoder_outputs = model.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                else:
+                    encoder_outputs = model.module.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                # we optionnally project last_hidden_state to avoid recomputing every time
+                encoder_hidden_states = encoder_outputs.last_hidden_state
+                if (
+                    config.text_encoder.hidden_size != config.decoder.hidden_size
+                    and config.decoder.cross_attention_hidden_size is None
+                ):
+                    encoder_hidden_states = (
+                        model.enc_to_dec_proj(encoder_hidden_states)
+                        if training_args.parallel_mode.value != "distributed"
+                        else model.module.enc_to_dec_proj(encoder_hidden_states)
+                    )
+
+                if batch.get("attention_mask", None) is not None:
+                    encoder_hidden_states = encoder_hidden_states * batch.get("attention_mask", None)[..., None]
+
+                encoder_outputs.last_hidden_state = encoder_hidden_states
+                batch["encoder_outputs"] = encoder_outputs
+
+        outputs = model(**batch, loss_reduction="sum")
+        # CE (data) loss
+        ce_loss = (outputs.loss * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch
+
+        metrics = {"loss": ce_loss}
+        
+        # per CE loss
+        per_codebook_losses = outputs.per_codebook_losses
+        metrics.update({f"codebook_{i}_loss": ((l  * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch) for (i,l) in enumerate(per_codebook_losses)})
+        return ce_loss, metrics
+
+    # Define eval fn
+    def eval_step(
+        batch,
+        accelerator,
+        autocast_kwargs,
+    ):
+        eval_model = model if not training_args.torch_compile else model._orig_mod
+
+        if mixed_precision == "fp16":
+            # fp16 doesn't work with T5-like models
+            with accelerator.autocast(autocast_handler=autocast_kwargs):
+                if training_args.parallel_mode.value != "distributed":
+                    encoder_outputs = model.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                else:
+                    encoder_outputs = model.module.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                # we optionnally project last_hidden_state to avoid recomputing every time
+                encoder_hidden_states = encoder_outputs.last_hidden_state
+                if (
+                    config.text_encoder.hidden_size != config.decoder.hidden_size
+                    and config.decoder.cross_attention_hidden_size is None
+                ):
+                    encoder_hidden_states = (
+                        model.enc_to_dec_proj(encoder_hidden_states)
+                        if training_args.parallel_mode.value != "distributed"
+                        else model.module.enc_to_dec_proj(encoder_hidden_states)
+                    )
+
+                if batch.get("attention_mask", None) is not None:
+                    encoder_hidden_states = encoder_hidden_states * batch.get("attention_mask", None)[..., None]
+
+                encoder_outputs.last_hidden_state = encoder_hidden_states
+                batch["encoder_outputs"] = encoder_outputs
+
+        with torch.no_grad():
+            outputs = eval_model(**batch)
+        # CE (data) loss
+        ce_loss = outputs.loss
+        metrics = {"loss": ce_loss}
+        
+        # per CE loss
+        per_codebook_losses = outputs.per_codebook_losses
+        metrics.update({f"codebook_{i}_loss": l for (i,l) in enumerate(per_codebook_losses)})
+        return metrics
+
+    def generate_step(batch, accelerator):
+        batch.pop("decoder_attention_mask", None)
+        eval_model = accelerator.unwrap_model(model, keep_fp32_wrapper=True)
+        if training_args.torch_compile:
+            # if the model is compiled, we use the original model bc compile is not compatible with .generate
+            eval_model = model._orig_mod
+
+        # since we've might have loaded the weights in fp32, we have to autocast to ensure FA2 weights are in half-precision.
+        # with accelerator.autocast(autocast_handler=AutocastKwargs(enabled=(attn_implementation=="flash_attention_2"))):
+        output_audios = eval_model.generate(**batch, **gen_kwargs)
+        output_audios = accelerator.pad_across_processes(output_audios, dim=1, pad_index=0)
+        return output_audios
+
+    model.train()
+
+    total_batched_samples = resume_step if resume_step is not None else 0
+    for epoch in range(epochs_trained, num_epochs):
+        with accelerator.local_main_process_first():
+            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+        sampler = None
+        if training_args.group_by_length:
+            sampler = LengthGroupedSampler(train_batch_size, lengths=vectorized_datasets["train"]["target_length"])
+        train_dataloader = DataLoader(
+            vectorized_datasets["train"],
+            collate_fn=data_collator,
+            batch_size=per_device_train_batch_size,
+            sampler=sampler,
+            shuffle=not training_args.group_by_length,
+            num_workers=training_args.dataloader_num_workers,
+            pin_memory=training_args.dataloader_pin_memory,
+        )
+        train_dataloader = accelerator.prepare(train_dataloader)
+        if hasattr(train_dataloader, "dataset") and isinstance(train_dataloader.dataset, IterableDataset):
+            train_dataloader.dataset.set_epoch(epoch)
+
+        if resume_step is not None:
+            # Skip the first N batches in the dataloader when resuming from a checkpoint
+            logger.info(f"  Skip first {resume_step} batches")
+            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
+            resume_step = None
+            accelerator.wait_for_everyone()
+
+        # We chunkify the epoch iterator into gradient accumulation steps `n` batches
+        train_iterator = iter(train_dataloader)
+        num_steps_in_epoch = len(train_dataloader)
+        remainder = num_steps_in_epoch % gradient_accumulation_steps
+        remainder = remainder if remainder != 0 else gradient_accumulation_steps
+        total_updates = math.ceil(num_steps_in_epoch / gradient_accumulation_steps)
+        
+        update_step = -1
+        for _ in range(total_updates):
+            update_step += 1
+            
+            # preload the total batch per step
+            batch_samples = []
+            num_batches_in_step = gradient_accumulation_steps if update_step != (total_updates - 1) else remainder
+            for _ in range(num_batches_in_step):
+                batch_samples += [next(train_iterator)]
+                
+            # get num items in batch - if different than BOS and than -100
+            num_items_in_batch = sum([(batch["labels"].ne(audio_encoder_bos_token_id) | batch["labels"].ne(-100) | batch["labels"].ne(audio_encoder_eos_token_id)).sum((0,1))[0] for batch in batch_samples])
+            num_items_in_batch = accelerator.gather(num_items_in_batch).sum().item()
+            
+            # losses = []
+            for i,batch in enumerate(batch_samples):
+                total_batched_samples += 1
+                ctx = model.no_sync if (i < len(batch_samples) - 1 and accelerator.num_processes > 1) else contextlib.nullcontext
+                
+                with ctx():
+                    loss, train_metric = train_step(batch, accelerator, autocast_kwargs, num_items_in_batch, gradient_accumulation_steps)
+                    accelerator.backward(loss)
+                    # losses.append(loss.detach())
+            
+            grad_norm = accelerator.clip_grad_norm_(model.parameters(), training_args.max_grad_norm)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+            # The accelerator has performed an optimization step behind the scenes
+            steps_trained_progress_bar.update(1)
+            cur_step += 1
+
+            # losses = accelerator.gather(sum(losses)).sum().item() / (accelerator.num_processes * gradient_accumulation_steps)
+            
+            if cur_step % training_args.logging_steps == 0:
+                steps_trained_progress_bar.write(
+                    f"Step... ({cur_step} / {total_train_steps} | Loss:"
+                    f" {train_metric['loss']}, Learning Rate:"
+                    f" {lr_scheduler.get_last_lr()[0]})"
+                )
+                train_metric["grad_norm"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm
+                log_metric(
+                    accelerator,
+                    metrics=train_metric,
+                    learning_rate=lr_scheduler.get_last_lr()[0],
+                    train_time=train_time + time.time() - train_start,
+                    step=cur_step,
+                    epoch=epoch,
+                    prefix="train",
+                )
+
+            # save checkpoint and weights after each save_steps and at the end of training
+            if (cur_step % training_args.save_steps == 0) or cur_step == total_train_steps:
+                intermediate_dir = os.path.join(training_args.output_dir, f"checkpoint-{cur_step}-epoch-{epoch}")
+                # safe_serialization=False to avoid shared tensors saving issue (TODO(YL): it's a temporary fix)
+                # https://github.com/huggingface/transformers/issues/27293#issuecomment-1872560074
+                accelerator.save_state(output_dir=intermediate_dir, safe_serialization=False)
+                accelerator.wait_for_everyone()
+                if accelerator.is_main_process:
+                    rotate_checkpoints(
+                        training_args.save_total_limit, output_dir=training_args.output_dir, logger=logger
+                    )
+
+                    if cur_step == total_train_steps:
+                        # un-wrap student model for save
+                        unwrapped_model = accelerator.unwrap_model(model)
+                        unwrapped_model.save_pretrained(training_args.output_dir)
+
+                    if training_args.push_to_hub:
+                        api.upload_folder(
+                            repo_id=repo_id,
+                            folder_path=training_args.output_dir,
+                            commit_message=f"Saving train state of step {cur_step}",
+                            run_as_future=True,
+                        )
+                accelerator.wait_for_everyone()
+
+            if training_args.do_eval and (cur_step % eval_steps == 0 or cur_step == total_train_steps):
+                train_time += time.time() - train_start
+                # ======================== Evaluating ==============================
+                model.eval()
+                eval_metrics = []
+                eval_preds = []
+                eval_descriptions = []
+                eval_prompts = []
+                eval_start = time.time()
+
+                # release training input batch
+                batch = release_memory(batch)
+
+                validation_dataloader = DataLoader(
+                    vectorized_datasets["eval"],
+                    collate_fn=data_collator,
+                    batch_size=per_device_eval_batch_size,
+                    drop_last=False,
+                    num_workers=training_args.eval_dataloader_num_workers,
+                    pin_memory=training_args.dataloader_pin_memory,
+                )
+                validation_dataloader = accelerator.prepare(validation_dataloader)
+
+                for batch in tqdm(
+                    validation_dataloader,
+                    desc=f"Evaluating - Inference ...",
+                    position=2,
+                    disable=not accelerator.is_local_main_process,
+                ):
+                    # Model forward
+                    eval_metric = eval_step(batch, accelerator, autocast_kwargs)
+                    eval_metric = accelerator.gather_for_metrics(eval_metric)
+                    eval_metric = {key: val.unsqueeze(0) if val.ndim == 0 else val for (key,val) in eval_metric.items()}
+                    eval_metrics.append(eval_metric)
+
+                if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
+                    validation_dataloader = DataLoader(
+                        vectorized_datasets["eval"],
+                        collate_fn=data_collator,
+                        batch_size=per_device_eval_batch_size,
+                        drop_last=False,
+                        num_workers=training_args.eval_dataloader_num_workers,
+                        pin_memory=training_args.dataloader_pin_memory,
+                    )
+                    validation_dataloader = accelerator.prepare(validation_dataloader)
+                    # generation
+                    for batch in tqdm(
+                        validation_dataloader,
+                        desc=f"Evaluating - Generation ...",
+                        position=2,
+                        disable=not accelerator.is_local_main_process,
+                    ):
+                        generated_audios = generate_step(batch, accelerator)
+                        # Gather all predictions and targets
+                        generated_audios, input_ids, prompts = accelerator.pad_across_processes(
+                            (generated_audios, batch["input_ids"], batch["prompt_input_ids"]), dim=1, pad_index=0
+                        )
+                        generated_audios, input_ids, prompts = accelerator.gather_for_metrics(
+                            (generated_audios, input_ids, prompts)
+                        )
+                        eval_preds.extend(generated_audios.to("cpu"))
+                        eval_descriptions.extend(input_ids.to("cpu"))
+                        eval_prompts.extend(prompts.to("cpu"))
+
+                eval_time = time.time() - eval_start
+                # normalize eval metrics
+                eval_metrics = {
+                    key: torch.mean(torch.cat([d[key] for d in eval_metrics])).to("cpu") for key in eval_metrics[0]
+                }
+
+                # compute metrics
+                metrics_desc = ""
+                if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
+                    if accelerator.is_local_main_process:
+                        (
+                            metric_values,
+                            pred_descriptions,
+                            pred_prompts,
+                            audios,
+                            transcriptions,
+                            si_sdr_measures,
+                        ) = compute_metrics(
+                            eval_preds,
+                            eval_descriptions,
+                            eval_prompts,
+                            accelerator.device,
+                            training_args.compute_clap_similarity_metric,
+                            training_args.compute_noise_level_metric,
+                            training_args.noise_level_to_compute_clean_wer,
+                        )
+                        eval_metrics.update(metric_values)
+                        metrics_desc = " ".join([f"Eval {key}: {value} |" for key, value in metric_values.items()])
+                        if "wandb" in training_args.report_to:
+                            log_pred(
+                                accelerator,
+                                pred_descriptions,
+                                pred_prompts,
+                                transcriptions,
+                                audios,
+                                si_sdr_measures,
+                                sampling_rate=sampling_rate,
+                                step=cur_step,
+                                prefix="eval",
+                            )
+                    accelerator.wait_for_everyone()
+
+                # Print metrics and update progress bar
+                if accelerator.is_local_main_process:
+                    steps_trained_progress_bar.write(
+                        f"Eval results for step ({cur_step} / {total_train_steps} | Eval Loss: {eval_metrics['loss']} |"
+                        f" {metrics_desc})"
+                    )
+
+                log_metric(
+                    accelerator,
+                    metrics=eval_metrics,
+                    train_time=eval_time,
+                    step=cur_step,
+                    epoch=epoch,
+                    prefix="eval",
+                )
+
+                # release eval batch and relax metrics
+                eval_metrics, eval_preds, eval_descriptions, eval_prompts, batch, eval_metric = release_memory(
+                    eval_metrics, eval_preds, eval_descriptions, eval_prompts, batch, eval_metric
+                )
+                if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
+                    generated_audios, input_ids, prompts = release_memory(generated_audios, input_ids, prompts)
+
+                # train mode
+                model.train()
+
+                # flush the train metrics
+                train_start = time.time()
+
+            # break condition
+            if cur_step == total_train_steps:
+                continue_training = False
+                break
+
+        if not continue_training:
+            break
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main()

capspeech/ar/training/run_parler_tts_training.py

@@ -0,0 +1,1279 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+""" Train Parler-TTS using 🤗 Accelerate"""
+
+import logging
+import os
+import re
+import sys
+import time
+import math
+import contextlib
+from multiprocess import set_start_method
+from datetime import timedelta
+import inspect
+from tqdm import tqdm
+from pathlib import Path
+import wandb
+
+import torch
+from torch.utils.data import DataLoader
+
+import datasets
+from datasets import DatasetDict, Dataset, IterableDataset, concatenate_datasets
+
+from huggingface_hub import HfApi
+
+import transformers
+from transformers import AutoFeatureExtractor, AutoTokenizer, HfArgumentParser
+from transformers.trainer_pt_utils import LengthGroupedSampler
+from transformers.optimization import get_scheduler
+from transformers.utils import send_example_telemetry
+
+
+from accelerate import Accelerator, skip_first_batches
+from accelerate.utils import set_seed, AutocastKwargs, InitProcessGroupKwargs, TorchDynamoPlugin, DistributedDataParallelKwargs
+from accelerate.utils.memory import release_memory
+
+from parler_tts import (
+    ParlerTTSConfig,
+    ParlerTTSForConditionalGeneration,
+    build_delay_pattern_mask,
+)
+
+from training.utils import (
+    get_last_checkpoint,
+    rotate_checkpoints,
+    log_pred,
+    log_metric,
+    load_all_codec_checkpoints,
+    save_codec_checkpoint,
+    get_last_codec_checkpoint_step,
+)
+from training.arguments import ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments
+from training.data import load_multiple_datasets, DataCollatorParlerTTSWithPadding, DataCollatorEncodecWithPadding
+from training.eval import clap_similarity, wer, si_sdr
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+    # See all possible arguments in src/transformers/training_args.py
+    # or by passing the --help flag to this script.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments))
+    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.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
+    # information sent is the one passed as arguments along with your Python/PyTorch versions.
+    send_example_telemetry("run_parler_tts", model_args, data_args)
+    
+    if data_args.wandb_key is not None:
+        wandb.login(key=data_args.wandb_key)
+
+    if training_args.dtype == "float16":
+        mixed_precision = "fp16"
+        torch_dtype = torch.float16
+    elif training_args.dtype == "bfloat16":
+        mixed_precision = "bf16"
+        torch_dtype = torch.bfloat16
+    else:
+        mixed_precision = "no"
+        torch_dtype = torch.float32
+
+    if data_args.pad_to_max_length and (
+        data_args.max_duration_in_seconds is None
+        or data_args.max_prompt_token_length is None
+        or data_args.max_description_token_length is None
+    ):
+        raise ValueError(
+            "`pad_to_max_length` is `True` but one of the following parameters has not been set: `max_duration_in_seconds`, `max_prompt_token_length`, `max_description_token_length`"
+        )
+
+    padding = "max_length" if data_args.pad_to_max_length else "longest"
+
+    ####### A. Preparation
+    kwargs_handlers = [InitProcessGroupKwargs(timeout=timedelta(minutes=120)), DistributedDataParallelKwargs(find_unused_parameters=False)]
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
+        mixed_precision=mixed_precision,
+        log_with=training_args.report_to,
+        project_dir=training_args.output_dir,
+        kwargs_handlers=kwargs_handlers,
+    )
+
+    accelerator.init_trackers(
+        project_name=data_args.wandb_project,
+        config={
+            "learning_rate": training_args.learning_rate,
+            "model_name_or_path": model_args.model_name_or_path,
+            "num_train_epochs": training_args.num_train_epochs,
+            "gradient_accumulation_steps": training_args.gradient_accumulation_steps,
+            "per_device_train_batch_size": training_args.per_device_train_batch_size,
+            "global_batch_size": training_args.per_device_train_batch_size * accelerator.num_processes,
+            "mixed_precision": mixed_precision,
+            "lr_scheduler_type": training_args.lr_scheduler_type,
+            "warmup_steps": training_args.warmup_steps,
+            "freeze_text_encoder": model_args.freeze_text_encoder,
+            "max_duration_in_seconds": data_args.max_duration_in_seconds,
+            "weight_decay": training_args.weight_decay,
+            "adam_beta1": training_args.adam_beta1,
+            "adam_beta2": training_args.adam_beta2,
+            "temperature": model_args.temperature,
+        },
+        init_kwargs={"wandb": {"name": data_args.wandb_run_name}} if data_args.wandb_run_name else {},
+    )
+
+    # Detecting last checkpoint and eventually continue from last checkpoint
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+
+    # Setup logging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+    logger.setLevel(logging.INFO if accelerator.is_main_process else logging.WARN)
+
+    # Log a small summary on each proces
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
+        f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
+    )
+
+    # Set the verbosity to info of the Transformers logger (on main process only)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+    num_workers = data_args.preprocessing_num_workers
+
+    # 1. First, lett's instantiate the feature extractor, tokenizers and model
+    # Note for distributed training, the .from_pretrained methods guarantee that only
+    # one local process can concurrently download model & vocab.
+
+    # load feature extractor
+    feature_extractor = AutoFeatureExtractor.from_pretrained(
+        model_args.feature_extractor_name or model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+    ) 
+    sampling_rate = feature_extractor.sampling_rate
+
+    # load prompt tokenizer
+    prompt_tokenizer = AutoTokenizer.from_pretrained(
+        model_args.prompt_tokenizer_name or model_args.description_tokenizer_name or model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+        use_fast=model_args.use_fast_tokenizer,
+        padding_side=model_args.prompt_padding_side,
+    )
+
+    # load description tokenizer
+    description_tokenizer = AutoTokenizer.from_pretrained(
+        model_args.description_tokenizer_name or model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+        use_fast=model_args.use_fast_tokenizer,
+    )
+
+    if model_args.use_fast_tokenizer:
+        logger.warning(
+            "Disabling fast tokenizer warning: https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3231-L3235"
+        )
+        prompt_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
+        description_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
+
+    # 2. Now, let's load the dataset
+
+    if data_args.save_to_disk is not None:
+        os.makedirs(data_args.save_to_disk, exist_ok=True)
+
+    # assume that the dataset has been saved to `save_to_disk` if the latter is not empty
+    dataset_was_precomputed = len(os.listdir(data_args.save_to_disk)) > 0
+    if dataset_was_precomputed:
+        with accelerator.local_main_process_first():
+            vectorized_datasets = datasets.load_from_disk(data_args.save_to_disk)
+    else:
+        raw_datasets = DatasetDict()
+
+        columns_to_keep = {
+            "target_audio_column_name": data_args.target_audio_column_name,
+            "prompt_column_name": data_args.prompt_column_name,
+            "source": data_args.source_column_name,
+        }
+        if data_args.description_column_name is not None:
+            columns_to_keep["description_column_name"] = data_args.description_column_name
+
+        if training_args.do_train:
+            raw_datasets["train"] = load_multiple_datasets(
+                accelerator,
+                data_args.train_dataset_name,
+                splits=data_args.train_split_name,
+                dataset_samples=data_args.train_dataset_samples,
+                seed=training_args.seed,
+                cache_dir=model_args.cache_dir,
+                num_proc=data_args.preprocessing_num_workers,
+                id_column_name=data_args.id_column_name,
+                columns_to_keep=columns_to_keep.values(),
+                prompt_column_name=data_args.prompt_column_name,
+                audio_column_name=data_args.target_audio_column_name,
+                sampling_rate=sampling_rate,
+                logger=logger,
+                mls_dir=data_args.mls_dir,
+                librittsrmix_dir=data_args.librittsrmix_dir,
+                gigaspeech_dir=data_args.gigaspeech_dir,
+                commonvoice_dir=data_args.commonvoice_dir,
+                emilia_dir=data_args.emilia_dir,
+                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
+            )
+
+            for key in columns_to_keep:
+                if columns_to_keep[key] not in raw_datasets["train"].column_names:
+                    raise ValueError(
+                        f"--{key} '{columns_to_keep[key]}' not found in dataset '{data_args.train_dataset_name}'."
+                        f" Make sure to set `--{key}` to the correct audio column - one of"
+                        f" {', '.join(raw_datasets['train'].column_names)}."
+                    )
+
+            if data_args.max_train_samples is not None:
+                raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))
+
+        if training_args.do_eval:
+            raw_datasets["eval"] = load_multiple_datasets(
+                accelerator,
+                data_args.eval_dataset_name if data_args.eval_dataset_name else data_args.train_dataset_name,
+                splits=data_args.eval_split_name,
+                cache_dir=model_args.cache_dir,
+                num_proc=data_args.preprocessing_num_workers,
+                id_column_name=data_args.id_column_name,
+                columns_to_keep=columns_to_keep.values(),
+                prompt_column_name=data_args.prompt_column_name,
+                audio_column_name=data_args.target_audio_column_name,
+                sampling_rate=sampling_rate,
+                logger=logger,
+                mls_dir=data_args.mls_dir,
+                librittsrmix_dir=data_args.librittsrmix_dir,
+                gigaspeech_dir=data_args.gigaspeech_dir,
+                commonvoice_dir=data_args.commonvoice_dir,
+                emilia_dir=data_args.emilia_dir
+                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
+            )
+
+            if data_args.max_eval_samples is not None:
+                with accelerator.local_main_process_first():
+                    raw_datasets["eval"] = (
+                        raw_datasets["eval"].shuffle(seed=training_args.seed).select(range(data_args.max_eval_samples))
+                    )
+
+    # 3. Next, let's load the config.
+    config = ParlerTTSConfig.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+    )
+ 
+    if training_args.codebook_weights is not None and len(training_args.codebook_weights) != config.decoder.num_codebooks:
+        raise ValueError(f"`codebook_weights` has length {len(training_args.codebook_weights)} when it should be of length {config.decoder.num_codebooks}.")
+
+    # update pad token id and decoder_start_token_id
+    config.decoder.update(
+        {
+            "cross_attention_implementation_strategy": model_args.cross_attention_implementation_strategy
+            if model_args.cross_attention_implementation_strategy is not None
+            else None,
+            "codebook_weights": training_args.codebook_weights if training_args.codebook_weights is not None else config.decoder.codebook_weights
+        }
+    )
+    config.update(
+        {
+            "pad_token_id": model_args.pad_token_id if model_args.pad_token_id is not None else config.pad_token_id,
+            "decoder_start_token_id": model_args.decoder_start_token_id
+            if model_args.decoder_start_token_id is not None
+            else config.decoder_start_token_id,
+        }
+    )
+
+    with open("events.txt", "r") as f:
+        events = [line.strip() for line in f]
+    events = ["<"+event.lower().replace(" ", "_")+">" for event in events]
+    events.append("<B_start>")
+    events.append("<B_end>")
+    events.append("<I_start>")
+    events.append("<I_end>")
+
+    special_tokens = {"additional_special_tokens": events}
+    prompt_tokenizer.add_special_tokens(special_tokens)
+    description_tokenizer.add_special_tokens(special_tokens)
+    padded_vocab_size = ((len(prompt_tokenizer) + 127) // 128) * 128 
+    config.vocab_size = padded_vocab_size 
+
+    # create model
+    model = ParlerTTSForConditionalGeneration.from_pretrained(
+        model_args.model_name_or_path,
+        ignore_mismatched_sizes=True,
+        cache_dir=model_args.cache_dir,
+        config=config,
+        token=data_args.token,
+        trust_remote_code=data_args.trust_remote_code,
+        attn_implementation={"decoder": model_args.attn_implementation, "text_encoder": "eager"},
+    )
+    model.text_encoder.resize_token_embeddings(padded_vocab_size)
+
+    # enable gradient checkpointing if necessary
+    if training_args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+
+    # 4. Now we preprocess the datasets including loading the audio, resampling and normalization
+    # Thankfully, `datasets` takes care of automatically loading and resampling the audio,
+    # so that we just need to set the correct target sampling rate and normalize the input
+    # via the `feature_extractor`
+
+    # derive max & min input length for sample rate & max duration
+    sampling_rate = feature_extractor.sampling_rate
+    max_target_length = int(data_args.max_duration_in_seconds * sampling_rate)
+    min_target_length = int(data_args.min_duration_in_seconds * sampling_rate)
+    target_audio_column_name = data_args.target_audio_column_name
+    description_column_name = data_args.description_column_name
+    prompt_column_name = data_args.prompt_column_name
+    feature_extractor_input_name = feature_extractor.model_input_names[0]
+    audio_encoder_pad_token_id = config.decoder.pad_token_id
+    audio_encoder_eos_token_id = config.decoder.eos_token_id
+    audio_encoder_bos_token_id = model.generation_config.decoder_start_token_id
+    max_length = model.generation_config.max_length
+    num_codebooks = model.decoder.config.num_codebooks
+    bandwidth = model_args.bandwidth
+    attn_implementation = model_args.attn_implementation
+
+    # Freeze Encoders
+    model.freeze_encoders(model_args.freeze_text_encoder)
+
+    # Test all gather - used for warmout and avoiding timeout
+    logger.debug(str(accelerator.process_index), main_process_only=False, in_order=True)
+    test_tensor = torch.tensor([accelerator.process_index], device=accelerator.device)
+    gathered_tensor = accelerator.gather(test_tensor)
+    print("gathered_tensor", gathered_tensor)
+    accelerator.wait_for_everyone()
+
+    if not dataset_was_precomputed:
+        # Filter on text length
+        if description_column_name is not None and data_args.max_text_length is not None:
+            with accelerator.local_main_process_first():
+                # filter description that is shorter than max_text_length
+                raw_datasets = raw_datasets.filter(
+                    lambda x: len(x) < data_args.max_text_length,
+                    num_proc=num_workers,
+                    input_columns=[description_column_name],
+                )
+
+        # Preprocessing the dataset.
+        # We need to tokenize the texts.
+        def pass_through_processors(description, prompt):
+            batch = {}
+
+            batch["input_ids"] = description_tokenizer(description.strip())["input_ids"]
+            batch["prompt_input_ids"] = prompt_tokenizer(prompt.strip())["input_ids"]
+
+            return batch
+
+        with accelerator.local_main_process_first():
+            # this is a trick to avoid to rewrite the entire audio column which takes ages
+            vectorized_datasets = raw_datasets.map(
+                pass_through_processors,
+                remove_columns=next(iter(raw_datasets.values())).column_names,
+                input_columns=[description_column_name, prompt_column_name],
+                num_proc=num_workers,
+                desc="preprocess datasets",
+            )
+
+        # We use Accelerate to perform distributed inference
+        # T5 doesn't support fp16
+        autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))
+
+        # Now we encode the audio labels with encodec.
+        ####### B. Encode audio
+
+        logger.info("*** Encode target audio with encodec ***")
+
+        # no need to prepare audio_decoder because used for inference without mixed precision
+        # see: https://huggingface.co/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare
+        if training_args.torch_compile:
+            audio_decoder = accelerator.prepare_model(model.audio_encoder, evaluation_mode=True)
+        else:
+            audio_decoder = model.audio_encoder
+
+        encoder_data_collator = DataCollatorEncodecWithPadding(
+            feature_extractor,
+            audio_column_name=target_audio_column_name,
+            mls_dir=data_args.mls_dir,
+            librittsrmix_dir=data_args.librittsrmix_dir,
+            gigaspeech_dir=data_args.gigaspeech_dir,
+            commonvoice_dir=data_args.commonvoice_dir,
+            emilia_dir=data_args.emilia_dir,
+            feature_extractor_input_name=feature_extractor_input_name,
+            max_length=max_target_length,
+            padding=padding,
+        )
+        encoder_signature = set(inspect.signature(audio_decoder.forward).parameters)
+
+        def apply_audio_decoder(batch):
+            len_audio = batch.pop("len_audio")
+            audio_decoder.to(batch["input_values"].device).eval()
+            if bandwidth is not None:
+                batch["bandwidth"] = bandwidth
+            elif "num_quantizers" in encoder_signature:
+                batch["num_quantizers"] = num_codebooks
+            elif "num_codebooks" in encoder_signature:
+                batch["num_codebooks"] = num_codebooks
+            elif "n_quantizers" in encoder_signature:
+                batch["n_quantizers"] = num_codebooks
+
+            with torch.no_grad():
+                labels = audio_decoder.encode(**batch)["audio_codes"]
+            output = {}
+            output["len_audio"] = len_audio
+            # (1, bsz, codebooks, seq_len) -> (bsz, seq_len, codebooks)
+            output["labels"] = labels.squeeze(0).transpose(1, 2)
+
+            # if `pad_to_max_length`, the maximum corresponding audio length of the current batch is max_duration*sampling_rate
+            max_length = len_audio.max() if padding != "max_length" else max_target_length
+            output["ratio"] = torch.ones_like(len_audio) * labels.shape[-1] / max_length
+            return output
+
+        # (1, codebooks, seq_len) where seq_len=1
+        bos_labels = torch.ones((1, num_codebooks, 1)) * audio_encoder_bos_token_id
+
+        def postprocess_dataset(labels):
+            # (1, codebooks, seq_len)
+            labels = torch.tensor(labels).unsqueeze(0)
+            # add bos
+            labels = torch.cat([bos_labels, labels], dim=-1)
+
+            labels, delay_pattern_mask = build_delay_pattern_mask(
+                labels,
+                bos_token_id=audio_encoder_bos_token_id,
+                pad_token_id=audio_encoder_eos_token_id,
+                max_length=labels.shape[-1] + num_codebooks,
+                num_codebooks=num_codebooks,
+            )
+
+            # the first ids of the delay pattern mask are precisely labels, we use the rest of the labels mask
+            # to take care of EOS
+            # we want labels to look like this:
+            #  - [B, a, b, E, E, E, E]
+            #  - [B, B, c, d, E, E, E]
+            #  - [B, B, B, e, f, E, E]
+            #  - [B, B, B, B, g, h, E]
+            labels = torch.where(delay_pattern_mask == -1, audio_encoder_eos_token_id, delay_pattern_mask)
+
+            # the first timestamp is associated to a row full of BOS, let's get rid of it
+            # we also remove the last timestampts (full of PAD)
+            output = {"labels": labels[:, 1:]}
+            return output
+
+        for split in vectorized_datasets:
+            data_loader = DataLoader(
+                raw_datasets[split],
+                batch_size=training_args.audio_encoder_per_device_batch_size,
+                collate_fn=encoder_data_collator,
+                num_workers=training_args.dataloader_num_workers,
+                pin_memory=True,
+            )
+            data_loader = accelerator.prepare(data_loader)
+            total_inference_steps = len(data_loader)
+
+            start_step = get_last_codec_checkpoint_step(os.path.join(data_args.temporary_save_to_disk, split))
+            accelerator.wait_for_everyone()
+            if start_step > 0:
+                logger.info(f"Resuming {split} from step {start_step}")
+                # efficiently skip the first n batches
+                start_step += 1
+                data_loader = skip_first_batches(data_loader, start_step)
+
+            all_generated_labels = []
+            all_lens = []
+            if start_step < total_inference_steps:
+                for i, batch in enumerate(tqdm(data_loader, disable=not accelerator.is_local_main_process)):
+                    cur_step = start_step + i
+                    generate_labels = apply_audio_decoder(batch)
+                    generate_labels = accelerator.pad_across_processes(generate_labels, dim=1, pad_index=0)
+                    generate_labels = accelerator.gather_for_metrics(generate_labels)
+
+                    if accelerator.is_main_process:
+                        lab = generate_labels["labels"].cpu().transpose(1, 2).to(torch.int16)
+                        rat = generate_labels["ratio"].cpu().squeeze(1)
+                        lens = generate_labels["len_audio"].cpu().squeeze(1)
+                        lab = [l[:, : int(ratio * length)] for (l, ratio, length) in zip(lab, rat, lens)]
+
+                        all_generated_labels.extend(lab)
+                        all_lens.extend(lens)
+
+                        if ((cur_step + 1) % data_args.save_codec_steps == 0) or (
+                            cur_step == total_inference_steps - 1
+                        ):
+                            tmp_labels = Dataset.from_dict({"labels": all_generated_labels, "target_length": all_lens})
+                            tmp_labels = tmp_labels.map(
+                                postprocess_dataset,
+                                num_proc=data_args.preprocessing_num_workers,  # this one is resource consuming if many processor.
+                                input_columns=["labels"],
+                                desc="Postprocessing labeling",
+                            )
+                            save_codec_checkpoint(
+                                os.path.join(data_args.temporary_save_to_disk, split), tmp_labels, cur_step
+                            )
+                            all_generated_labels = []
+                            all_lens = []
+
+                accelerator.wait_for_everyone()
+
+            if accelerator.is_main_process and len(all_generated_labels) > 0:
+                tmp_labels = Dataset.from_dict({"labels": all_generated_labels, "target_length": all_lens})
+                tmp_labels = tmp_labels.map(
+                    postprocess_dataset,
+                    num_proc=data_args.preprocessing_num_workers,  # this one is resource consuming if many processor.
+                    input_columns=["labels"],
+                    desc="Postprocessing labeling",
+                )
+                save_codec_checkpoint(os.path.join(data_args.temporary_save_to_disk, split), tmp_labels, cur_step)
+                all_generated_labels = []
+                all_lens = []
+            accelerator.wait_for_everyone()
+
+            del all_generated_labels
+            accelerator.wait_for_everyone()
+
+            with accelerator.local_main_process_first():
+                tmp_labels = load_all_codec_checkpoints(os.path.join(data_args.temporary_save_to_disk, split)).select(
+                    range(len(vectorized_datasets[split]))
+                )
+                logger.info(f"Concatenating {split}: {tmp_labels} with {vectorized_datasets[split]}")
+                vectorized_datasets[split] = concatenate_datasets([vectorized_datasets[split], tmp_labels], axis=1)
+
+        accelerator.free_memory()
+        del generate_labels, all_lens
+
+        with accelerator.local_main_process_first():
+            # NOTE: filtering is done at the end because in the `datasets` library, caching audio files is done after most operations
+            # caching audio files is time and disk-space consuming, so we want to avoid it at all costs, especially for large (>1Kh) audio datasets.
+            # That's also why we avoid to concat the processed datasets (vectorized_datasets) with the audio column present in raw_datasets.
+
+            def is_audio_in_length_range(length):
+                return length > min_target_length and length < max_target_length
+
+            # filter data that is shorter than min_target_length
+            vectorized_datasets = vectorized_datasets.filter(
+                is_audio_in_length_range,
+                num_proc=num_workers,
+                input_columns=["target_length"],
+            )
+
+        if description_column_name is not None and data_args.max_description_token_length is not None:
+            with accelerator.local_main_process_first():
+                # filter description that is shorter than max_text_length
+                vectorized_datasets = vectorized_datasets.filter(
+                    lambda x: len(x) < data_args.max_description_token_length,
+                    num_proc=num_workers,
+                    input_columns=["input_ids"],
+                )
+
+        if data_args.max_prompt_token_length is not None:
+            with accelerator.local_main_process_first():
+                # filter description that is shorter than max_text_length
+                vectorized_datasets = vectorized_datasets.filter(
+                    lambda x: len(x) < data_args.max_prompt_token_length,
+                    num_proc=num_workers,
+                    input_columns=["prompt_input_ids"],
+                )
+
+    if data_args.save_to_disk is not None and not dataset_was_precomputed:
+        if accelerator.is_main_process:
+            vectorized_datasets.save_to_disk(
+                data_args.save_to_disk,
+                num_proc=min(data_args.preprocessing_num_workers, len(vectorized_datasets["eval"]) - 1),
+            )
+        accelerator.wait_for_everyone()
+        logger.info(f"Dataset saved at {data_args.save_to_disk}")
+
+    audio_max_length = None
+    if padding == "max_length":
+        audio_max_length = max(vectorized_datasets["train"]["target_length"])
+        with accelerator.local_main_process_first():
+            max_sample = vectorized_datasets["train"].filter(
+                lambda x: x == audio_max_length,
+                num_proc=num_workers,
+                input_columns=["target_length"],
+            )
+        audio_max_length = max([len(l[0]) for l in max_sample["labels"]])
+
+    if description_column_name is not None and data_args.max_description_token_length is not None:
+        with accelerator.local_main_process_first():
+            # filter description that is shorter than max_text_length
+            vectorized_datasets = vectorized_datasets.filter(
+                lambda x: len(x) < data_args.max_description_token_length,
+                num_proc=num_workers,
+                input_columns=["input_ids"],
+            )
+
+    if data_args.max_prompt_token_length is not None:
+        with accelerator.local_main_process_first():
+            # filter description that is shorter than max_text_length
+            vectorized_datasets = vectorized_datasets.filter(
+                lambda x: len(x) < data_args.max_prompt_token_length,
+                num_proc=num_workers,
+                input_columns=["prompt_input_ids"],
+            )
+
+    if training_args.group_by_length:
+        # apply a simple heuristic to take into account audio and text lengths
+        def add_target_lengths(target_length, prompt, description):
+            return {"target_length": target_length + len(prompt) + len(description)}
+
+        with accelerator.local_main_process_first():
+            vectorized_datasets = vectorized_datasets.map(
+                add_target_lengths,
+                num_proc=num_workers,
+                input_columns=["target_length", "prompt_input_ids", "input_ids"],
+            )
+
+    # for large datasets it is advised to run the preprocessing on a
+    # single machine first with ``args.preprocessing_only`` since there will mostly likely
+    # be a timeout when running the script in distributed mode.
+    # In a second step ``args.preprocessing_only`` can then be set to `False` to load the
+    # cached dataset
+    if data_args.preprocessing_only and data_args.save_to_disk is None:
+        raise ValueError(
+            "`preprocessing_only=True` but `save_to_disk` is not set. The latter should indicates where to save the dataset locally."
+        )
+    elif data_args.preprocessing_only:
+        logger.info(f"Data preprocessing finished. Files save at {data_args.save_to_disk}")
+        return
+
+    # 6. Next, we can prepare the training.
+
+    # Let's use word CLAP similary and WER metrics as our evaluation metrics,
+    def compute_metrics(
+        audios,
+        descriptions,
+        prompts,
+        device="cpu",
+        compute_clap_similarity_metric=False,
+        compute_noise_level_metric=False,
+        noise_level_to_compute_clean_wer=None,
+    ):
+        results = {}
+        input_ids = descriptions
+        texts = description_tokenizer.batch_decode(input_ids, skip_special_tokens=True)
+        prompts = prompt_tokenizer.batch_decode(prompts, skip_special_tokens=True)
+        audios = [a.float().cpu().numpy() for a in audios]
+
+        if compute_clap_similarity_metric:
+            clap_score = clap_similarity(
+                model_args.clap_model_name_or_path, texts, audios, device, input_sampling_rate=sampling_rate
+            )
+            results["clap"] = clap_score
+
+        si_sdr_measures = None
+        if compute_noise_level_metric:
+            si_sdr_measures = si_sdr(audios, device, input_sampling_rate=sampling_rate)
+
+        word_error, transcriptions, clean_word_error, noisy_word_error, percent_clean_samples = wer(
+            model_args.asr_model_name_or_path,
+            prompts,
+            audios,
+            device,
+            training_args.per_device_eval_batch_size,
+            sampling_rate,
+            noise_level_to_compute_clean_wer,
+            si_sdr_measures,
+        )
+        results["wer"] = word_error
+        if clean_word_error is not None:
+            results["clean_wer"] = clean_word_error
+            results["noisy_word_error"] = noisy_word_error
+            results["percent_clean_samples"] = percent_clean_samples
+
+        return results, texts, prompts, audios, transcriptions, si_sdr_measures
+
+    # Define Training Schedule
+    # Store some constants
+    per_device_train_batch_size = int(training_args.per_device_train_batch_size)
+    train_batch_size = per_device_train_batch_size * accelerator.num_processes
+    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
+    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
+
+    if training_args.max_steps < 0:
+        num_epochs = int(training_args.num_train_epochs)
+        steps_per_epoch = len(vectorized_datasets["train"]) // (train_batch_size * gradient_accumulation_steps)
+        total_train_steps = steps_per_epoch * num_epochs
+    elif training_args.max_steps > 0:
+        logger.info("max_steps is given, it will override any value given in num_train_epochs")
+        total_train_steps = int(training_args.max_steps)
+        # Setting a very large number of epochs so we go as many times as necessary over the iterator.
+        num_epochs = sys.maxsize
+        steps_per_epoch = total_train_steps
+
+    if training_args.eval_steps is None:
+        logger.info(f"eval_steps is not set, evaluating at the end of each epoch")
+        eval_steps = steps_per_epoch
+    else:
+        eval_steps = training_args.eval_steps
+        
+    if training_args.eval_generation_steps is None:
+        eval_generation_steps = eval_steps
+    else:
+        eval_generation_steps = training_args.eval_generation_steps
+
+    # T5 doesn't support fp16
+    autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))
+
+    # Define optimizer, LR scheduler, collator
+    optimizer = torch.optim.AdamW(
+        params=model.parameters(),
+        lr=training_args.learning_rate,
+        betas=(training_args.adam_beta1, training_args.adam_beta2),
+        eps=training_args.adam_epsilon,
+        weight_decay=training_args.weight_decay,
+    )
+
+    # LR scheduler gets stepped by `num_processes` each time -> account for this in warmup / total steps
+    lr_scheduler = get_scheduler(
+        name=training_args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=training_args.get_warmup_steps(total_train_steps) * accelerator.num_processes,
+        num_training_steps=total_train_steps * accelerator.num_processes,
+    )
+
+    # Instantiate custom data collator
+    data_collator = DataCollatorParlerTTSWithPadding(
+        prompt_tokenizer=prompt_tokenizer,
+        description_tokenizer=description_tokenizer,
+        pad_to_multiple_of=data_args.pad_to_multiple_of,
+        padding=padding,
+        prompt_max_length=data_args.max_prompt_token_length,
+        description_max_length=data_args.max_description_token_length,
+        audio_max_length=audio_max_length,
+    )
+
+    # Prepare everything with accelerate
+    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
+
+    num_examples = total_train_steps * train_batch_size * gradient_accumulation_steps
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {num_examples}")
+    logger.info("  Instantaneous batch size per device =" f" {per_device_train_batch_size}")
+    logger.info("  Gradient accumulation steps =" f" {gradient_accumulation_steps}")
+    logger.info(
+        f"  Total train batch size (w. parallel & distributed) = {train_batch_size * gradient_accumulation_steps}"
+    )
+    logger.info(f"  Total optimization steps = {total_train_steps}")
+
+    # ======================== Training ================================
+    train_time = 0
+    train_start = time.time()
+    steps_trained_progress_bar = tqdm(
+        range(total_train_steps), desc="Train steps ... ", position=0, disable=not accelerator.is_local_main_process
+    )
+    continue_training = True
+    epochs_trained = 0
+    cur_step = 0
+
+    checkpoint = None
+    if training_args.resume_from_checkpoint is not None:
+        checkpoint = training_args.resume_from_checkpoint
+    elif last_checkpoint is not None:
+        checkpoint = last_checkpoint
+
+    if accelerator.is_main_process:
+        if training_args.push_to_hub:
+            api = HfApi(token=training_args.hub_token)
+
+            # Create repo (repo_name from args or inferred)
+            repo_name = training_args.hub_model_id
+            if repo_name is None:
+                repo_name = Path(training_args.output_dir).absolute().name
+            repo_id = api.create_repo(repo_name, exist_ok=True).repo_id
+
+            with open(os.path.join(training_args.output_dir, ".gitignore"), "w+") as gitignore:
+                if "wandb" not in gitignore:
+                    gitignore.write("wandb\n")
+        elif training_args.output_dir is not None:
+            os.makedirs(training_args.output_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    # Now save everything to be able to create a single processor later
+    # make sure all processes wait until data is saved
+    # only the main process saves them
+    if accelerator.is_main_process:
+        # save feature extractor, tokenizer and config
+        if (
+            model_args.prompt_tokenizer_name is None
+            and model_args.description_tokenizer_name
+            or (model_args.prompt_tokenizer_name == model_args.description_tokenizer_name)
+        ):
+            prompt_tokenizer.save_pretrained(training_args.output_dir)
+        else:
+            logger.warning(
+                f"Prompt tokenizer ('{model_args.prompt_tokenizer_name}') and description tokenizer ('{model_args.description_tokenizer_name}') are not the same. Saving only the prompt tokenizer."
+            )
+            prompt_tokenizer.save_pretrained(training_args.output_dir)
+
+        feature_extractor.save_pretrained(training_args.output_dir)
+        config.save_pretrained(training_args.output_dir)
+    accelerator.wait_for_everyone()
+
+    if checkpoint is not None:
+        accelerator.load_state(checkpoint)
+        # Find num steps and epoch from saved state string pattern
+        pattern = r"checkpoint-(\d+)-epoch-(\d+)"
+        match = re.search(pattern, checkpoint)
+        cur_step = int(match.group(1))
+        epochs_trained = int(match.group(2))
+
+        logger.info("  Continuing training from checkpoint, will skip to saved global_step")
+        logger.info(f"  Continuing training from epoch {epochs_trained}")
+        logger.info(f"  Continuing training from global step {cur_step}")
+
+        steps_trained_progress_bar.update(cur_step)
+
+        for epoch in range(0, epochs_trained):
+            with accelerator.local_main_process_first():
+                vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+
+        if training_args.max_steps < 0:
+            # we know exactly the number of steps per epoch, so can skip through the required number of batches
+            resume_step = (cur_step - epochs_trained * steps_per_epoch) * gradient_accumulation_steps
+        else:
+            # Currently we don't know how many steps we've taken in the current epoch
+            # So we just shuffle the dataset one extra time and start from a fresh epoch
+            # This is "good enough" for our purposes but not fully correct
+            resume_step = None
+            with accelerator.local_main_process_first():
+                vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+    else:
+        resume_step = None
+
+    gen_kwargs = {
+        "do_sample": model_args.do_sample,
+        "temperature": model_args.temperature,
+        "max_length": model_args.max_length,
+        # Because of the delayed pattern mask, generation might stop earlier because of unexpected behaviour
+        # on the first tokens of the codebooks that are delayed.
+        # This fix the issue.
+        "min_new_tokens": num_codebooks + 1,
+    }
+
+    # Define gradient update step fn
+    def train_step(
+        batch,
+        accelerator,
+        autocast_kwargs,
+        num_items_in_batch,
+        gradient_accumulation_steps,
+    ):
+        if mixed_precision == "fp16":
+            # fp16 doesn't work with T5-like models
+            with accelerator.autocast(autocast_handler=autocast_kwargs):
+                if training_args.parallel_mode.value != "distributed":
+                    encoder_outputs = model.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                else:
+                    encoder_outputs = model.module.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                # we optionnally project last_hidden_state to avoid recomputing every time
+                encoder_hidden_states = encoder_outputs.last_hidden_state
+                if (
+                    config.text_encoder.hidden_size != config.decoder.hidden_size
+                    and config.decoder.cross_attention_hidden_size is None
+                ):
+                    encoder_hidden_states = (
+                        model.enc_to_dec_proj(encoder_hidden_states)
+                        if training_args.parallel_mode.value != "distributed"
+                        else model.module.enc_to_dec_proj(encoder_hidden_states)
+                    )
+
+                if batch.get("attention_mask", None) is not None:
+                    encoder_hidden_states = encoder_hidden_states * batch.get("attention_mask", None)[..., None]
+
+                encoder_outputs.last_hidden_state = encoder_hidden_states
+                batch["encoder_outputs"] = encoder_outputs
+
+        outputs = model(**batch, loss_reduction="sum")
+        # CE (data) loss
+        ce_loss = (outputs.loss * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch
+
+        metrics = {"loss": ce_loss}
+        
+        # per CE loss
+        per_codebook_losses = outputs.per_codebook_losses
+        metrics.update({f"codebook_{i}_loss": ((l  * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch) for (i,l) in enumerate(per_codebook_losses)})
+        return ce_loss, metrics
+
+    # Define eval fn
+    def eval_step(
+        batch,
+        accelerator,
+        autocast_kwargs,
+    ):
+        eval_model = model if not training_args.torch_compile else model._orig_mod
+
+        if mixed_precision == "fp16":
+            # fp16 doesn't work with T5-like models
+            with accelerator.autocast(autocast_handler=autocast_kwargs):
+                if training_args.parallel_mode.value != "distributed":
+                    encoder_outputs = model.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                else:
+                    encoder_outputs = model.module.text_encoder(
+                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
+                    )
+                # we optionnally project last_hidden_state to avoid recomputing every time
+                encoder_hidden_states = encoder_outputs.last_hidden_state
+                if (
+                    config.text_encoder.hidden_size != config.decoder.hidden_size
+                    and config.decoder.cross_attention_hidden_size is None
+                ):
+                    encoder_hidden_states = (
+                        model.enc_to_dec_proj(encoder_hidden_states)
+                        if training_args.parallel_mode.value != "distributed"
+                        else model.module.enc_to_dec_proj(encoder_hidden_states)
+                    )
+
+                if batch.get("attention_mask", None) is not None:
+                    encoder_hidden_states = encoder_hidden_states * batch.get("attention_mask", None)[..., None]
+
+                encoder_outputs.last_hidden_state = encoder_hidden_states
+                batch["encoder_outputs"] = encoder_outputs
+
+        with torch.no_grad():
+            outputs = eval_model(**batch)
+        # CE (data) loss
+        ce_loss = outputs.loss
+        metrics = {"loss": ce_loss}
+        
+        # per CE loss
+        per_codebook_losses = outputs.per_codebook_losses
+        metrics.update({f"codebook_{i}_loss": l for (i,l) in enumerate(per_codebook_losses)})
+        return metrics
+
+    def generate_step(batch, accelerator):
+        batch.pop("decoder_attention_mask", None)
+        eval_model = accelerator.unwrap_model(model, keep_fp32_wrapper=True)
+        if training_args.torch_compile:
+            # if the model is compiled, we use the original model bc compile is not compatible with .generate
+            eval_model = model._orig_mod
+
+        # since we've might have loaded the weights in fp32, we have to autocast to ensure FA2 weights are in half-precision.
+        # with accelerator.autocast(autocast_handler=AutocastKwargs(enabled=(attn_implementation=="flash_attention_2"))):
+        output_audios = eval_model.generate(**batch, **gen_kwargs)
+        output_audios = accelerator.pad_across_processes(output_audios, dim=1, pad_index=0)
+        return output_audios
+
+    model.train()
+
+    total_batched_samples = resume_step if resume_step is not None else 0
+    for epoch in range(epochs_trained, num_epochs):
+        with accelerator.local_main_process_first():
+            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+        sampler = None
+        if training_args.group_by_length:
+            sampler = LengthGroupedSampler(train_batch_size, lengths=vectorized_datasets["train"]["target_length"])
+        train_dataloader = DataLoader(
+            vectorized_datasets["train"],
+            collate_fn=data_collator,
+            batch_size=per_device_train_batch_size,
+            sampler=sampler,
+            shuffle=not training_args.group_by_length,
+            num_workers=training_args.dataloader_num_workers,
+            pin_memory=training_args.dataloader_pin_memory,
+        )
+        train_dataloader = accelerator.prepare(train_dataloader)
+        if hasattr(train_dataloader, "dataset") and isinstance(train_dataloader.dataset, IterableDataset):
+            train_dataloader.dataset.set_epoch(epoch)
+
+        if resume_step is not None:
+            # Skip the first N batches in the dataloader when resuming from a checkpoint
+            logger.info(f"  Skip first {resume_step} batches")
+            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
+            resume_step = None
+            accelerator.wait_for_everyone()
+
+        # We chunkify the epoch iterator into gradient accumulation steps `n` batches
+        train_iterator = iter(train_dataloader)
+        num_steps_in_epoch = len(train_dataloader)
+        remainder = num_steps_in_epoch % gradient_accumulation_steps
+        remainder = remainder if remainder != 0 else gradient_accumulation_steps
+        total_updates = math.ceil(num_steps_in_epoch / gradient_accumulation_steps)
+        
+        update_step = -1
+        for _ in range(total_updates):
+            update_step += 1
+            
+            # preload the total batch per step
+            batch_samples = []
+            num_batches_in_step = gradient_accumulation_steps if update_step != (total_updates - 1) else remainder
+            for _ in range(num_batches_in_step):
+                batch_samples += [next(train_iterator)]
+                
+            # get num items in batch - if different than BOS and than -100
+            num_items_in_batch = sum([(batch["labels"].ne(audio_encoder_bos_token_id) | batch["labels"].ne(-100) | batch["labels"].ne(audio_encoder_eos_token_id)).sum((0,1))[0] for batch in batch_samples])
+            num_items_in_batch = accelerator.gather(num_items_in_batch).sum().item()
+            
+            # losses = []
+            for i,batch in enumerate(batch_samples):
+                total_batched_samples += 1
+                ctx = model.no_sync if (i < len(batch_samples) - 1 and accelerator.num_processes > 1) else contextlib.nullcontext
+                
+                with ctx():
+                    loss, train_metric = train_step(batch, accelerator, autocast_kwargs, num_items_in_batch, gradient_accumulation_steps)
+                    accelerator.backward(loss)
+                    # losses.append(loss.detach())
+            
+            grad_norm = accelerator.clip_grad_norm_(model.parameters(), training_args.max_grad_norm)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+            # The accelerator has performed an optimization step behind the scenes
+            steps_trained_progress_bar.update(1)
+            cur_step += 1
+
+            # losses = accelerator.gather(sum(losses)).sum().item() / (accelerator.num_processes * gradient_accumulation_steps)
+            
+            if cur_step % training_args.logging_steps == 0:
+                steps_trained_progress_bar.write(
+                    f"Step... ({cur_step} / {total_train_steps} | Loss:"
+                    f" {train_metric['loss']}, Learning Rate:"
+                    f" {lr_scheduler.get_last_lr()[0]})"
+                )
+                train_metric["grad_norm"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm
+                log_metric(
+                    accelerator,
+                    metrics=train_metric,
+                    learning_rate=lr_scheduler.get_last_lr()[0],
+                    train_time=train_time + time.time() - train_start,
+                    step=cur_step,
+                    epoch=epoch,
+                    prefix="train",
+                )
+
+            # save checkpoint and weights after each save_steps and at the end of training
+            if (cur_step % training_args.save_steps == 0) or cur_step == total_train_steps:
+                intermediate_dir = os.path.join(training_args.output_dir, f"checkpoint-{cur_step}-epoch-{epoch}")
+                # safe_serialization=False to avoid shared tensors saving issue (TODO(YL): it's a temporary fix)
+                # https://github.com/huggingface/transformers/issues/27293#issuecomment-1872560074
+                accelerator.save_state(output_dir=intermediate_dir, safe_serialization=False)
+                accelerator.wait_for_everyone()
+                if accelerator.is_main_process:
+                    rotate_checkpoints(
+                        training_args.save_total_limit, output_dir=training_args.output_dir, logger=logger
+                    )
+
+                    if cur_step == total_train_steps:
+                        # un-wrap student model for save
+                        unwrapped_model = accelerator.unwrap_model(model)
+                        unwrapped_model.save_pretrained(training_args.output_dir)
+
+                    if training_args.push_to_hub:
+                        api.upload_folder(
+                            repo_id=repo_id,
+                            folder_path=training_args.output_dir,
+                            commit_message=f"Saving train state of step {cur_step}",
+                            run_as_future=True,
+                        )
+                accelerator.wait_for_everyone()
+
+            if training_args.do_eval and (cur_step % eval_steps == 0 or cur_step == total_train_steps):
+                train_time += time.time() - train_start
+                # ======================== Evaluating ==============================
+                model.eval()
+                eval_metrics = []
+                eval_preds = []
+                eval_descriptions = []
+                eval_prompts = []
+                eval_start = time.time()
+
+                # release training input batch
+                batch = release_memory(batch)
+
+                validation_dataloader = DataLoader(
+                    vectorized_datasets["eval"],
+                    collate_fn=data_collator,
+                    batch_size=per_device_eval_batch_size,
+                    drop_last=False,
+                    num_workers=training_args.eval_dataloader_num_workers,
+                    pin_memory=training_args.dataloader_pin_memory,
+                )
+                validation_dataloader = accelerator.prepare(validation_dataloader)
+
+                for batch in tqdm(
+                    validation_dataloader,
+                    desc=f"Evaluating - Inference ...",
+                    position=2,
+                    disable=not accelerator.is_local_main_process,
+                ):
+                    # Model forward
+                    eval_metric = eval_step(batch, accelerator, autocast_kwargs)
+                    eval_metric = accelerator.gather_for_metrics(eval_metric)
+                    eval_metric = {key: val.unsqueeze(0) if val.ndim == 0 else val for (key,val) in eval_metric.items()}
+                    eval_metrics.append(eval_metric)
+
+                if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
+                    validation_dataloader = DataLoader(
+                        vectorized_datasets["eval"],
+                        collate_fn=data_collator,
+                        batch_size=per_device_eval_batch_size,
+                        drop_last=False,
+                        num_workers=training_args.eval_dataloader_num_workers,
+                        pin_memory=training_args.dataloader_pin_memory,
+                    )
+                    validation_dataloader = accelerator.prepare(validation_dataloader)
+                    # generation
+                    for batch in tqdm(
+                        validation_dataloader,
+                        desc=f"Evaluating - Generation ...",
+                        position=2,
+                        disable=not accelerator.is_local_main_process,
+                    ):
+                        generated_audios = generate_step(batch, accelerator)
+                        # Gather all predictions and targets
+                        generated_audios, input_ids, prompts = accelerator.pad_across_processes(
+                            (generated_audios, batch["input_ids"], batch["prompt_input_ids"]), dim=1, pad_index=0
+                        )
+                        generated_audios, input_ids, prompts = accelerator.gather_for_metrics(
+                            (generated_audios, input_ids, prompts)
+                        )
+                        eval_preds.extend(generated_audios.to("cpu"))
+                        eval_descriptions.extend(input_ids.to("cpu"))
+                        eval_prompts.extend(prompts.to("cpu"))
+
+                eval_time = time.time() - eval_start
+                # normalize eval metrics
+                eval_metrics = {
+                    key: torch.mean(torch.cat([d[key] for d in eval_metrics])).to("cpu") for key in eval_metrics[0]
+                }
+
+                # compute metrics
+                metrics_desc = ""
+                if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
+                    if accelerator.is_local_main_process:
+                        (
+                            metric_values,
+                            pred_descriptions,
+                            pred_prompts,
+                            audios,
+                            transcriptions,
+                            si_sdr_measures,
+                        ) = compute_metrics(
+                            eval_preds,
+                            eval_descriptions,
+                            eval_prompts,
+                            accelerator.device,
+                            training_args.compute_clap_similarity_metric,
+                            training_args.compute_noise_level_metric,
+                            training_args.noise_level_to_compute_clean_wer,
+                        )
+                        eval_metrics.update(metric_values)
+                        metrics_desc = " ".join([f"Eval {key}: {value} |" for key, value in metric_values.items()])
+                        if "wandb" in training_args.report_to:
+                            log_pred(
+                                accelerator,
+                                pred_descriptions,
+                                pred_prompts,
+                                transcriptions,
+                                audios,
+                                si_sdr_measures,
+                                sampling_rate=sampling_rate,
+                                step=cur_step,
+                                prefix="eval",
+                            )
+                    accelerator.wait_for_everyone()
+
+                # Print metrics and update progress bar
+                if accelerator.is_local_main_process:
+                    steps_trained_progress_bar.write(
+                        f"Eval results for step ({cur_step} / {total_train_steps} | Eval Loss: {eval_metrics['loss']} |"
+                        f" {metrics_desc})"
+                    )
+
+                log_metric(
+                    accelerator,
+                    metrics=eval_metrics,
+                    train_time=eval_time,
+                    step=cur_step,
+                    epoch=epoch,
+                    prefix="eval",
+                )
+
+                # release eval batch and relax metrics
+                eval_metrics, eval_preds, eval_descriptions, eval_prompts, batch, eval_metric = release_memory(
+                    eval_metrics, eval_preds, eval_descriptions, eval_prompts, batch, eval_metric
+                )
+                if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
+                    generated_audios, input_ids, prompts = release_memory(generated_audios, input_ids, prompts)
+
+                # train mode
+                model.train()
+
+                # flush the train metrics
+                train_start = time.time()
+
+            # break condition
+            if cur_step == total_train_steps:
+                continue_training = False
+                break
+
+        if not continue_training:
+            break
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main()

capspeech/ar/training/utils.py

@@ -0,0 +1,203 @@
+import os
+import re
+import shutil
+from dataclasses import field
+from pathlib import Path
+from typing import Dict, List
+
+import torch
+from datasets import concatenate_datasets, load_from_disk
+from wandb import Audio
+from datasets import load_from_disk, concatenate_datasets
+
+
+def list_field(default=None, metadata=None):
+    return field(default_factory=lambda: default, metadata=metadata)
+
+
+_RE_CHECKPOINT = re.compile(r"^checkpoint-(\d+)-epoch-(\d+)$")
+CHECKPOINT_CODEC_PREFIX = "checkpoint"
+_RE_CODEC_CHECKPOINT = re.compile(r"^checkpoint-(\d+)$")
+
+
+def get_last_checkpoint(folder):
+    content = os.listdir(folder)
+    checkpoints = [
+        path
+        for path in content
+        if _RE_CHECKPOINT.search(path) is not None and os.path.isdir(os.path.join(folder, path))
+    ]
+    if len(checkpoints) == 0:
+        return
+    return os.path.join(folder, max(checkpoints, key=lambda x: int(_RE_CHECKPOINT.search(x).groups()[0])))
+
+
+def sorted_checkpoints(output_dir=None, checkpoint_prefix="checkpoint") -> List[str]:
+    """Helper function to sort saved checkpoints from oldest to newest."""
+    ordering_and_checkpoint_path = []
+
+    glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}-*") if os.path.isdir(x)]
+
+    for path in glob_checkpoints:
+        regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path)
+        if regex_match is not None and regex_match.groups() is not None:
+            ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))
+
+    checkpoints_sorted = sorted(ordering_and_checkpoint_path)
+    checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
+    return checkpoints_sorted
+
+
+def rotate_checkpoints(save_total_limit=None, output_dir=None, checkpoint_prefix="checkpoint", logger=None) -> None:
+    """Helper function to delete old checkpoints."""
+    if save_total_limit is None or save_total_limit <= 0:
+        return
+    # Check if we should delete older checkpoint(s)
+    checkpoints_sorted = sorted_checkpoints(output_dir=output_dir, checkpoint_prefix=checkpoint_prefix)
+    if len(checkpoints_sorted) <= save_total_limit:
+        return
+
+    number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
+    checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
+    for checkpoint in checkpoints_to_be_deleted:
+        logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
+        shutil.rmtree(checkpoint, ignore_errors=True)
+
+
+def save_codec_checkpoint(output_dir, dataset, step):
+    checkpoint_path = f"{CHECKPOINT_CODEC_PREFIX}-{step}"
+    output_path = os.path.join(output_dir, checkpoint_path)
+    dataset.save_to_disk(output_path)
+
+
+def load_codec_checkpoint(checkpoint_path):
+    dataset = load_from_disk(checkpoint_path)
+    return dataset
+
+
+def sorted_codec_checkpoints(output_dir=None) -> List[str]:
+    """Helper function to sort saved checkpoints from oldest to newest."""
+    ordering_and_checkpoint_path = []
+
+    glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{CHECKPOINT_CODEC_PREFIX}-*")]
+
+    for path in glob_checkpoints:
+        regex_match = re.match(f".*{CHECKPOINT_CODEC_PREFIX}-([0-9]+)", path)
+        if regex_match is not None and regex_match.groups() is not None:
+            ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))
+
+    checkpoints_sorted = sorted(ordering_and_checkpoint_path)
+    checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
+    return checkpoints_sorted
+
+
+def load_all_codec_checkpoints(output_dir=None) -> List[str]:
+    """Helper function to load and concat all checkpoints."""
+    checkpoints_sorted = sorted_codec_checkpoints(output_dir=output_dir)
+    datasets = [load_from_disk(checkpoint) for checkpoint in checkpoints_sorted]
+    datasets = concatenate_datasets(datasets, axis=0)
+    return datasets
+
+
+def get_last_codec_checkpoint_step(folder) -> int:
+    if not os.path.exists(folder) or not os.path.isdir(folder):
+        os.makedirs(folder, exist_ok=True)
+        return 0
+    content = os.listdir(folder)
+    checkpoints = [path for path in content if _RE_CODEC_CHECKPOINT.search(path) is not None]
+    if len(checkpoints) == 0:
+        return 0
+    last_checkpoint = os.path.join(
+        folder, max(checkpoints, key=lambda x: int(_RE_CODEC_CHECKPOINT.search(x).groups()[0]))
+    )
+    # Find num steps saved state string pattern
+    pattern = r"checkpoint-(\d+)"
+    match = re.search(pattern, last_checkpoint)
+    cur_step = int(match.group(1))
+    return cur_step
+
+
+def log_metric(
+    accelerator,
+    metrics: Dict,
+    train_time: float,
+    step: int,
+    epoch: int,
+    learning_rate: float = None,
+    prefix: str = "train",
+):
+    """Helper function to log all training/evaluation metrics with the correct prefixes and styling."""
+    log_metrics = {}
+    for k, v in metrics.items():
+        if "codebook" in k:
+            log_metrics[f"codebook_{prefix}/{k}"] = v
+        else:
+            log_metrics[f"{prefix}/{k}"] = v
+    log_metrics[f"{prefix}/time"] = train_time
+    log_metrics[f"{prefix}/epoch"] = epoch
+    if learning_rate is not None:
+        log_metrics[f"{prefix}/learning_rate"] = learning_rate
+    accelerator.log(log_metrics, step=step)
+
+
+def log_pred(
+    accelerator,
+    pred_descriptions: List[str],
+    pred_prompts: List[str],
+    transcriptions: List[str],
+    audios: List[torch.Tensor],
+    si_sdr_measures: List[float],
+    sampling_rate: int,
+    step: int,
+    prefix: str = "eval",
+    num_lines: int = 200000,
+):
+    """Helper function to log target/predicted transcriptions to weights and biases (wandb)."""
+    if accelerator.is_main_process:
+        wandb_tracker = accelerator.get_tracker("wandb")
+        # pretty name for current step: step 50000 -> step 50k
+        cur_step_pretty = f"{int(step // 1000)}k" if step > 1000 else step
+        prefix_pretty = prefix.replace("/", "-")
+
+        if si_sdr_measures is None:
+            # convert str data to a wandb compatible format
+            str_data = [
+                [pred_descriptions[i], pred_prompts[i], transcriptions[i]] for i in range(len(pred_descriptions))
+            ]
+            # log as a table with the appropriate headers
+            wandb_tracker.log_table(
+                table_name=f"predictions/{prefix_pretty}-step-{cur_step_pretty}",
+                columns=["Target descriptions", "Target prompts", "Predicted transcriptions"],
+                data=str_data[:num_lines],
+                step=step,
+                commit=False,
+            )
+        else:
+            # convert str data to a wandb compatible format
+            str_data = [
+                [pred_descriptions[i], pred_prompts[i], transcriptions[i], si_sdr_measures[i]]
+                for i in range(len(pred_descriptions))
+            ]
+            # log as a table with the appropriate headers
+            wandb_tracker.log_table(
+                table_name=f"predictions/{prefix_pretty}-step-{cur_step_pretty}",
+                columns=["Target descriptions", "Target prompts", "Predicted transcriptions", "Noise estimation"],
+                data=str_data[:num_lines],
+                step=step,
+                commit=False,
+            )
+
+        # wandb can only loads 100 audios per step
+        wandb_tracker.log(
+            {
+                "Speech samples": [
+                    Audio(
+                        audio,
+                        caption=f"{pred_prompts[i]} --- DESCRIPTION: {pred_descriptions[i]}",
+                        sample_rate=sampling_rate,
+                    )
+                    for (i, audio) in enumerate(audios[: min(len(audios), 100)])
+                ]
+            },
+            step=step,
+        )

capspeech/eval/README.md

@@ -0,0 +1,42 @@
+# CapSpeech Evaluation Tools
+
+## Get Start
+Install dependicies:
+```bash
+conda create -n capeval python=3.9
+conda activate capeval
+pip install -r requirements.txt
+pip install git+https://github.com/sarulab-speech/UTMOSv2.git
+```
+
+For ASR, we need:
+```bash
+conda install ffmpeg
+```
+
+## Evaluate pitch, monotony, speed, age, gender
+RUN:
+```bash
+python base_eval.py
+```
+
+## Evaluate UTMOSv2
+RUN:
+```bash
+python mos_eval.py
+```
+
+## Evaluate ASR Results
+RUN:
+```bash
+python asr_eval.py
+```
+
+## Evaluate emotion, accent
+RUN:
+```bash
+cd src/example/
+python categorized_emotion.py
+python dialect_world_dialect.py
+```
+Please refer to [Vox-profile](https://github.com/tiantiaf0627/vox-profile-release.git) for more evaluation tools.

capspeech/eval/age_gender.py

@@ -0,0 +1,35 @@
+import audeer
+import audonnx
+import numpy as np
+
+def age_gender_apply(waveform):
+    age_labels = ['child', 'teenager', 'young adult', 'middle-aged adult', 'elderly']
+    gender_labels = ['female', 'male']
+    url = 'https://zenodo.org/record/7761387/files/w2v2-L-robust-6-age-gender.25c844af-1.1.1.zip'
+    cache_root = audeer.mkdir('cache')
+    model_root = audeer.mkdir('model')
+    sampling_rate = 16000
+    archive_path = audeer.download_url(url, cache_root, verbose=True)
+    audeer.extract_archive(archive_path, model_root)
+    model = audonnx.load(model_root)
+
+    result = model(waveform, sampling_rate)
+    # Process age
+    age_label = result['logits_age'].squeeze() * 100.0
+    if age_label <= 12:
+        age_label = 'child'
+    elif age_label <= 19:
+        age_label = 'teenager'
+    elif age_label <= 39:
+        age_label = 'young adult'
+    elif age_label <= 64:
+        age_label = 'middle-aged adult'
+    else:
+        age_label = 'elderly'
+
+    # Process gender
+    gender_label = result['logits_gender'].squeeze()
+    gender_label = gender_label[:2]  # Remove child
+    gender_label = np.argmax(gender_label)
+
+    return age_label, gender_labels[gender_label]

capspeech/eval/asr_eval.py

@@ -0,0 +1,24 @@
+from jiwer import wer as calculate_wer
+from jiwer import cer as calculate_cer
+from whisper.normalizers import EnglishTextNormalizer
+import whisper
+import torch
+
+normalizer = EnglishTextNormalizer()
+device = "cuda" if torch.cuda.is_available() else "cpu"
+whisper_model = whisper.load_model("large-v3-turbo", device=device)
+
+def asr(wav_path):
+    result = whisper_model.transcribe(wav_path)
+    pred = result['text'].strip()
+    pred = normalizer(pred)
+    return pred
+
+if __name__ == '__main__':
+    gt_text="Hey, how are you doing today? I like it."
+    wav_path="your-audio"
+    gt_text = normalizer(gt_text.strip())
+    pred_asr = asr(wav_path)
+    wer = round(calculate_wer(gt_text, pred_asr), 3)
+    cer = round(calculate_cer(gt_text, pred_asr), 3)
+    print(wer, cer)

capspeech/eval/base_eval.py

@@ -0,0 +1,32 @@
+from pitch import pitch_apply
+from speed import speed_apply
+from age_gender import age_gender_apply
+import librosa
+import json
+import bisect
+
+SPEAKER_RATE_BINS = ["very slowly", "slowly", "slightly slowly", "moderate speed", "slightly fast", "fast", "very fast"]
+UTTERANCE_LEVEL_STD = ["very monotone", "monotone", "slightly expressive and animated", "expressive and animated", "very expressive and animated"]
+SPEAKER_LEVEL_PITCH_BINS = ["very low-pitch", "low-pitch", "slightly low-pitch", "moderate pitch", "slightly high-pitch", "high-pitch", "very high-pitch"]
+with open("bin.json") as json_file:
+    text_bins_dict = json.load(json_file)
+
+audiopath = "YOUR_AUDIO_PATH"
+waveform, _ = librosa.load(audiopath, sr=16000)
+age, gender = age_gender_apply(waveform)
+pitch_mean, pitch_std = pitch_apply(waveform)
+if gender == "male":
+    index = bisect.bisect_right(text_bins_dict["pitch_bins_male"], pitch_mean) - 1
+    pitch = SPEAKER_LEVEL_PITCH_BINS[index]
+else:
+    index = bisect.bisect_right(text_bins_dict["pitch_bins_female"], pitch_mean) - 1
+    pitch = SPEAKER_LEVEL_PITCH_BINS[index]
+
+index = bisect.bisect_right(text_bins_dict["speech_monotony"], pitch_std) - 1
+monotony = UTTERANCE_LEVEL_STD[index]
+speech_duration = speed_apply(waveform)
+
+index = bisect.bisect_right(text_bins_dict["speaking_rate"], speech_duration) - 1
+speed = SPEAKER_RATE_BINS[index]
+
+print(pitch, monotony, speed, age, gender)

capspeech/eval/bin.json

@@ -0,0 +1,10 @@
+{
+    "speaking_rate": [0.0, 3.8258038258038254, 7.651607651607651, 11.477411477411476, 15.303215303215302, 19.129019129019127, 22.95482295482295, 26.78062678062678], 
+    "noise": [17.12751579284668, 25.4012325831822, 33.67494937351772, 41.94866616385323, 50.22238295418875, 58.49609974452427, 66.76981653485979, 75.04353332519531], 
+    "reverberation": [10, 35, 45, 55, 59, 60], 
+    "speech_monotony": [0.0, 20.37920924595424, 40.75841849190848, 70, 90, 142.6544647216797], 
+    "pitch_bins_male": [64.6531982421875, 81.66683959960938, 98.68048095703125, 115.69412231445312, 132.707763671875, 149.72140502929688, 166.73504638671875, 183.74868774414062], 
+    "pitch_bins_female": [120.17855072021484, 141.6242690945264, 163.06998746883795, 184.51570584314953, 205.96142421746106, 227.40714259177264, 248.8528609660842, 270.29857934039575], 
+    "si-sdr": [-17.804332733154297, -0.40644073486328125, 10, 20, 25, 28, 34.38934326171875], 
+    "pesq": [1, 1.7, 2.4, 3.1, 3.6, 4, 4.499948978424072]
+}

capspeech/eval/pitch.py

@@ -0,0 +1,30 @@
+import torch 
+import penn
+
+def pitch_apply(waveform):
+    hopsize = .01
+    fmin = 30.
+    fmax = 1000.
+    checkpoint = None
+    center = 'half-hop'
+    interp_unvoiced_at = .065
+    sampling_rate = 16000
+    penn_batch_size = 4096
+    waveform = torch.Tensor(waveform).unsqueeze(0)
+    pitch, periodicity = penn.from_audio(
+        waveform.float(),
+        sampling_rate,
+        hopsize=hopsize,
+        fmin=fmin,
+        fmax=fmax,
+        checkpoint=checkpoint,
+        batch_size=penn_batch_size,
+        center=center,
+        interp_unvoiced_at=interp_unvoiced_at,
+        gpu=None
+        )     
+    
+    pitch_mean = pitch.mean().cpu().numpy()
+    pitch_std = pitch.std().cpu().numpy()
+
+    return pitch_mean, pitch_std

capspeech/eval/requirements.txt

@@ -0,0 +1,16 @@
+datasets[audio]
+https://github.com/marianne-m/brouhaha-vad/archive/main.zip
+penn
+g2p
+demucs
+transformers
+bitsandbytes
+git+https://github.com/sarulab-speech/UTMOSv2.git
+-U openai-whisper
+jiwer
+numpy==1.26.4
+audeer
+audonnx
+laion_clap
+numpy==1.26.4
+onnxruntime

capspeech/eval/speed.py

@@ -0,0 +1,29 @@
+from pyannote.audio import Model
+from pathlib import Path
+from brouhaha.pipeline import RegressiveActivityDetectionPipeline
+import torch 
+from huggingface_hub import hf_hub_download
+import numpy as np
+
+def speed_apply(waveform):
+    ratio = 16000/270
+    sampling_rate = 16000
+    device = "cpu"
+    waveform = torch.Tensor(waveform).unsqueeze(0)
+    model = Model.from_pretrained(
+            Path(hf_hub_download(repo_id="ylacombe/brouhaha-best", filename="best.ckpt")),
+            strict=False,
+        )
+    model.to(device)
+
+    pipeline = RegressiveActivityDetectionPipeline(segmentation=model, batch_size=1)
+    pipeline.to(torch.device(device))
+
+    device = pipeline._models["segmentation"].device
+
+    res = pipeline({"sample_rate": sampling_rate,
+                    "waveform": waveform.to(device).float()})
+
+    speech_duration = sum(map(lambda x: x[0].duration, res["annotation"].itertracks()))     
+        
+    return speech_duration

capspeech/eval/src/example/categorized_emotion.py

@@ -0,0 +1,92 @@
+import torch
+import logging
+import sys, os, pdb
+import torch.nn.functional as F
+
+from pathlib import Path
+
+sys.path.append(os.path.join(str(Path(os.path.realpath(__file__)).parents[1])))
+sys.path.append(os.path.join(str(Path(os.path.realpath(__file__)).parents[1]), 'model', 'emotion'))
+
+from wavlm_emotion import WavLMWrapper
+from whisper_emotion import WhisperWrapper
+
+
+# define logging console
+import logging
+logging.basicConfig(
+    format='%(asctime)s %(levelname)-3s ==> %(message)s', 
+    level=logging.INFO, 
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+
+os.environ["MKL_NUM_THREADS"] = "1" 
+os.environ["NUMEXPR_NUM_THREADS"] = "1" 
+os.environ["OMP_NUM_THREADS"] = "1" 
+
+
+if __name__ == '__main__':
+
+    label_list = [
+        'Anger', 
+        'Contempt', 
+        'Disgust', 
+        'Fear', 
+        'Happiness', 
+        'Neutral', 
+        'Sadness', 
+        'Surprise', 
+        'Other'
+    ]
+
+    # Find device
+    device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+    if torch.cuda.is_available(): print('GPU available, use GPU')
+
+    # Define the model
+    # Note that ensemble yields the better performance than the single model
+    # Define the model wrapper
+    model_path = "model"
+    wavlm_model = model = WavLMWrapper(
+        pretrain_model="wavlm_large", 
+        finetune_method="finetune",
+        output_class_num=9,
+        freeze_params=True, 
+        use_conv_output=True,
+        detailed_class_num=17
+    ).to(device)
+    
+    whisper_model = WhisperWrapper(
+        pretrain_model="whisper_large", 
+        finetune_method="lora",
+        lora_rank=16, 
+        output_class_num=9,
+        freeze_params=True, 
+        use_conv_output=True,
+        detailed_class_num=17
+    ).to(device)
+        
+    whisper_model.load_state_dict(torch.load(os.path.join(model_path, f"whisper_emotion.pt"), weights_only=True), strict=False)
+    whisper_model.load_state_dict(torch.load(os.path.join(model_path, f"whisper_emotion_lora.pt")), strict=False)
+    wavlm_model.load_state_dict(torch.load(os.path.join(model_path, f"wavlm_emotion.pt"), weights_only=True), strict=False)
+
+    wavlm_model.eval()
+    whisper_model.eval()
+    
+    # Audio must be 16k Hz
+    data = torch.zeros([1, 16000]).to(device)
+    whisper_logits, whisper_embedding, _, _, _, _   = whisper_model(
+        data, return_feature=True
+    )
+    wavlm_logits, wavlm_embedding, _, _, _, _       = wavlm_model(
+        data, return_feature=True
+    )
+    
+    ensemble_logits = (whisper_logits + wavlm_logits) / 2
+    ensemble_prob   = F.softmax(ensemble_logits, dim=1)
+
+    print(ensemble_prob.shape)
+    print(whisper_embedding.shape)
+    print(wavlm_embedding.shape)
+    print(label_list[torch.argmax(ensemble_prob).detach().cpu().item()])
+

capspeech/eval/src/example/dialect_world_dialect.py

@@ -0,0 +1,87 @@
+import torch
+import sys, os, pdb
+import argparse, logging
+import torch.nn.functional as F
+
+from pathlib import Path
+
+
+sys.path.append(os.path.join(str(Path(os.path.realpath(__file__)).parents[1])))
+sys.path.append(os.path.join(str(Path(os.path.realpath(__file__)).parents[1]), 'model', 'dialect'))
+
+from wavlm_dialect import WavLMWrapper
+from whisper_dialect import WhisperWrapper
+
+
+# define logging console
+import logging
+logging.basicConfig(
+    format='%(asctime)s %(levelname)-3s ==> %(message)s', 
+    level=logging.INFO, 
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+
+os.environ["MKL_NUM_THREADS"] = "1" 
+os.environ["NUMEXPR_NUM_THREADS"] = "1" 
+os.environ["OMP_NUM_THREADS"] = "1" 
+
+
+if __name__ == '__main__':
+
+   
+    label_list = [
+        'East Asia', 'English', 'Germanic', 'Irish', 
+        'North America', 'Northern Irish', 'Oceania', 
+        'Other', 'Romance', 'Scottish', 'Semitic', 'Slavic', 
+        'South African', 'Southeast Asia', 'South Asia', 'Welsh'
+    ]
+    
+    # Find device
+    device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+    if torch.cuda.is_available(): print('GPU available, use GPU')
+
+    # Define the model
+    # Note that ensemble yields the better performance than the single model
+    model_path = "YOUR_PATH"
+    # Define the model wrapper
+    wavlm_model = model = WavLMWrapper(
+        pretrain_model="wavlm_large", 
+        finetune_method="lora",
+        lora_rank=16,
+        output_class_num=16,
+        freeze_params=False, 
+        use_conv_output=True,
+        apply_gradient_reversal=False, 
+        num_dataset=3
+    ).to(device)
+    
+    whisper_model = WhisperWrapper(
+        pretrain_model="whisper_large", 
+        finetune_method="lora",
+        lora_rank=16,
+        output_class_num=16,
+        freeze_params=False, 
+        use_conv_output=True,
+        apply_gradient_reversal=False, 
+        num_dataset=11
+    ).to(device)
+    
+    wavlm_model.load_state_dict(torch.load(os.path.join(model_path, f"wavlm_world_dialect.pt"), weights_only=True), strict=False)
+    wavlm_model.load_state_dict(torch.load(os.path.join(model_path, f"wavlm_world_dialect_lora.pt")), strict=False)
+    
+    whisper_model.load_state_dict(torch.load(os.path.join(model_path, f"whisper_world_dialect.pt"), weights_only=True), strict=False)
+    whisper_model.load_state_dict(torch.load(os.path.join(model_path, f"whisper_world_dialect_lora.pt")), strict=False)
+
+    wavlm_model.eval()
+    whisper_model.eval()
+        
+    data = torch.zeros([1, 16000]).to(device)
+    wavlm_logits, wavlm_embeddings      = wavlm_model(data, return_feature=True)
+    whisper_logits, whisper_embeddings  = whisper_model(data, return_feature=True)
+
+    ensemble_logits = (wavlm_logits + whisper_logits) / 2
+    ensemble_prob   = F.softmax(ensemble_logits, dim=1)
+
+    pred = label_list[ensemble_prob.argmax(-1)]
+    print(pred)
+

capspeech/eval/src/model/adapter.py

@@ -0,0 +1,73 @@
+# --------------------------------------------------------
+# References:
+# https://github.com/jxhe/unify-parameter-efficient-tuning
+# --------------------------------------------------------
+
+import math
+import torch
+import torch.nn as nn
+
+
+class Adapter(nn.Module):
+    def __init__(
+        self,
+        config=None,
+        d_model=768,
+        bottleneck=None,
+        dropout=0.0,
+        init_option="lora",
+        adapter_scalar="1.0",
+        adapter_layernorm_option="none"
+    ):
+        super().__init__()
+        self.n_embd = config.d_model if d_model is None else d_model
+        self.down_size = config.attn_bn if bottleneck is None else bottleneck
+
+        #_before
+        self.adapter_layernorm_option = adapter_layernorm_option
+
+        self.adapter_layer_norm_before = None
+        if adapter_layernorm_option == "in" or adapter_layernorm_option == "out":
+            self.adapter_layer_norm_before = nn.LayerNorm(self.n_embd)
+
+        if adapter_scalar == "learnable_scalar":
+            self.scale = nn.Parameter(torch.ones(1))
+        else:
+            self.scale = float(adapter_scalar)
+
+        self.down_proj = nn.Linear(self.n_embd, self.down_size)
+        self.non_linear_func = nn.ReLU()
+        self.up_proj = nn.Linear(self.down_size, self.n_embd)
+
+        self.dropout = dropout
+        if init_option == "bert":
+            raise NotImplementedError
+        elif init_option == "lora":
+            with torch.no_grad():
+                nn.init.kaiming_uniform_(self.down_proj.weight, a=math.sqrt(5))
+                nn.init.zeros_(self.up_proj.weight)
+                nn.init.zeros_(self.down_proj.bias)
+                nn.init.zeros_(self.up_proj.bias)
+
+    def forward(self, x, add_residual=True, residual=None):
+        residual = x if residual is None else residual
+        if self.adapter_layernorm_option == 'in':
+            x = self.adapter_layer_norm_before(x)
+        
+        down = self.down_proj(x)
+        
+        down = self.non_linear_func(down)
+        down = nn.functional.dropout(down, p=self.dropout, training=self.training)
+        up = self.up_proj(down)
+        
+        up = up * self.scale
+
+        if self.adapter_layernorm_option == 'out':
+            up = self.adapter_layer_norm_before(up)
+
+        if add_residual:
+            output = up + residual
+        else:
+            output = up
+
+        return output

capspeech/eval/src/model/dialect/wavlm_dialect.py

@@ -0,0 +1,300 @@
+import os
+import pdb
+import copy
+import torch
+import argparse
+import loralib as lora
+import transformers.models.wavlm.modeling_wavlm as wavlm
+from speechbrain.nnet.normalization import LayerNorm
+from speechbrain.lobes.models.huggingface_transformers.huggingface import make_padding_masks
+
+from torch import nn
+from torch.nn import functional as F
+from transformers import Wav2Vec2FeatureExtractor
+from transformers import WavLMModel
+
+import sys
+from pathlib import Path
+sys.path.append(os.path.join(str(Path(os.path.realpath(__file__)).parents[1])))
+from revgrad import RevGrad
+
+class WavLMEncoderLayer(nn.Module):
+    def __init__(self, layer_idx, config, has_relative_position_bias: bool = True):
+        super().__init__()
+        self.attention = wavlm.WavLMAttention(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            dropout=config.attention_dropout,
+            num_buckets=config.num_buckets,
+            max_distance=config.max_bucket_distance,
+            has_relative_position_bias=has_relative_position_bias,
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.feed_forward = wavlm.WavLMFeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.config = config
+        
+        if layer_idx > config.num_hidden_layers // 2:
+            if self.config.finetune_method == "lora" or self.config.finetune_method == "combined":
+                self.feed_forward.intermediate_dense    = lora.Linear(config.hidden_size, config.intermediate_size, r=config.lora_rank)
+                self.feed_forward.output_dense          = lora.Linear(config.intermediate_size, config.hidden_size, r=config.lora_rank)
+
+    def forward(self, hidden_states, attention_mask=None, position_bias=None, output_attentions=False, index=0):
+        
+        attn_residual = hidden_states
+        hidden_states, attn_weights, position_bias = self.attention(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_bias=position_bias,
+            output_attentions=output_attentions,
+            index=index,
+        )
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = attn_residual + hidden_states
+        
+        hidden_states = self.layer_norm(hidden_states)
+        hidden_states = hidden_states + self.feed_forward(hidden_states)
+        hidden_states = self.final_layer_norm(hidden_states)
+        outputs = (hidden_states, position_bias)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+class WavLMEncoderLayerStableLayerNorm(nn.Module):
+    def __init__(self, layer_idx, config, has_relative_position_bias: bool = True):
+        super().__init__()
+        self.attention = wavlm.WavLMAttention(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            dropout=config.attention_dropout,
+            num_buckets=config.num_buckets,
+            max_distance=config.max_bucket_distance,
+            has_relative_position_bias=has_relative_position_bias,
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.feed_forward = wavlm.WavLMFeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.config = config
+
+        if layer_idx > config.num_hidden_layers // 2:
+            if self.config.finetune_method == "lora" or self.config.finetune_method == "combined":
+                self.feed_forward.intermediate_dense    = lora.Linear(config.hidden_size, config.intermediate_size, r=config.lora_rank)
+                self.feed_forward.output_dense          = lora.Linear(config.intermediate_size, config.hidden_size, r=config.lora_rank)
+            
+
+    def forward(self, hidden_states, attention_mask=None, position_bias=None, output_attentions=False):
+        attn_residual = hidden_states
+        hidden_states = self.layer_norm(hidden_states)
+        hidden_states, attn_weights, position_bias = self.attention(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_bias=position_bias,
+            output_attentions=output_attentions,
+        )
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = attn_residual + hidden_states
+        hidden_states = hidden_states + self.feed_forward(self.final_layer_norm(hidden_states))
+
+        outputs = (hidden_states, position_bias)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+   
+class WavLMWrapper(nn.Module):
+    def __init__(
+        self, 
+        pretrain_model="wavlm_large", 
+        hidden_dim=256,
+        finetune_method="lora",
+        lora_rank=16,
+        freeze_params=True,
+        output_class_num=4,
+        use_conv_output=True,
+        apply_gradient_reversal=False,
+        num_dataset=4
+    ):
+        super(WavLMWrapper, self).__init__()
+        # 1. We Load the model first with weights
+        if pretrain_model == "wavlm":
+            self.backbone_model = WavLMModel.from_pretrained(
+                "microsoft/wavlm-base-plus",
+                output_hidden_states=True,
+            )
+        elif pretrain_model == "wavlm_large":
+            self.processor = Wav2Vec2FeatureExtractor.from_pretrained('microsoft/wavlm-large')
+            self.backbone_model = WavLMModel.from_pretrained(
+                "microsoft/wavlm-large",
+                output_hidden_states=True,
+            )
+        self.pretrain_model             = pretrain_model
+        self.finetune_method            = finetune_method
+        self.apply_gradient_reversal    = apply_gradient_reversal
+        self.use_conv_output            = use_conv_output
+        
+        state_dict = self.backbone_model.state_dict()
+        # 2. Read the model config
+        self.model_config = self.backbone_model.config
+        self.model_config.finetune_method        = finetune_method
+        self.model_config.lora_rank              = lora_rank
+        
+        # 3. Config encoder layers with adapter or embedding prompt
+        if self.pretrain_model == "wavlm":
+            self.backbone_model.encoder.layers = nn.ModuleList(
+                [WavLMEncoderLayer(i, self.model_config, has_relative_position_bias=(i == 0)) for i in range(self.model_config.num_hidden_layers)]
+            )
+        elif self.pretrain_model == "wavlm_large":
+            self.backbone_model.encoder.layers = nn.ModuleList(
+                [WavLMEncoderLayerStableLayerNorm(i, self.model_config, has_relative_position_bias=(i == 0)) for i in range(self.model_config.num_hidden_layers)]
+            )
+        # 4. Load the weights back
+        msg = self.backbone_model.load_state_dict(state_dict, strict=False)
+
+        # 5. Freeze the weights
+        self.freeze_params = freeze_params
+        if self.freeze_params and self.finetune_method != "lora":
+            for _, p in self.backbone_model.named_parameters(): p.requires_grad = False
+        elif self.freeze_params and self.finetune_method == "lora":
+            for name, p in self.backbone_model.named_parameters():
+                if name in msg.missing_keys: p.requires_grad = True
+                else: p.requires_grad = False
+        else:
+            for _, p in self.backbone_model.named_parameters(): p.requires_grad = True
+
+        # 6. Downstream models
+        self.model_seq = nn.Sequential(
+            nn.Conv1d(self.model_config.hidden_size, hidden_dim, 1, padding=0),
+            nn.ReLU(),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(hidden_dim, hidden_dim, 1, padding=0),
+            nn.ReLU(),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(hidden_dim, hidden_dim, 1, padding=0)
+        )
+
+        if self.use_conv_output:
+            num_layers = self.model_config.num_hidden_layers + 1  # transformer layers + input embeddings
+            self.weights = nn.Parameter(torch.ones(num_layers)/num_layers)
+        else:
+            num_layers = self.model_config.num_hidden_layers
+            self.weights = nn.Parameter(torch.zeros(num_layers))
+        
+        if apply_gradient_reversal:
+            self.dataset_layer = nn.Sequential(
+                RevGrad(),
+                nn.Linear(hidden_dim, hidden_dim),
+                nn.ReLU(),
+                nn.Linear(hidden_dim, num_dataset),
+            )
+        
+        self.out_layer = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, output_class_num),
+        )
+        
+    def forward(self, x, length=None, return_feature=False):
+        # 1. feature extraction and projections
+        if self.pretrain_model == "wavlm_large":  
+            with torch.no_grad():
+                signal, attention_mask = list(), list()
+                if length is not None: attention_mask = make_padding_masks(x, wav_len=length/length.max()).to(x.device)
+                else: attention_mask = make_padding_masks(x, wav_len=torch.tensor([1]).to(x.device)).to(x.device)
+
+                for idx in range(len(x)):
+                    input = self.processor(x[idx], sampling_rate=16_000, return_tensors="pt", padding=True)
+                    signal.append(input["input_values"][0].to(x.device))
+                signal = torch.stack(signal)
+
+        # 2. get length and mask
+        if length is not None:
+            length = self.get_feat_extract_output_lengths(length.detach().cpu())
+            length = length.cuda()
+
+        if self.pretrain_model == "wavlm": 
+            x = self.backbone_model(
+                x, output_hidden_states=True
+            ).hidden_states
+        else: 
+            x = self.backbone_model(
+                signal, 
+                attention_mask=attention_mask, 
+                output_hidden_states=True
+            ).hidden_states
+        
+        # 4. stacked feature
+        if self.use_conv_output: stacked_feature = torch.stack(x, dim=0)
+        else: stacked_feature = torch.stack(x, dim=0)[1:]
+        
+        # 5. Weighted sum
+        _, *origin_shape = stacked_feature.shape
+        # Return transformer enc outputs [num_enc_layers, B, T, D]
+        if self.use_conv_output:
+            stacked_feature = stacked_feature.view(self.backbone_model.config.num_hidden_layers+1, -1)
+        else:
+            stacked_feature = stacked_feature.view(self.backbone_model.config.num_hidden_layers, -1)
+        norm_weights = F.softmax(self.weights, dim=-1)
+        
+        # Perform weighted average
+        weighted_feature = (norm_weights.unsqueeze(-1) * stacked_feature).sum(dim=0)
+        features = weighted_feature.view(*origin_shape)
+        
+        # 6. Pass the weighted average to point-wise 1D Conv
+        # B x T x D
+        features = features.transpose(1, 2)
+        features = self.model_seq(features)
+        features = features.transpose(1, 2)
+        
+        # 7. Pooling
+        if length is not None:
+            mean, std = list(), list()
+            for snt_id in range(features.shape[0]):
+                # Avoiding padded time steps
+                actual_size = length[snt_id]
+                mean.append(torch.mean(features[snt_id, 0:actual_size, ...], dim=0))
+            features = torch.stack(mean)
+        else:
+            features = torch.mean(features, dim=1)
+
+        # 8. Output predictions
+        # B x D
+        predicted = self.out_layer(features)
+        if self.apply_gradient_reversal: 
+            dataset_predicted = self.dataset_layer(features)
+            if return_feature: return predicted, dataset_predicted, features
+            return predicted, dataset_predicted
+        if return_feature: return predicted, features
+        return predicted
+    
+    # From huggingface
+    def get_feat_extract_output_lengths(self, input_length):
+        """
+        Computes the output length of the convolutional layers
+        """
+        def _conv_out_length(input_length, kernel_size, stride):
+            # 1D convolutional layer output length formula taken
+            # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html
+            return (input_length - kernel_size) // stride + 1
+        for kernel_size, stride in zip(self.backbone_model.config.conv_kernel, self.backbone_model.config.conv_stride):
+            input_length = _conv_out_length(input_length, kernel_size, stride)
+        return input_length
+
+def prepare_mask(length, shape, dtype):
+    # Modified from huggingface
+    mask = torch.zeros(
+        shape, dtype=dtype
+    )
+    # these two operations makes sure that all values
+    # before the output lengths indices are attended to
+    mask[(torch.arange(mask.shape[0]), length.cpu() - 1)] = 1
+    mask = mask.flip([-1]).cumsum(-1).flip([-1]).bool()
+    return mask
+    
+    

capspeech/eval/src/model/dialect/whisper_dialect.py

@@ -0,0 +1,301 @@
+import os
+import pdb
+import copy
+import torch
+import argparse
+import numpy as np
+import loralib as lora
+import transformers.models.whisper.modeling_whisper as whisper
+
+from torch import nn
+from torch.nn import functional as F
+from transformers.activations import ACT2FN
+from transformers import WhisperModel, AutoFeatureExtractor
+
+import sys
+from pathlib import Path
+sys.path.append(os.path.join(str(Path(os.path.realpath(__file__)).parents[1])))
+from revgrad import RevGrad
+
+class WhisperEncoderLayer(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.embed_dim = config.d_model
+        self.self_attn = whisper.WhisperAttention(
+            embed_dim=self.embed_dim,
+            num_heads=config.encoder_attention_heads,
+            dropout=config.attention_dropout,
+        )
+        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.activation_function]
+        self.activation_dropout = config.activation_dropout
+        self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.config = config
+        
+        if layer_idx > config.encoder_layers // 2:
+            if self.config.finetune_method == "lora" or self.config.finetune_method == "combined":
+                self.fc1 = lora.Linear(self.embed_dim, config.encoder_ffn_dim, r=config.lora_rank)
+                self.fc2 = lora.Linear(config.encoder_ffn_dim, self.embed_dim, r=config.lora_rank)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        layer_head_mask: torch.Tensor,
+        output_attentions: bool = False,
+    ):
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(seq_len, batch, embed_dim)`
+            attention_mask (`torch.FloatTensor`): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
+                `(encoder_attention_heads,)`.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        residual = hidden_states
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+        hidden_states, attn_weights, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            layer_head_mask=layer_head_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.activation_fn(self.fc1(hidden_states))
+        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+        
+        if hidden_states.dtype == torch.float16 and (
+            torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()
+        ):
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+   
+class WhisperWrapper(nn.Module):
+    def __init__(
+        self, 
+        pretrain_model="whisper_large",
+        output_class_num=4, 
+        hidden_dim=256, 
+        finetune_method="lora",
+        lora_rank=16,
+        freeze_params=True,
+        use_conv_output=True,
+        apply_gradient_reversal=False, 
+        num_dataset=4
+    ):
+        super(WhisperWrapper, self).__init__()
+        # 1. We Load the model first with weights
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-tiny", chunk_length=15)
+        self.pretrain_model = pretrain_model
+        if self.pretrain_model == "whisper_tiny":
+            self.backbone_model = WhisperModel.from_pretrained(
+                "openai/whisper-tiny",
+                output_hidden_states=True,
+                ignore_mismatched_sizes=True,
+                max_source_positions=750,
+            )
+        elif self.pretrain_model == "whisper_base":
+            self.backbone_model = WhisperModel.from_pretrained(
+                "openai/whisper-base",
+                output_hidden_states=True,
+                ignore_mismatched_sizes=True,
+                max_source_positions=750,
+            )
+        elif self.pretrain_model == "whisper_small":
+            self.backbone_model = WhisperModel.from_pretrained(
+                "openai/whisper-small",
+                output_hidden_states=True,
+                max_source_positions=750,
+                ignore_mismatched_sizes=True
+            )
+        elif self.pretrain_model == "whisper_medium":
+            self.backbone_model = WhisperModel.from_pretrained(
+                "openai/whisper-medium",
+                output_hidden_states=True,
+                ignore_mismatched_sizes=True
+            )
+        elif self.pretrain_model == "whisper_large":
+            self.feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-large-v3", chunk_length=15)
+            self.backbone_model = WhisperModel.from_pretrained(
+                "openai/whisper-large-v3",
+                output_hidden_states=True,
+                ignore_mismatched_sizes=True,
+                max_source_positions=750,
+            )
+        self.embed_positions = copy.deepcopy(self.backbone_model.encoder.embed_positions.weight)
+        self.embed_positions.requires_grad = False
+
+        state_dict = self.backbone_model.state_dict()
+        # 2. Read the model config
+        self.model_config = self.backbone_model.config
+        self.model_config.finetune_method       = finetune_method
+        self.model_config.lora_rank             = lora_rank
+        self.finetune_method                    = finetune_method
+        self.apply_gradient_reversal            = apply_gradient_reversal
+        self.use_conv_output                    = use_conv_output
+        
+        if self.finetune_method == "lora":
+            # 3. Config encoder layers with adapter or embedding prompt
+            self.backbone_model.encoder.layers = nn.ModuleList(
+                [WhisperEncoderLayer(self.model_config, layer_idx) for layer_idx in range(self.model_config.encoder_layers)]
+            )
+            # 4. Load the weights back
+            msg = self.backbone_model.load_state_dict(state_dict, strict=False)
+        
+        # 2. Freeze the weights
+        self.freeze_params = freeze_params
+        if self.freeze_params and self.finetune_method != "lora":
+            for _, p in self.backbone_model.named_parameters(): p.requires_grad = False
+        elif self.freeze_params and self.finetune_method == "lora":
+            for name, p in self.backbone_model.named_parameters():
+                if name in msg.missing_keys: p.requires_grad = True
+                else: p.requires_grad = False
+        else:
+            for name, p in self.backbone_model.named_parameters(): 
+                if "decoder" not in name and "conv1" not in name and "conv2" not in name and "embed_positions" not in name: p.requires_grad = True
+                else: p.requires_grad = False
+        
+        # 6. Downstream models
+        self.model_seq = nn.Sequential(
+            nn.Conv1d(self.model_config.hidden_size, hidden_dim, 1, padding=0),
+            nn.ReLU(),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(hidden_dim, hidden_dim, 1, padding=0),
+            nn.ReLU(),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(hidden_dim, hidden_dim, 1, padding=0)
+        )
+
+        if use_conv_output:
+            num_layers = self.model_config.num_hidden_layers + 1  # transformer layers + input embeddings
+            self.weights = nn.Parameter(torch.ones(num_layers)/num_layers)
+        else:
+            num_layers = self.model_config.num_hidden_layers
+            self.weights = nn.Parameter(torch.zeros(num_layers))
+        
+        if apply_gradient_reversal:
+            self.dataset_layer = nn.Sequential(
+                RevGrad(),
+                nn.Linear(hidden_dim, hidden_dim),
+                nn.ReLU(),
+                nn.Linear(hidden_dim, num_dataset),
+            )
+        self.out_layer = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, output_class_num),
+        )
+            
+        
+    def forward(self, x, length=None, return_feature=False):
+        # 1. feature extraction and projections
+        if length is not None:
+            max_audio_len = 15*16000
+            # Append to list for feature_extractor to work
+            new_x = list()
+            for idx in range(len(length)):
+                new_x.append(x[idx].detach().cpu().numpy())
+            
+            # Max length is max audio len in a batch
+            features = self.feature_extractor(
+                new_x,
+                return_tensors="pt", 
+                sampling_rate=16000,
+                max_length=max_audio_len
+            )
+            features = features.input_features.cuda()
+        else:
+            max_audio_len = 15*16000
+            features = self.feature_extractor(
+                x[0].detach().cpu(), 
+                return_tensors="pt", 
+                sampling_rate=16000,
+                max_length=max_audio_len
+            )
+            features = features.input_features.cuda()
+        
+        # 2. get length and mask
+        if length is not None:
+            length = self._get_feat_extract_output_lengths(length.detach().cpu())
+            # Replace positional embeddings
+            self.backbone_model.encoder.embed_positions = self.backbone_model.encoder.embed_positions.from_pretrained(self.embed_positions[:750])
+        else:
+            # Replace positional embeddings
+            length = torch.tensor([len(x[0])])
+            length = self._get_feat_extract_output_lengths(length)
+            self.backbone_model.encoder.embed_positions = self.backbone_model.encoder.embed_positions.from_pretrained(self.embed_positions[:750])
+            
+        # 3. transformer encoding features
+        # compute reduced attention_mask corresponding to feature vectors
+        features = self.backbone_model.encoder(
+            features, output_hidden_states=True
+        ).hidden_states
+
+        features = torch.stack(features, dim=0)[-1]
+
+        # 6. Pass the weighted average to point-wise 1D Conv
+        # B x T x D
+        features = features.transpose(1, 2)
+        features = self.model_seq(features)
+        features = features.transpose(1, 2)
+        
+        # 7. Pooling
+        if length is not None:
+            mean, std = list(), list()
+            for snt_id in range(features.shape[0]):
+                # Avoiding padded time steps
+                actual_size = length[snt_id]
+                mean.append(torch.mean(features[snt_id, 0:actual_size, ...], dim=0))
+            features = torch.stack(mean)
+        else:
+            features = torch.mean(features, dim=1)
+            
+        # 8. Output predictions
+        # B x D
+        predicted = self.out_layer(features)
+        if self.apply_gradient_reversal: 
+            dataset_predicted = self.dataset_layer(features)
+            if return_feature: return predicted, dataset_predicted, features
+            return predicted, dataset_predicted
+        if return_feature: return predicted, features
+        return predicted
+        
+    # From huggingface
+    def _get_feat_extract_output_lengths(self, input_lengths):
+        """
+        Computes the output length of the convolutional layers
+        """
+        input_lengths = input_lengths // 160
+        input_lengths = (input_lengths - 1) // 2 + 1
+        return input_lengths
+
+def prepare_mask(length, shape, dtype):
+    # Modified from huggingface
+    mask = torch.zeros(
+        shape, dtype=dtype
+    )
+    # these two operations makes sure that all values
+    # before the output lengths indices are attended to
+    mask[(torch.arange(mask.shape[0]), length.cpu() - 1)] = 1
+    mask = mask.flip([-1]).cumsum(-1).flip([-1]).bool()
+    return mask
+    

capspeech/eval/src/model/emotion/wavlm_emotion.py

@@ -0,0 +1,315 @@
+import os
+import torch
+import loralib as lora
+import transformers.models.wavlm.modeling_wavlm as wavlm
+from speechbrain.lobes.models.huggingface_transformers.huggingface import make_padding_masks
+
+from torch import nn
+from torch.nn import functional as F
+from transformers import Wav2Vec2FeatureExtractor
+from transformers import WavLMModel
+
+import sys
+from pathlib import Path
+sys.path.append(os.path.join(str(Path(os.path.realpath(__file__)).parents[1])))
+
+class WavLMEncoderLayer(nn.Module):
+    def __init__(self, layer_idx, config, has_relative_position_bias: bool = True):
+        super().__init__()
+        self.attention = wavlm.WavLMAttention(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            dropout=config.attention_dropout,
+            num_buckets=config.num_buckets,
+            max_distance=config.max_bucket_distance,
+            has_relative_position_bias=has_relative_position_bias,
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.feed_forward = wavlm.WavLMFeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.config = config
+        
+        if layer_idx > config.num_hidden_layers // 2:
+            if self.config.finetune_method == "lora" or self.config.finetune_method == "combined":
+                self.feed_forward.intermediate_dense    = lora.Linear(config.hidden_size, config.intermediate_size, r=config.lora_rank)
+                self.feed_forward.output_dense          = lora.Linear(config.intermediate_size, config.hidden_size, r=config.lora_rank)
+        
+    def forward(self, hidden_states, attention_mask=None, position_bias=None, output_attentions=False, index=0):
+        attn_residual = hidden_states
+        hidden_states, attn_weights, position_bias = self.attention(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_bias=position_bias,
+            output_attentions=output_attentions,
+            index=index,
+        )
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = attn_residual + hidden_states
+        
+        # Adapter
+        if self.config.finetune_method == "adapter":
+            adapt_h = self.adapter(hidden_states)
+
+        hidden_states = self.layer_norm(hidden_states)
+        hidden_states = hidden_states + self.feed_forward(hidden_states)
+        hidden_states = self.final_layer_norm(hidden_states)
+        outputs = (hidden_states, position_bias)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+class WavLMEncoderLayerStableLayerNorm(nn.Module):
+    def __init__(self, layer_idx, config, has_relative_position_bias: bool = True):
+        super().__init__()
+        self.attention = wavlm.WavLMAttention(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            dropout=config.attention_dropout,
+            num_buckets=config.num_buckets,
+            max_distance=config.max_bucket_distance,
+            has_relative_position_bias=has_relative_position_bias,
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.feed_forward = wavlm.WavLMFeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.config = config
+
+        if layer_idx > config.num_hidden_layers // 2:
+            if self.config.finetune_method == "lora" or self.config.finetune_method == "combined":
+                self.feed_forward.intermediate_dense    = lora.Linear(config.hidden_size, config.intermediate_size, r=config.lora_rank)
+                self.feed_forward.output_dense          = lora.Linear(config.intermediate_size, config.hidden_size, r=config.lora_rank)
+            
+
+    def forward(self, hidden_states, attention_mask=None, position_bias=None, output_attentions=False):
+        attn_residual = hidden_states
+        hidden_states = self.layer_norm(hidden_states)
+        hidden_states, attn_weights, position_bias = self.attention(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_bias=position_bias,
+            output_attentions=output_attentions,
+        )
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = attn_residual + hidden_states
+        hidden_states = hidden_states + self.feed_forward(self.final_layer_norm(hidden_states))
+
+        outputs = (hidden_states, position_bias)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+   
+class WavLMWrapper(nn.Module):
+    def __init__(
+        self, 
+        pretrain_model="wavlm_large", 
+        hidden_dim=256,
+        finetune_method="lora",
+        lora_rank=16,
+        freeze_params=True,
+        output_class_num=4,
+        use_conv_output=True,
+        detailed_class_num=17
+    ):
+        super(WavLMWrapper, self).__init__()
+        # 1. We Load the model first with weights
+        self.pretrain_model     = pretrain_model
+        self.finetune_method    = finetune_method
+        self.freeze_params      = freeze_params
+        self.use_conv_output    = use_conv_output
+        self.lora_rank          = lora_rank
+        if self.pretrain_model == "wavlm":
+            self.backbone_model = WavLMModel.from_pretrained(
+                "microsoft/wavlm-base-plus",
+                output_hidden_states=True,
+            )
+        elif self.pretrain_model == "wavlm_large":
+            self.processor = Wav2Vec2FeatureExtractor.from_pretrained('microsoft/wavlm-large')
+            self.backbone_model = WavLMModel.from_pretrained(
+                "microsoft/wavlm-large",
+                output_hidden_states=True,
+            )
+        state_dict = self.backbone_model.state_dict()
+        # 2. Read the model config
+        self.model_config = self.backbone_model.config
+        self.model_config.finetune_method        = self.finetune_method
+        self.model_config.lora_rank              = self.lora_rank
+        
+        # 3. Config encoder layers with adapter or embedding prompt
+        if self.pretrain_model == "wavlm":
+            self.backbone_model.encoder.layers = nn.ModuleList(
+                [WavLMEncoderLayer(i, self.model_config, has_relative_position_bias=(i == 0)) for i in range(self.model_config.num_hidden_layers)]
+            )
+        elif self.pretrain_model == "wavlm_large":
+            self.backbone_model.encoder.layers = nn.ModuleList(
+                [WavLMEncoderLayerStableLayerNorm(i, self.model_config, has_relative_position_bias=(i == 0)) for i in range(self.model_config.num_hidden_layers)]
+            )
+        # 4. Load the weights back
+        msg = self.backbone_model.load_state_dict(state_dict, strict=False)
+
+        # 5. Freeze the weights
+        self.freeze_params = freeze_params
+        if self.freeze_params and self.finetune_method != "lora":
+            for _, p in self.backbone_model.named_parameters(): p.requires_grad = False
+        elif self.freeze_params and self.finetune_method == "lora":
+            for name, p in self.backbone_model.named_parameters():
+                if name in msg.missing_keys: p.requires_grad = True
+                else: p.requires_grad = False
+        else:
+            for _, p in self.backbone_model.named_parameters(): p.requires_grad = True
+
+        # 6. Downstream models
+        self.model_seq = nn.Sequential(
+            nn.Conv1d(self.model_config.hidden_size, hidden_dim, 1, padding=0),
+            nn.ReLU(),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(hidden_dim, hidden_dim, 1, padding=0),
+            nn.ReLU(),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(hidden_dim, hidden_dim, 1, padding=0)
+        )
+
+        if self.use_conv_output:
+            num_layers = self.model_config.num_hidden_layers + 1  # transformer layers + input embeddings
+            self.weights = nn.Parameter(torch.ones(num_layers)/num_layers)
+        else:
+            num_layers = self.model_config.num_hidden_layers
+            self.weights = nn.Parameter(torch.zeros(num_layers))
+        
+        self.emotion_layer = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, output_class_num),
+        )
+
+        self.detailed_out_layer = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, detailed_class_num),
+        )
+        
+        self.arousal_layer = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 1),
+            nn.Sigmoid()
+        )
+
+        self.valence_layer = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 1),
+            nn.Sigmoid()
+        )
+
+        self.dominance_layer = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 1),
+            nn.Sigmoid()
+        )
+        
+    def forward(self, x, length=None, return_feature=False):
+        # 1. feature extraction and projections
+        if self.pretrain_model == "wavlm_large":  
+            with torch.no_grad():
+                signal, attention_mask = list(), list()
+                if length is not None: attention_mask = make_padding_masks(x, wav_len=length/length.max()).to(x.device)
+                else: attention_mask = make_padding_masks(x, wav_len=torch.tensor([1]).to(x.device)).to(x.device)
+
+                for idx in range(len(x)):
+                    input = self.processor(x[idx], sampling_rate=16_000, return_tensors="pt", padding=True)
+                    signal.append(input["input_values"][0].to(x.device))
+                signal = torch.stack(signal)
+
+        # 2. get length and mask
+        if length is not None:
+            length = self.get_feat_extract_output_lengths(length.detach().cpu())
+            length = length.cuda()
+
+        if self.pretrain_model == "wavlm": 
+            x = self.backbone_model(
+                x, output_hidden_states=True
+            ).hidden_states
+        else: 
+            x = self.backbone_model(
+                signal, 
+                attention_mask=attention_mask, 
+                output_hidden_states=True
+            ).hidden_states
+        
+        # 4. stacked feature
+        if self.use_conv_output: stacked_feature = torch.stack(x, dim=0)
+        else: stacked_feature = torch.stack(x, dim=0)[1:]
+        
+        # 5. Weighted sum
+        _, *origin_shape = stacked_feature.shape
+        # Return transformer enc outputs [num_enc_layers, B, T, D]
+        if self.use_conv_output:
+            stacked_feature = stacked_feature.view(self.backbone_model.config.num_hidden_layers+1, -1)
+        else:
+            stacked_feature = stacked_feature.view(self.backbone_model.config.num_hidden_layers, -1)
+        norm_weights = F.softmax(self.weights, dim=-1)
+        
+        # Perform weighted average
+        weighted_feature = (norm_weights.unsqueeze(-1) * stacked_feature).sum(dim=0)
+        features = weighted_feature.view(*origin_shape)
+        
+        # 6. Pass the weighted average to point-wise 1D Conv
+        # B x T x D
+        features = features.transpose(1, 2)
+        features = self.model_seq(features)
+        features = features.transpose(1, 2)
+        
+        # 7. Pooling
+        if length is not None:
+            mean, std = list(), list()
+            for snt_id in range(features.shape[0]):
+                # Avoiding padded time steps
+                actual_size = length[snt_id]
+                mean.append(torch.mean(features[snt_id, 0:actual_size, ...], dim=0))
+            features = torch.stack(mean)
+        else:
+            features = torch.mean(features, dim=1)
+
+        # Output predictions
+        # B x D
+        predicted           = self.emotion_layer(features)
+        detailed_predicted  = self.detailed_out_layer(features)
+        arousal             = self.arousal_layer(features)
+        valence             = self.valence_layer(features)
+        dominance           = self.dominance_layer(features)
+        if return_feature: return predicted, features, detailed_predicted, arousal, valence, dominance
+        return predicted, detailed_predicted, arousal, valence, dominance
+    
+    # From huggingface
+    def get_feat_extract_output_lengths(self, input_length):
+        """
+        Computes the output length of the convolutional layers
+        """
+        def _conv_out_length(input_length, kernel_size, stride):
+            # 1D convolutional layer output length formula taken
+            # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html
+            return (input_length - kernel_size) // stride + 1
+        for kernel_size, stride in zip(self.backbone_model.config.conv_kernel, self.backbone_model.config.conv_stride):
+            input_length = _conv_out_length(input_length, kernel_size, stride)
+        return input_length
+
+def prepare_mask(length, shape, dtype):
+    # Modified from huggingface
+    mask = torch.zeros(
+        shape, dtype=dtype
+    )
+    # these two operations makes sure that all values
+    # before the output lengths indices are attended to
+    mask[(torch.arange(mask.shape[0]), length.cpu() - 1)] = 1
+    mask = mask.flip([-1]).cumsum(-1).flip([-1]).bool()
+    return mask
+    

capspeech/eval/src/model/emotion/wavlm_emotion_dim.py

@@ -0,0 +1,318 @@
+import os
+import pdb
+import torch
+import argparse
+import numpy as np
+import loralib as lora
+import transformers.models.wav2vec2.modeling_wav2vec2 as w2v2
+import transformers.models.wavlm.modeling_wavlm as wavlm
+from speechbrain.lobes.models.huggingface_transformers.huggingface import make_padding_masks
+
+from torch import nn
+from torch.nn import functional as F
+from transformers import Wav2Vec2FeatureExtractor
+from transformers import WavLMModel
+
+import sys
+from pathlib import Path
+sys.path.append(os.path.join(str(Path(os.path.realpath(__file__)).parents[1])))
+
+class WavLMEncoderLayer(nn.Module):
+    def __init__(self, layer_idx, config, has_relative_position_bias: bool = True):
+        super().__init__()
+        self.attention = wavlm.WavLMAttention(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            dropout=config.attention_dropout,
+            num_buckets=config.num_buckets,
+            max_distance=config.max_bucket_distance,
+            has_relative_position_bias=has_relative_position_bias,
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.feed_forward = wavlm.WavLMFeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.config = config
+        
+        if layer_idx > config.num_hidden_layers // 2:
+            if self.config.finetune_method == "lora" or self.config.finetune_method == "combined":
+                self.feed_forward.intermediate_dense    = lora.Linear(config.hidden_size, config.intermediate_size, r=config.lora_rank)
+                self.feed_forward.output_dense          = lora.Linear(config.intermediate_size, config.hidden_size, r=config.lora_rank)
+        
+    def forward(self, hidden_states, attention_mask=None, position_bias=None, output_attentions=False, index=0):
+        attn_residual = hidden_states
+        hidden_states, attn_weights, position_bias = self.attention(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_bias=position_bias,
+            output_attentions=output_attentions,
+            index=index,
+        )
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = attn_residual + hidden_states
+        
+        # Adapter
+        if self.config.finetune_method == "adapter":
+            adapt_h = self.adapter(hidden_states)
+
+        hidden_states = self.layer_norm(hidden_states)
+        hidden_states = hidden_states + self.feed_forward(hidden_states)
+        hidden_states = self.final_layer_norm(hidden_states)
+        outputs = (hidden_states, position_bias)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+class WavLMEncoderLayerStableLayerNorm(nn.Module):
+    def __init__(self, layer_idx, config, has_relative_position_bias: bool = True):
+        super().__init__()
+        self.attention = wavlm.WavLMAttention(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            dropout=config.attention_dropout,
+            num_buckets=config.num_buckets,
+            max_distance=config.max_bucket_distance,
+            has_relative_position_bias=has_relative_position_bias,
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.feed_forward = wavlm.WavLMFeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.config = config
+
+        if layer_idx > config.num_hidden_layers // 2:
+            if self.config.finetune_method == "lora" or self.config.finetune_method == "combined":
+                self.feed_forward.intermediate_dense    = lora.Linear(config.hidden_size, config.intermediate_size, r=config.lora_rank)
+                self.feed_forward.output_dense          = lora.Linear(config.intermediate_size, config.hidden_size, r=config.lora_rank)
+            
+
+    def forward(self, hidden_states, attention_mask=None, position_bias=None, output_attentions=False):
+        attn_residual = hidden_states
+        hidden_states = self.layer_norm(hidden_states)
+        hidden_states, attn_weights, position_bias = self.attention(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_bias=position_bias,
+            output_attentions=output_attentions,
+        )
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = attn_residual + hidden_states
+        hidden_states = hidden_states + self.feed_forward(self.final_layer_norm(hidden_states))
+
+        outputs = (hidden_states, position_bias)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+   
+class WavLMWrapper(nn.Module):
+    def __init__(
+        self, 
+        pretrain_model="wavlm_large", 
+        hidden_dim=256,
+        finetune_method="lora",
+        lora_rank=16,
+        freeze_params=True,
+        output_class_num=4,
+        use_conv_output=True,
+        detailed_class_num=17,
+        predict_gender=False
+    ):
+        super(WavLMWrapper, self).__init__()
+        # 1. We Load the model first with weights
+        self.pretrain_model     = pretrain_model
+        self.finetune_method    = finetune_method
+        self.freeze_params      = freeze_params
+        self.use_conv_output    = use_conv_output
+        self.lora_rank          = lora_rank
+        self.predict_gender     = predict_gender
+        if self.pretrain_model == "wavlm":
+            self.backbone_model = WavLMModel.from_pretrained(
+                "microsoft/wavlm-base-plus",
+                output_hidden_states=True,
+            )
+        elif self.pretrain_model == "wavlm_large":
+            self.processor = Wav2Vec2FeatureExtractor.from_pretrained('microsoft/wavlm-large')
+            self.backbone_model = WavLMModel.from_pretrained(
+                "microsoft/wavlm-large",
+                output_hidden_states=True,
+            )
+        state_dict = self.backbone_model.state_dict()
+        # 2. Read the model config
+        self.model_config = self.backbone_model.config
+        self.model_config.finetune_method        = self.finetune_method
+        self.model_config.lora_rank              = self.lora_rank
+        
+        # 3. Config encoder layers with adapter or embedding prompt
+        if self.pretrain_model == "wavlm":
+            self.backbone_model.encoder.layers = nn.ModuleList(
+                [WavLMEncoderLayer(i, self.model_config, has_relative_position_bias=(i == 0)) for i in range(self.model_config.num_hidden_layers)]
+            )
+        elif self.pretrain_model == "wavlm_large":
+            self.backbone_model.encoder.layers = nn.ModuleList(
+                [WavLMEncoderLayerStableLayerNorm(i, self.model_config, has_relative_position_bias=(i == 0)) for i in range(self.model_config.num_hidden_layers)]
+            )
+        # 4. Load the weights back
+        msg = self.backbone_model.load_state_dict(state_dict, strict=False)
+
+        # 5. Freeze the weights
+        self.freeze_params = freeze_params
+        if self.freeze_params and self.finetune_method != "lora":
+            for _, p in self.backbone_model.named_parameters(): p.requires_grad = False
+        elif self.freeze_params and self.finetune_method == "lora":
+            for name, p in self.backbone_model.named_parameters():
+                if name in msg.missing_keys: p.requires_grad = True
+                else: p.requires_grad = False
+        else:
+            for _, p in self.backbone_model.named_parameters(): p.requires_grad = True
+
+        # 6. Downstream models
+        self.model_seq = nn.Sequential(
+            nn.Conv1d(self.model_config.hidden_size, hidden_dim, 1, padding=0),
+            nn.ReLU(),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(hidden_dim, hidden_dim, 1, padding=0),
+            nn.ReLU(),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(hidden_dim, hidden_dim, 1, padding=0)
+        )
+
+        if self.use_conv_output:
+            num_layers = self.model_config.num_hidden_layers + 1  # transformer layers + input embeddings
+            self.weights = nn.Parameter(torch.ones(num_layers)/num_layers)
+        else:
+            num_layers = self.model_config.num_hidden_layers
+            self.weights = nn.Parameter(torch.zeros(num_layers))
+        
+        self.arousal_layer = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 1),
+            nn.Sigmoid()
+        )
+
+        self.valence_layer = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 1),
+            nn.Sigmoid()
+        )
+
+        self.dominance_layer = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 1),
+            nn.Sigmoid()
+        )
+        
+        if self.predict_gender:
+            self.gender_layer = nn.Sequential(
+                nn.Linear(hidden_dim, hidden_dim),
+                nn.ReLU(),
+                nn.Linear(hidden_dim, 2)
+            )
+        
+    def forward(self, x, length=None, return_feature=False):
+        # 1. feature extraction and projections
+        if self.pretrain_model == "wavlm_large":  
+            with torch.no_grad():
+                signal, attention_mask = list(), list()
+                if length is not None: attention_mask = make_padding_masks(x, wav_len=length/length.max()).to(x.device)
+                else: attention_mask = make_padding_masks(x, wav_len=torch.tensor([1]).to(x.device)).to(x.device)
+
+                for idx in range(len(x)):
+                    input = self.processor(x[idx], sampling_rate=16_000, return_tensors="pt", padding=True)
+                    signal.append(input["input_values"][0].to(x.device))
+                signal = torch.stack(signal)
+
+        # 2. get length and mask
+        if length is not None:
+            length = self.get_feat_extract_output_lengths(length.detach().cpu())
+            length = length.cuda()
+
+        if self.pretrain_model == "wavlm": 
+            x = self.backbone_model(
+                x, output_hidden_states=True
+            ).hidden_states
+        else: 
+            x = self.backbone_model(
+                signal, 
+                attention_mask=attention_mask, 
+                output_hidden_states=True
+            ).hidden_states
+        
+        # 4. stacked feature
+        if self.use_conv_output: stacked_feature = torch.stack(x, dim=0)
+        else: stacked_feature = torch.stack(x, dim=0)[1:]
+        
+        # 5. Weighted sum
+        _, *origin_shape = stacked_feature.shape
+        # Return transformer enc outputs [num_enc_layers, B, T, D]
+        if self.use_conv_output:
+            stacked_feature = stacked_feature.view(self.backbone_model.config.num_hidden_layers+1, -1)
+        else:
+            stacked_feature = stacked_feature.view(self.backbone_model.config.num_hidden_layers, -1)
+        norm_weights = F.softmax(self.weights, dim=-1)
+        
+        # Perform weighted average
+        weighted_feature = (norm_weights.unsqueeze(-1) * stacked_feature).sum(dim=0)
+        features = weighted_feature.view(*origin_shape)
+        
+        # 6. Pass the weighted average to point-wise 1D Conv
+        # B x T x D
+        features = features.transpose(1, 2)
+        features = self.model_seq(features)
+        features = features.transpose(1, 2)
+        
+        # 7. Pooling
+        if length is not None:
+            mean, std = list(), list()
+            for snt_id in range(features.shape[0]):
+                # Avoiding padded time steps
+                actual_size = length[snt_id]
+                mean.append(torch.mean(features[snt_id, 0:actual_size, ...], dim=0))
+            features = torch.stack(mean)
+        else:
+            features = torch.mean(features, dim=1)
+
+        # 8. Output predictions
+        # B x D
+        arousal             = self.arousal_layer(features)
+        valence             = self.valence_layer(features)
+        dominance           = self.dominance_layer(features)
+        
+        if(self.predict_gender):
+            gender_outputs = self.gender_layer(features)
+            return arousal, valence, dominance, gender_outputs
+        
+        return arousal, valence, dominance
+    
+    # From huggingface
+    def get_feat_extract_output_lengths(self, input_length):
+        """
+        Computes the output length of the convolutional layers
+        """
+        def _conv_out_length(input_length, kernel_size, stride):
+            # 1D convolutional layer output length formula taken
+            # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html
+            return (input_length - kernel_size) // stride + 1
+        for kernel_size, stride in zip(self.backbone_model.config.conv_kernel, self.backbone_model.config.conv_stride):
+            input_length = _conv_out_length(input_length, kernel_size, stride)
+        return input_length
+
+def prepare_mask(length, shape, dtype):
+    # Modified from huggingface
+    mask = torch.zeros(
+        shape, dtype=dtype
+    )
+    # these two operations makes sure that all values
+    # before the output lengths indices are attended to
+    mask[(torch.arange(mask.shape[0]), length.cpu() - 1)] = 1
+    mask = mask.flip([-1]).cumsum(-1).flip([-1]).bool()
+    return mask
+