# Hyperparameters toggles
prefix: ""

## SSL features Selection
pretrained_models_path: pretrained_models/
# pretrained_models:
# {
#     "wav2vec2_base": "facebook/wav2vec2-base", # 768
#     "hubert_base": "facebook/hubert-base-ls960", # 768
#     "wavlm_base": "microsoft/wavlm-base", # 768
#     "wavlm_base_plus": "microsoft/wavlm-base-plus", # 768
#     "hubert_multilingual": "utter-project/mHuBERT-147", # 768
#     "clap" : "laion/clap-htsat-fused", # 768
#     "data2vec_base": "facebook/data2vec-audio-base", # 768
    
#     "wav2vec2_large": "facebook/wav2vec2-large", # 1024
#     "hubert_large": "facebook/hubert-large-ls960", # 1024
#     "wavlm_large": "microsoft/wavlm-large-plus", # 1024
#     "data2vec_large": "facebook/data2vec-audio-large", #1024
#     "whisper_medium": "openai/whisper-medium", # 1024
    
#     "whisper_large_v3_turbo": "openai/whisper-large-v3-turbo", # 1280
# }


# select pretrained SSL models
perceived_ssl_model: "wavlm_large" # in pretrained_models
canonical_ssl_model: Null

# # models hidden size, varies by model
ENCODER_DIM: 1024

# # How to fuse the features
feature_fusion: "mono"        # Options: "mono" for single ssl, "dual_ssl_enc" for dual ssl encoder, "dual_loss" for single SSL dual ssl loss
blend_alpha: 0.5              # If using "blend" fusion

# Input files
# Data files
data_folder_save: "./data"
train_annotation: !ref <data_folder_save>/train-train.json
valid_annotation: !ref <data_folder_save>/train-dev.json
test_annotation: !ref <data_folder_save>/test.json
# Extra data
train_annotation_extra: !ref <data_folder_save>/train-train_with_extra.json
use_extra_train_data: False

evaluate_key: "PER" # use "mpd_f1_seq" for Transformer decoder path best mpd f1
                            # "PER_seq" for Transformer decoder's best error rate
                            # "PER" for ctc path best error rate
                            # "mpd_f1" for ctc path best mpd f1
max_save_models: 3 # Maximum number of saved models for each metrics 
# generate training id for output folder
# generate_training_id: !apply:trainer.generate_training_id.generate_training_id [!ref <perceived_ssl_model_id>, !ref <canonical_ssl_model_id>, !ref <feature_fusion>, !ref <prefix>]

# output files
output_folder: !ref exp_l2arctic/<perceived_ssl_model>_<canonical_ssl_model>_<feature_fusion>_<prefix>
per_file: !ref <output_folder>/per.txt
mpd_file: !ref <output_folder>/mpd.txt
save_folder: !ref <output_folder>/save
train_log: !ref <output_folder>/train_log.txt

on_training_test_wer_folder: !ref <output_folder>/on_training_test_wer
on_training_test_mpd_folder: !ref <output_folder>/on_training_test_mpd

# Training Target 
training_target: "target" # "target": deduplicated canonical phoneme sequence; "target_with_repeats": with repeats
                          # "canonical"
                          # "perceived": deduplicated perceived phoneme sequence

perceived_ssl: !new:speechbrain.lobes.models.huggingface_transformers.wavlm.WavLM
    source: "microsoft/wavlm-large"
    freeze: !ref <freeze_perceived_ssl>
    freeze_feature_extractor: !ref <freeze_perceived_feature_extractor>
    save_path: !ref <pretrained_models_path>
    output_all_hiddens: False
preceived_ssl_emb_layer: -1

enc: !new:torch.nn.Sequential
  - !new:speechbrain.lobes.models.VanillaNN.VanillaNN
     input_shape: [null, null, !ref <ENCODER_DIM>]
     activation: !ref <activation>
     dnn_blocks: !ref <dnn_layers>
     dnn_neurons: !ref <dnn_neurons>
  - !new:torch.nn.LayerNorm
     normalized_shape: !ref <dnn_neurons>

