# ============================================================
# Smoke Test Config — 1000 steps with checkpoint save/resume
# ============================================================

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_smoke.parquet"
  max_mel_frames: 3500
  max_tokens: 384
  max_batch_mel_frames: 200000
  max_batch_utterances: 4         # hard cap to keep decoder logits and attention memory bounded
  temperature: 5.0
  shuffle_buffer_size: 2000       # smaller buffer for faster startup
  sort_buffer_for_batching: true
  seed: 42

num_workers: 2                    # reduced for smoke test (2 * 8 GPUs = 16 workers)

# --- Training ---
epochs: 2
learning_rate: 5.0e-5
min_lr: 1.0e-6
warmup_steps: 100
weight_decay: 0.01
gradient_clip: 1.0
gradient_accumulation_steps: 2
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                    # disabled for smoke test to save memory (EMA copies full model)
ema_decay: 0.999

# --- Precision ---
bf16: true

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

# --- Checkpointing ---
output_dir: "/workspace/maya-asr/checkpoints_smoke"
save_every_steps: 250
eval_every_steps: 500
log_every_steps: 10
max_checkpoints: 5

# --- Logging ---
wandb_project: "maya-asr-cohere-transcribe"
wandb_run_name: "smoke-test-1000steps"

# --- Max steps (smoke test) ---
max_steps: 1000

# --- Resume ---
resume_from: null
