# ============================================================
# Cohere Transcribe Finetuning — Indic Language Extension
# Hardware: 8x H200 (141GB HBM3e each)
# Data: 73M utterances / 145K hours (12 languages)
# Model: 59,554 vocab (16K original + 43K Indic subwords)
# ============================================================

model_name: "/workspace/training/tokenizer_extension/extended_model"
trust_remote_code: true

# --- Data ---
data:
  mel_shards_dir: "/workspace/maya-asr/mel_shards_128ch/per_shard"
  manifest_path: "/workspace/maya-asr/manifests/training_manifest_cleaned.parquet"
  max_mel_frames: 3500
  max_tokens: 384
  max_batch_mel_frames: 1200000   # was 800K — fatter batches, fewer sync points
  max_batch_utterances: 128       # was 64 — use more of the 140GB/GPU
  temperature: 5.0
  shuffle_buffer_size: 4000
  sort_buffer_for_batching: true
  seed: 42

num_workers: 8                    # more workers to keep pipeline fed

# --- Training ---
epochs: 2
learning_rate: 2.0e-4
min_lr: 1.0e-6
warmup_steps: 0
weight_decay: 0.01
gradient_clip: 1.0
gradient_accumulation_steps: 1    # was 2 — single fat step, no accumulation overhead
label_smoothing: 0.1

# --- Layer-wise LR Decay ---
llrd_factor: 0.9

# --- Regularization ---
dropout_rate: 0.05
use_spec_augment: true

# --- EMA ---
use_ema: false
ema_decay: 0.999

# --- Precision ---
bf16: true

# --- Distributed ---
strategy: "ddp"
gradient_checkpointing: true

# --- Checkpointing ---
output_dir: "/workspace/maya-asr/checkpoints"
save_every_steps: 2000
eval_every_steps: 10000
log_every_steps: 50
max_checkpoints: 7

# --- Logging ---
wandb_project: "maya-asr-cohere-transcribe"
wandb_run_name: "indic-lang-extension-v1"

# --- Max steps ---
max_steps: 0

# --- Resume ---
resume_from: null