ctc_lin:  !new:speechbrain.nnet.linear.Linear
    input_size: !ref <dnn_neurons>
    n_neurons: !ref <output_neurons>  # 40 phonemes + 1 blank + 1 err

# Model parameters
activation: !name:torch.nn.LeakyReLU
dnn_layers: 2
dnn_neurons: 384  
freeze_perceived_ssl: False
freeze_canonical_ssl: False
freeze_perceived_feature_extractor: True  # freeze the CNN extractor in wav2vec
freeze_canonical_feature_extractor: True         # Freeze Whisper encoder?

log_softmax: !new:speechbrain.nnet.activations.Softmax
    apply_log: True

ctc_cost: !name:speechbrain.nnet.losses.ctc_loss
    blank_index: !ref <blank_index>

ctc_cost_mispro: !name:speechbrain.nnet.losses.ctc_loss
    blank_index: !ref <blank_index>

# Outputs
output_neurons: 44 # l2arctic: 40phns(sil)+err+blank + eos + bos =44
blank_index: 0

model: !new:torch.nn.ModuleList
    - [!ref <enc>, !ref <ctc_lin>]

adam_opt_class: !name:torch.optim.Adam
    lr: !ref <lr>

pretrained_opt_class: !name:torch.optim.Adam
    lr: !ref <lr_pretrained>

checkpointer: !new:speechbrain.utils.checkpoints.Checkpointer
    checkpoints_dir: !ref <save_folder>
    recoverables:
        model: !ref <model>
        perceived_ssl: !ref <perceived_ssl>
        counter: !ref <epoch_counter>
    allow_partial_load: True
# canonical_ssl: !ref <canonical_ssl>

augmentation: !new:speechbrain.augment.time_domain.SpeedPerturb
    orig_freq: !ref <sample_rate>
    speeds: [95, 100, 105]

epoch_counter: !new:speechbrain.utils.epoch_loop.EpochCounter
    limit: !ref <number_of_epochs>

train_logger: !new:speechbrain.utils.train_logger.FileTrainLogger
    save_file: !ref <train_log>

ctc_stats: !name:speechbrain.utils.metric_stats.MetricStats
    metric: !name:speechbrain.nnet.losses.ctc_loss
        blank_index: !ref <blank_index>
        reduction: batch

per_stats: !name:speechbrain.utils.metric_stats.ErrorRateStats


# # TIMIT
# timit_local_data_folder: "/common/db/TIMIT"  # Path to TIMIT datase

seed: 3047
__set_seed: !apply:torch.manual_seed [!ref <seed>]

# training parameters
number_of_epochs: 100
batch_size: 16
lr: 0.0003
sorting: ascending
sample_rate: 16000
gradient_accumulation: 2
lr_pretrained: 0.00001

# Mix-Precision Training
auto_mix_prec: true
# or
precision: fp16         # 支持 "fp32"、"fp16" 或 "bf16"
eval_precision: fp32    # 推理同样切换到 FP16

# Dataloader options
train_dataloader_opts:
    batch_size: !ref <batch_size>
    

valid_dataloader_opts:
    batch_size: !ref <batch_size>
    

test_dataloader_opts:
    batch_size: 1

pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
    collect_in: !ref <save_folder>/
    loadables:
        perceived_ssl: !ref <perceived_ssl>
        model: !ref <model>
        tokenizer: !ref <tokenizer>

encoder: !new:speechbrain.nnet.containers.LengthsCapableSequential
    perceived_ssl: !ref <perceived_ssl>
    enc: !ref <enc>
    ctc_lin: !ref <ctc_lin>
    log_softmax: !ref <log_softmax>

decoding_function: !name:speechbrain.decoders.ctc_greedy_decode
    blank_id: !ref <blank_index>

tokenizer: !new:speechbrain.dataio.encoder.CTCTextEncoder

modules:
  encoder: !ref <encoder>