add latest changes

example_configs/librispeech_tts_r_100.json

 {
-    "model_name_or_path": "/raid/yoach/tmp/small-stable-speech-untrained/",
-    "feature_extractor_name":"facebook/encodec_32khz",
-    "description_tokenizer_name":"t5-base",
-    "prompt_tokenizer_name":"t5-base",
+    "model_name_or_path": "/raid/yoach/tmp/artefacts/small-stable-speech-untrained/",
+    "feature_extractor_name":"facebook/encodec_24khz",
+    "description_tokenizer_name":"google-t5/t5-small",
+    "prompt_tokenizer_name":"google-t5/t5-small",
 
-    "push_to_hub": false,
-    "hub_model_id": "stable-speech-mini",
+    "push_to_hub": true,
+    "hub_model_id": "ylacombe/stable-speech-mini",
     "report_to": ["wandb"],
     "overwrite_output_dir": true,
-    "output_dir": "/raid/yoach/tmp/artefacts/training/",
+    "output_dir": "/raid/yoach/tmp/artefacts/training-mini/",
 
 
     "train_dataset_name": "blabble-io/libritts_r",
@@ -28,22 +28,24 @@
 
     "max_eval_samples": 24,
     
-    "max_duration_in_seconds": 32,
+    "max_duration_in_seconds": 29,
     "min_duration_in_seconds": 2.0,
 
     "add_audio_samples_to_wandb": true,
     "id_column_name": "id",
 
-    "preprocessing_num_workers": 8,
+    "preprocessing_num_workers": 1,
+
+
+    "pad_token_id": 1024,
+    "decoder_start_token_id": 1025,
 
-    "pad_token_id": 2050,
-    "decoder_start_token_id": 2048,
 
     "do_train": true,
-    "num_train_epochs": 200,
+    "num_train_epochs": 15,
     "gradient_accumulation_steps": 1,
-    "gradient_checkpointing": false,
-    "per_device_train_batch_size": 50,
+    "gradient_checkpointing": true,
+    "per_device_train_batch_size": 40,
     "learning_rate": 1e-4,
     "adam_beta1": 0.9,
     "adam_beta2": 0.999,
@@ -61,11 +63,13 @@
     "predict_with_generate": true,
     "include_inputs_for_metrics": true,
     "evaluation_strategy": "steps",
-    "eval_steps": 600,
+    "eval_steps": 3000,
+    "save_steps": 3000,
+
     "per_device_eval_batch_size": 8,
     "generation_max_length": 400,
 
-    "audio_encode_per_device_eval_batch_size":24,
+    "audio_encode_per_device_eval_batch_size":32,
     "dtype": "float16",
     "seed": 456,
 

example_configs/librispeech_tts_r_dummy.json

 {
-    "model_name_or_path": "/home/yoach/dataspeech/artefacts/tiny-model/",
-    "feature_extractor_name":"facebook/encodec_32khz",
-    "description_tokenizer_name":"t5-base",
-    "prompt_tokenizer_name":"t5-base",
+    "model_name_or_path": "/raid/yoach/tmp/artefacts/tiny-model/",
+    "feature_extractor_name":"facebook/encodec_24khz",
+    "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": "/home/yoach/dataspeech/artefacts/training/",
+    "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",
@@ -36,11 +36,11 @@
 
     "preprocessing_num_workers": 1,
 
-    "pad_token_id": 2048,
-    "decoder_start_token_id": 2049,
+    "pad_token_id": 1024,
+    "decoder_start_token_id": 1025,
 
     "do_train": true,
-    "num_train_epochs": 120,
+    "num_train_epochs": 180,
     "gradient_accumulation_steps": 1,
     "gradient_checkpointing": false,
     "per_device_train_batch_size": 2,
@@ -62,7 +62,7 @@
     "evaluation_strategy": "steps",
     "eval_steps": 30,
     "per_device_eval_batch_size": 2,
-    "generation_max_length": 400,
+    "generation_max_length": 800,
     "do_sample": false,
 
     "logging_steps": 15,

init_dummy_model.py

@@ -2,7 +2,18 @@ from stable_speech import StableSpeechConfig, StableSpeechForCausalLM, StableSpe
 from transformers import T5Config, EncodecConfig
 from transformers import AutoConfig
 
+text_model = "google-t5/t5-small"
+encodec_version = "facebook/encodec_24khz"
+num_codebooks = 8
+
+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,
@@ -14,6 +25,10 @@ decoder_config = StableSpeechDecoderConfig(
     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 ?
         
@@ -21,26 +36,25 @@ decoder_config = StableSpeechDecoderConfig(
         
 decoder = StableSpeechForCausalLM(decoder_config)
 
-decoder.save_pretrained("/home/yoach/dataspeech/artefacts/decoder/")
+decoder.save_pretrained("/raid/yoach/tmp/artefacts/decoder/")
 
 
-t5 = AutoConfig.from_pretrained("t5-base")
 
 model = StableSpeechForConditionalGeneration.from_sub_models_pretrained(
-    text_encoder_pretrained_model_name_or_path="t5-base",
-    audio_encoder_pretrained_model_name_or_path="facebook/encodec_32khz",
-    decoder_pretrained_model_name_or_path="/home/yoach/dataspeech/artefacts/decoder/",
+    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 = 2049
-model.generation_config.pad_token_id = 2048
-model.generation_config.eos_token_id = 2048
+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("/home/yoach/dataspeech/artefacts/tiny-model/")
\ No newline at end of file
+model.save_pretrained("/raid/yoach/tmp/artefacts/tiny-model/")
\ No newline at end of file

init_model.py

@@ -2,9 +2,22 @@ from stable_speech import StableSpeechConfig, StableSpeechForCausalLM, StableSpe
 from transformers import T5Config, EncodecConfig
 from transformers import AutoConfig
 
+
+text_model = "google-t5/t5-small"
+encodec_version = "facebook/encodec_24khz"
+num_codebooks = 8
+
+
+t5 = AutoConfig.from_pretrained(text_model)
+encodec = AutoConfig.from_pretrained(encodec_version)
+
+encodec_vocab_size = encodec.codebook_size
+
+
 decoder_config = StableSpeechDecoderConfig(
-    max_position_embeddings=2048,
-    num_hidden_layers=24,
+    vocab_size=encodec_vocab_size+1,
+    max_position_embeddings=2250, # 30 s
+    num_hidden_layers=12,
     ffn_dim=4096,
     num_attention_heads=16,
     layerdrop=0.0,
@@ -14,30 +27,35 @@ decoder_config = StableSpeechDecoderConfig(
     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("/home/yoach/dataspeech/artefacts/decoder/")
+decoder.save_pretrained("/raid/yoach/tmp/artefacts/decoder/")
+
 
 
-t5 = AutoConfig.from_pretrained("t5-base")
 
 model = StableSpeechForConditionalGeneration.from_sub_models_pretrained(
-    text_encoder_pretrained_model_name_or_path="t5-base",
-    audio_encoder_pretrained_model_name_or_path="facebook/encodec_32khz",
-    decoder_pretrained_model_name_or_path="/home/yoach/dataspeech/artefacts/decoder/",
+    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 = 2049
-model.generation_config.pad_token_id = 2048
-model.generation_config.eos_token_id = 2048
+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("/home/yoach/dataspeech/artefacts/small-stable-speech-untrained/")
\ No newline at end of file
+
+model.save_pretrained("/raid/yoach/tmp/artefacts/small-stable-speech-untrained/")
\ No newline at end of file

run_stable_speech_training.py

@@ -257,7 +257,7 @@ class ModelArguments:
         metadata={"help": "Whether to do sampling or greedy decoding."},
     )
     bandwidth: float = field(
-        default=3, # TODO
+        default=6, # TODO
         metadata={"help": "Audio encoder bandwidth."},
     )