add dac config, init, and temporary datasets saving

example_configs/librispeech_tts_r.json

@@ -63,7 +63,7 @@
     "evaluation_strategy": "steps",
     "eval_steps": 600,
     "per_device_eval_batch_size": 8,
-    "generation_max_length": 400,
+    "generation_max_length": 2250,
 
     "fp16": false,
     "seed": 456,

example_configs/librispeech_tts_r_100.json

 {
     "model_name_or_path": "/raid/yoach/tmp/artefacts/small-stable-speech-untrained/",
-    "feature_extractor_name":"facebook/encodec_24khz",
+    "feature_extractor_name":"ylacombe/dac_44khZ_8kbps",
     "description_tokenizer_name":"google-t5/t5-small",
     "prompt_tokenizer_name":"google-t5/t5-small",
 
-    "push_to_hub": true,
+    "push_to_hub": false,
     "hub_model_id": "ylacombe/stable-speech-mini",
     "report_to": ["wandb"],
-    "overwrite_output_dir": true,
+    "overwrite_output_dir": false,
     "output_dir": "/raid/yoach/tmp/artefacts/training-mini/",
 
 
@@ -34,7 +34,7 @@
     "add_audio_samples_to_wandb": true,
     "id_column_name": "id",
 
-    "preprocessing_num_workers": 1,
+    "preprocessing_num_workers": 8,
 
 
     "pad_token_id": 1024,
@@ -45,7 +45,7 @@
     "num_train_epochs": 15,
     "gradient_accumulation_steps": 1,
     "gradient_checkpointing": true,
-    "per_device_train_batch_size": 40,
+    "per_device_train_batch_size": 28,
     "learning_rate": 1e-4,
     "adam_beta1": 0.9,
     "adam_beta2": 0.999,
@@ -63,11 +63,10 @@
     "predict_with_generate": true,
     "include_inputs_for_metrics": true,
     "evaluation_strategy": "steps",
-    "eval_steps": 3000,
-    "save_steps": 3000,
+    "eval_steps": 2500,
+    "save_steps": 2499,
 
     "per_device_eval_batch_size": 8,
-    "generation_max_length": 400,
 
     "audio_encode_per_device_eval_batch_size":32,
     "dtype": "float16",

example_configs/librispeech_tts_r_75M.json

+{
+    "model_name_or_path": "/raid/yoach/tmp/artefacts/stable-speech-untrained-75M/",
+    "save_to_disk":  "/raid/yoach/tmp/artefacts/libritts_r_1k_hours_processed/",
+    "preprocessing_only": false,
+
+
+    "feature_extractor_name":"ylacombe/dac_44khZ_8kbps",
+    "description_tokenizer_name":"google/t5-v1_1-small",
+    "prompt_tokenizer_name":"google/t5-v1_1-small",
+
+    "push_to_hub": false,
+    "hub_model_id": "ylacombe/stable-speech-75M",
+    "report_to": ["wandb"],
+    "overwrite_output_dir": false,
+    "output_dir": "/raid/yoach/tmp/artefacts/training-75M-0.1/",
+
+    "train_dataset_name": "blabble-io/libritts_r+blabble-io/libritts_r+blabble-io/libritts_r",
+    "train_metadata_dataset_name": "stable-speech/libritts-r-tags-and-text-generated+stable-speech/libritts-r-tags-and-text-generated+stable-speech/libritts-r-tags-and-text-generated",
+    "train_dataset_config_name": "clean+clean+other",
+    "train_split_name": "train.clean.360+train.clean.100+train.other.500",
+
+    "eval_dataset_name": "blabble-io/libritts_r+blabble-io/libritts_r",
+    "eval_metadata_dataset_name": "stable-speech/libritts-r-tags-and-text-generated+stable-speech/libritts-r-tags-and-text-generated",
+    "eval_dataset_config_name": "clean+other",
+    "eval_split_name": "test.clean+test.other",
+
+
+    "target_audio_column_name": "audio", 
+    "description_column_name": "text_description",
+    "prompt_column_name": "text",
+
+    "max_eval_samples": 24,
+    
+    "max_duration_in_seconds": 35,
+    "min_duration_in_seconds": 2.0,
+
+    "add_audio_samples_to_wandb": true,
+    "id_column_name": "id",
+
+    "preprocessing_num_workers": 16,
+
+    "pad_token_id": 1024,
+    "decoder_start_token_id": 1025,
+
+    "do_train": true,
+    "num_train_epochs": 1,
+    "gradient_accumulation_steps": 1,
+    "gradient_checkpointing": true,
+    "per_device_train_batch_size": 28,
+    "learning_rate": 1e-4,
+    "adam_beta1": 0.9,
+    "adam_beta2": 0.999,
+    "weight_decay": 0.03,
+
+    "lr_scheduler_type": "constant_with_warmup",
+    "warmup_steps":  5000,
+
+
+    "logging_steps": 102,
+    "freeze_text_encoder": true,
+
+
+    "do_eval": true, 
+    "predict_with_generate": true,
+    "include_inputs_for_metrics": true,
+    "evaluation_strategy": "steps",
+    "eval_steps": 2500,
+    "save_steps": 2499,
+
+    "per_device_eval_batch_size": 1,
+
+    "audio_encode_per_device_eval_batch_size":24,
+    "dtype": "bfloat16",
+    "seed": 456,
+
+    "dataloader_num_workers":16
+}

example_configs/librispeech_tts_r_dummy_dac.json

+{
+    "model_name_or_path": "/raid/yoach/tmp/artefacts/tiny-dac-model/",
+    "save_to_disk":  "/raid/yoach/tmp/artefacts/small_experiment_dataset/",
+
+    "feature_extractor_name":"ylacombe/dac_44khZ_8kbps",
+    "description_tokenizer_name":"google-t5/t5-small",
+    "prompt_tokenizer_name":"google-t5/t5-small",
+
+    "push_to_hub": false,
+    "hub_model_id": "stable-speech-mini",
+    "report_to": ["wandb"],
+    "overwrite_output_dir": true,
+    "output_dir": "/raid/yoach/tmp/artefacts/training/",
+
+    "train_dataset_name": "blabble-io/libritts_r",
+    "train_metadata_dataset_name": "stable-speech/libritts-r-tags-and-text-generated",
+    "train_dataset_config_name": "clean",
+    "train_split_name": "train.clean.360",
+
+    "eval_dataset_name": "blabble-io/libritts_r",
+    "eval_metadata_dataset_name": "stable-speech/libritts-r-tags-and-text-generated",
+    "eval_dataset_config_name": "clean",
+    "eval_split_name": "train.clean.360",
+
+    "target_audio_column_name": "audio", 
+    "description_column_name": "text_description",
+    "prompt_column_name": "text",
+
+    "max_train_samples": 4,
+    "max_eval_samples": 4,
+
+    
+    "max_duration_in_seconds": 30,
+    "min_duration_in_seconds": 1.0,
+
+    "add_audio_samples_to_wandb": true,
+    "id_column_name": "id",
+
+    "preprocessing_num_workers": 1,
+
+    "pad_token_id": 1024,
+    "decoder_start_token_id": 1025,
+
+    "do_train": true,
+    "num_train_epochs": 180,
+    "gradient_accumulation_steps": 1,
+    "gradient_checkpointing": false,
+    "per_device_train_batch_size": 2,
+    "learning_rate": 1e-3,
+    "adam_beta1": 0.9,
+    "adam_beta2": 0.999,
+    "weight_decay": 0.1,
+
+    "lr_scheduler_type": "cosine",
+    "warmup_ratio":  0.1,
+
+
+    "freeze_text_encoder": true,
+
+
+    "do_eval": true, 
+    "predict_with_generate": true,
+    "include_inputs_for_metrics": true,
+    "evaluation_strategy": "steps",
+    "eval_steps": 30,
+    "per_device_eval_batch_size": 2,
+    "generation_max_length": 800,
+    "do_sample": false,
+
+    "logging_steps": 15,
+
+    "dtype": "float32",
+    "seed": 456,
+
+    "dataloader_num_workers":8
+}

init_dummy_model_dac.py

+from stable_speech import StableSpeechConfig, StableSpeechForCausalLM, StableSpeechForConditionalGeneration, StableSpeechDecoderConfig
+from transformers import AutoConfig
+from transformers import AutoModel
+
+
+from transformers import AutoConfig, AutoModel
+from stable_speech import DACConfig, DACModel
+
+AutoConfig.register("dac", DACConfig)
+AutoModel.register(DACConfig, DACModel)
+
+text_model = "google-t5/t5-small"
+encodec_version = "ylacombe/dac_44khZ_8kbps"
+num_codebooks = 9
+
+t5 = AutoConfig.from_pretrained(text_model)
+encodec = AutoConfig.from_pretrained(encodec_version)
+
+encodec_vocab_size = encodec.codebook_size
+
+
+decoder_config = StableSpeechDecoderConfig(
+    vocab_size=encodec_vocab_size+1,
+    max_position_embeddings=2048,
+    num_hidden_layers=4,
+    ffn_dim=512,
+    num_attention_heads=8,
+    layerdrop=0.0,
+    use_cache=True,
+    activation_function="gelu",
+    hidden_size=512,
+    dropout=0.0,
+    attention_dropout=0.0,
+    activation_dropout=0.0,
+    pad_token_id=encodec_vocab_size,
+    eos_token_id=encodec_vocab_size,
+    bos_token_id=encodec_vocab_size+1,
+    num_codebooks=num_codebooks,
+)
+# TODO: ?? how to make it stop ?
+        
+
+        
+decoder = StableSpeechForCausalLM(decoder_config)
+
+decoder.save_pretrained("/raid/yoach/tmp/artefacts/decoder/")
+
+
+
+model = StableSpeechForConditionalGeneration.from_sub_models_pretrained(
+    text_encoder_pretrained_model_name_or_path=text_model,
+    audio_encoder_pretrained_model_name_or_path=encodec_version,
+    decoder_pretrained_model_name_or_path="/raid/yoach/tmp/artefacts/decoder/",
+    vocab_size = t5.vocab_size
+)
+
+# set the appropriate bos/pad token ids
+model.generation_config.decoder_start_token_id = encodec_vocab_size+1
+model.generation_config.pad_token_id = encodec_vocab_size
+model.generation_config.eos_token_id = encodec_vocab_size
+
+# set other default generation config params
+model.generation_config.max_length = int(30 * model.audio_encoder.config.frame_rate)
+model.generation_config.do_sample = False # True
+model.generation_config.guidance_scale = 1 # 3.0
+
+model.save_pretrained("/raid/yoach/tmp/artefacts/tiny-dac-model/")
\ No newline at end of file

init_model.py

@@ -3,9 +3,15 @@ from transformers import T5Config, EncodecConfig
 from transformers import AutoConfig
 
 
+from transformers import AutoConfig, AutoModel
+from stable_speech import DACConfig, DACModel
+
+AutoConfig.register("dac", DACConfig)
+AutoModel.register(DACConfig, DACModel)
+
 text_model = "google-t5/t5-small"
-encodec_version = "facebook/encodec_24khz"
-num_codebooks = 8
+encodec_version = "ylacombe/dac_44khZ_8kbps"
+num_codebooks = 9
 
 
 t5 = AutoConfig.from_pretrained(text_model)
@@ -16,7 +22,7 @@ encodec_vocab_size = encodec.codebook_size
 
 decoder_config = StableSpeechDecoderConfig(
     vocab_size=encodec_vocab_size+1,
-    max_position_embeddings=2250, # 30 s
+    max_position_embeddings=3000, # 30 s = 2580
     num_hidden_layers=12,
     ffn_dim=4096,
     num_attention_heads=16,

init_model_75M.py

+from stable_speech import StableSpeechConfig, StableSpeechForCausalLM, StableSpeechForConditionalGeneration, StableSpeechDecoderConfig
+from transformers import T5Config, EncodecConfig
+from transformers import AutoConfig
+
+
+from transformers import AutoConfig, AutoModel
+from stable_speech import DACConfig, DACModel
+
+AutoConfig.register("dac", DACConfig)
+AutoModel.register(DACConfig, DACModel)
+
+text_model = "google/t5-v1_1-small"
+encodec_version = "ylacombe/dac_44khZ_8kbps"
+num_codebooks = 9
+
+
+t5 = AutoConfig.from_pretrained(text_model)
+encodec = AutoConfig.from_pretrained(encodec_version)
+
+encodec_vocab_size = encodec.codebook_size
+
+
+decoder_config = StableSpeechDecoderConfig(
+    vocab_size=encodec_vocab_size+1,
+    max_position_embeddings=4096, # 30 s = 2580
+    num_hidden_layers=8,
+    ffn_dim=3072,
+    num_attention_heads=12,
+    layerdrop=0.0,
+    use_cache=True,
+    activation_function="gelu",
+    hidden_size=768,
+    dropout=0.0,
+    attention_dropout=0.0,
+    activation_dropout=0.0,
+    pad_token_id=encodec_vocab_size,
+    eos_token_id=encodec_vocab_size,
+    bos_token_id=encodec_vocab_size+1,
+    num_codebooks=num_codebooks,
+)
+
+        
+decoder = StableSpeechForCausalLM(decoder_config)
+decoder.save_pretrained("/raid/yoach/tmp/artefacts/decoder_small/")
+
+
+
+
+model = StableSpeechForConditionalGeneration.from_sub_models_pretrained(
+    text_encoder_pretrained_model_name_or_path=text_model,
+    audio_encoder_pretrained_model_name_or_path=encodec_version,
+    decoder_pretrained_model_name_or_path="/raid/yoach/tmp/artefacts/decoder_small/",
+    vocab_size = t5.vocab_size
+)
+
+# set the appropriate bos/pad token ids
+model.generation_config.decoder_start_token_id = encodec_vocab_size+1
+model.generation_config.pad_token_id = encodec_vocab_size
+model.generation_config.eos_token_id = encodec_vocab_size
+
+# set other default generation config params
+model.generation_config.max_length = int(30 * model.audio_encoder.config.frame_rate)
+model.generation_config.do_sample = False # True
+model.generation_config.guidance_scale = 1 # 3.0
+
+
+model.save_pretrained("/raid/yoach/tmp/artefacts/stable-speech-untrained-75M/")
\ No newline at end of file

run_stable_speech_training.py

@@ -397,7 +397,8 @@ class DataTrainingArguments:
                 "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"
+                " 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. "
             )
         },
     )
