# Stage 1 Production Config: FastConformer Hybrid TDT+CTC on full dataset
# Hardware: 8x NVIDIA H200 (140 GB HBM3 each)
# Target: ~200K steps, ~5-7 days
#
# Pretrained encoder: Parakeet TDT 1.1B (42-layer FastConformer, English)
# Tokenizer: encoder-side 32K BPE (quality-filtered text, 11 Indic + English)
# Decoder/joint: randomly initialized for 32K vocab, then finetuned
#
# Optimizations v2:
#   - TDT loss (matching pretrained encoder objective) instead of RNNT
#   - Differential LR: encoder=2e-5, decoder/joint/CTC=1e-3
#   - Encoder frozen for first 5K steps, then unfrozen at low LR
#   - CTC loss warmup: 0 -> 0.3 over first 3K steps
#   - Language embedding conditioning (lang_id -> encoder)

name: maya-asr-stage1-prod-v2

trainer:
  devices: 8
  num_nodes: 1
  accelerator: gpu
  strategy: ddp_find_unused_parameters_true
  precision: bf16-mixed
  max_steps: 200000
  val_check_interval: 10000
  accumulate_grad_batches: 4  # effective = 8 GPUs x batch x 4 accum
  gradient_clip_val: 1.0
  log_every_n_steps: 100
  enable_checkpointing: false  # handled by exp_manager
  default_root_dir: experiments/stage1_prod

model:
  model_defaults:
    asr_enc_hidden: 1024
    enc_hidden: 1024
    dec_hidden: 640
    pred_hidden: 640

  preprocessor:
    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
    sample_rate: 16000
    normalize: per_feature
    window_size: 0.025
    window_stride: 0.01
    features: 80
    n_fft: 512
    dither: 0.00001

  encoder:
    _target_: nemo.collections.asr.modules.ConformerEncoder
    feat_in: 80
    n_layers: 42
    d_model: ${model.model_defaults.enc_hidden}
    subsampling: dw_striding
    subsampling_factor: 8
    subsampling_conv_channels: 256
    ff_expansion_factor: 4
    self_attention_model: rel_pos
    n_heads: 8
    conv_kernel_size: 9
    dropout: 0.1
    dropout_pre_encoder: 0.1
    dropout_emb: 0.0

  decoder:
    _target_: nemo.collections.asr.modules.RNNTDecoder
    prednet:
      pred_hidden: ${model.model_defaults.pred_hidden}
      pred_rnn_layers: 2
      dropout: 0.1

  joint:
    _target_: nemo.collections.asr.modules.RNNTJoint
    num_extra_outputs: 5  # TDT duration heads (durations=[0,1,2,3,4])
    fuse_loss_wer: true
    fused_batch_size: 2   # sub-batch joint computation → caps peak VRAM
    jointnet:
      joint_hidden: 640
      activation: relu
      dropout: 0.1
      encoder_hidden: ${model.model_defaults.enc_hidden}
      pred_hidden: ${model.model_defaults.pred_hidden}

  loss:
    loss_name: tdt
    tdt_kwargs:
      durations: [0, 1, 2, 3, 4]
      sigma: 0.02
      omega: 0.1
      fastemit_lambda: 0.0
      clamp: -1.0

  decoding:
    strategy: greedy
    model_type: tdt
    durations: [0, 1, 2, 3, 4]
    greedy:
      max_symbols: 10

  aux_ctc:
    ctc_loss_weight: 0.3  # warmed up from 0->0.3 over 3K steps in train_prod.py
    use_cer: false
    decoder:
      _target_: nemo.collections.asr.modules.ConvASRDecoder
      feat_in: null
      num_classes: -1

  # Encoder-side tokenizer (32K BPE, quality-filtered, 11 Indic + English)
  tokenizer:
    dir: tokenizers/stage1_prod_bpe
    type: bpe

  train_ds:
    # Production: tar-offset loading via train_prod.py (not NeMo tarred format)
    manifest_filepath: data/manifests/stage1_prod_val_v2.jsonl  # placeholder for model init
    batch_size: 4
    shuffle: true
    num_workers: 8
    max_duration: 20.0
    min_duration: 0.5
    sample_rate: 16000
    trim_silence: false

  validation_ds:
    manifest_filepath: data/manifests/stage1_prod_val_v2.jsonl
    batch_size: 4
    shuffle: false
    num_workers: 4
    sample_rate: 16000

  # NOTE: optim LR is overridden by differential LR in train_prod.py
  # This base LR is used for NeMo's scheduler init but actual param group
  # LRs are set in the training script.
  optim:
    name: adamw
    lr: 0.001  # decoder/joint/CTC LR (encoder uses 2e-5 via param groups)
    betas: [0.9, 0.98]
    weight_decay: 0.0001
    sched:
      name: CosineAnnealing
      warmup_steps: 2000
      min_lr: 0.00001

exp_manager:
  exp_dir: experiments/stage1_prod
  name: maya_asr_stage1_tdt_v2
  create_tensorboard_logger: true
  create_wandb_logger: true
  wandb_logger_kwargs:
    project: maya-asr
    name: tdt-parakeet-finetune-stage1-v2
    resume: allow
  resume_if_exists: false
  resume_ignore_no_checkpoint: true
  create_checkpoint_callback: true
  checkpoint_callback_params:
    monitor: val_loss
    mode: min
    save_top_k: 3
    save_last: true
    every_n_train_steps: 5000
    every_n_epochs: 0
    save_nemo_on_train_end: true