@@ -442,6 +443,12 @@ class DataTrainingArguments:
         default="stable-speech",
         metadata={"help": "The name of the wandb project."},
     )
+    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."
+        }
+    )
     
 @dataclass
 class StableSpeechTrainingArguments(Seq2SeqTrainingArguments):
@@ -781,10 +788,19 @@ def main():
 
     # Set seed before initializing model.
     set_seed(training_args.seed)
+    num_workers = data_args.preprocessing_num_workers
 
     # 1. First, 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:
+        vectorized_datasets = datasets.load_from_disk(data_args.save_to_disk)
+    else:    
         raw_datasets = DatasetDict()
-    num_workers = data_args.preprocessing_num_workers
 
         columns_to_keep = {
             "target_audio_column_name": data_args.target_audio_column_name,
@@ -921,6 +937,7 @@ def main():
     num_codebooks = model.decoder.config.num_codebooks
     bandwidth = model_args.bandwidth
     
+    if not dataset_was_precomputed:
         # resample target audio
         raw_datasets = raw_datasets.cast_column(
             target_audio_column_name, datasets.features.Audio(sampling_rate=sampling_rate)
@@ -966,7 +983,6 @@ def main():
                 input_columns=["target_length"],
             )
 
-
         # 5. Now we encode the audio labels with encodec.
         # We use Accelerate to perform distributed inference
 
@@ -1054,7 +1070,7 @@ def main():
             with accelerator.main_process_first():
                 vectorized_datasets[split] = vectorized_datasets[split].map(
                     postprocess_dataset,
-                num_proc=num_workers,
+                    num_proc=1, # this one is resource consuming if many processor.
                     input_columns=["input_ids", "prompt_input_ids"],
                     desc="Postprocessing labeling",
                     with_indices=True,
@@ -1066,15 +1082,19 @@ def main():
         del generate_labels
 
     
+    if data_args.save_to_disk is not None and not dataset_was_precomputed:
+        vectorized_datasets.save_to_disk(data_args.save_to_disk)
+        logger.info(f"Dataset saved at {data_args.save_to_disk}")
 
     # 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:
-        # TODO: save to disk in this step instead of something else ??
-        logger.info(f"Data preprocessing finished. Files cached at {vectorized_datasets.cache_files}")
+    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