rustytorch/crates/models/rtx-csm/examples/lora_train_emotional.rs

//! Multi-stage LoRA fine-tune driven by a JSONL manifest.
//!
//! The manifest format is one JSON object per line:
//! ```jsonl
//! {"wav": "audiobook/0001.wav", "transcript": "Once upon a time...", "stage": "audiobook", "emotion_tag": "[neutral]"}
//! {"wav": "podcast/0123.wav",   "transcript": "So like, anyway...",   "stage": "podcast",   "emotion_tag": "[casual]"}
//! {"wav": "va/excited_07.wav",  "transcript": "We did it!",           "stage": "va",        "emotion_tag": "[excited]"}
//! ```
//! `emotion_tag` and `stage` are optional. Wav paths can be absolute or
//! relative to the manifest's directory.
//!
//! Curriculum (per `docs/personal_voice_training_guide.md` §4):
//!   - stage "audiobook": 3 epochs at peak_lr 1e-4 (clean priors)
//!   - stage "podcast":   1 epoch  at peak_lr 3e-5 (mixed)
//!   - stage "va":        1 epoch  at peak_lr 1e-5 (acted, noisier)
//! After each stage, an intermediate adapter snapshot is written so you can
//! A/B subsequent stages.
//!
//! Usage:
//!   cargo run -p rtx-csm --release --features metal \
//!     --example lora_train_emotional -- \
//!       --manifest /data/voice/manifest.jsonl \
//!       --output /tmp/voice_final.safetensors \
//!       --extended-lora

use anyhow::Result;
use candle_nn::VarMap;
use clap::Parser;
use rtx_csm::Generator;
use rtx_csm::lora::LoraConfig;
use rtx_csm::training::{
    CurriculumStage, CurriculumTrainer, LoraAdapterMetadata, TrainingConfig, TrainingDataset,
    save_lora_adapter_with_metadata,
};
use std::path::PathBuf;

#[derive(Debug, Parser)]
struct Cli {
    /// JSONL manifest of (wav, transcript, [stage], [emotion_tag], [speaker]) rows.
    #[arg(long)]
    manifest: PathBuf,

    /// Final adapter output (safetensors). Per-stage snapshots are written
    /// next to it as `<stem>.<stage>.safetensors`.
    #[arg(long)]
    output: PathBuf,

    /// Default speaker id when a manifest row doesn't specify one.
    #[arg(long, default_value_t = 0)]
    speaker: u32,

    #[arg(long, default_value_t = 8)]
    rank: usize,

    #[arg(long, default_value_t = 16.0)]
    alpha: f32,

    /// Frames sampled per training step.
    #[arg(long, default_value_t = 4)]
    frames_per_step: usize,

    #[arg(long, default_value_t = 1.0)]
    grad_clip: f64,

    #[arg(long, default_value_t = 42)]
    seed: u64,

    /// Use Phase 12.1 extended LoRA coverage (q+k+v+output_proj + MLP)
    /// instead of q+v only. Recommended for emotional fine-tunes —
    /// the FFN carries prosodic style and a wider adapter has the
    /// capacity to learn multiple emotion tags simultaneously.
    #[arg(long, default_value_t = false)]
    extended_lora: bool,

    /// Override the canned per-stage `peak_lr`. The default recipe
    /// (1e-4 / 3e-5 / 1e-5 across stages) over-fits noticeably on
    /// small (~100-200 clip) corpora at rank 8. Try `--peak-lr 2e-5`
    /// for a gentler run when held-out loss is going up.
    #[arg(long)]
    peak_lr: Option<f64>,

    /// Override the canned per-stage epoch count. Default is
    /// `audiobook: 3, podcast: 1, va: 1`. A single epoch is often
    /// enough on small corpora.
    #[arg(long)]
    epochs: Option<usize>,

    #[arg(long)]
    cpu: bool,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();

    let device = if cli.cpu {
        candle_core::Device::Cpu
    } else {
        Generator::default_device()?
    };
    println!("device: {device:?}");

    let mut generator = Generator::load_csm_1b(&device)?;
    println!("model loaded");

    println!("loading manifest from {}", cli.manifest.display());
    let dataset = TrainingDataset::load_from_manifest(&cli.manifest, cli.speaker, &mut generator)?;
    println!("dataset: {} examples", dataset.len());

    // Inject LoRA into the backbone.
    let base = if cli.extended_lora {
        LoraConfig::extended()
    } else {
        LoraConfig::default()
    };
    let lora_cfg = LoraConfig {
        rank: cli.rank,
        alpha: cli.alpha,
        ..base
    };
    println!("LoRA target_modules={:?}", lora_cfg.target_modules);
    let vm = VarMap::new();
    generator.model.inner.add_lora_to_backbone(&lora_cfg, &vm)?;
    let n_params: usize = vm.all_vars().iter().map(|v| v.shape().elem_count()).sum();
    println!(
        "LoRA injected: {} trainable params ({} adapter Vars, rank={} alpha={})",
        n_params,
        vm.all_vars().len(),
        cli.rank,
        cli.alpha,
    );

    // Build per-stage checkpoint paths from the final-output stem.
    let stem = cli
        .output
        .file_stem()
        .map(|s| s.to_string_lossy().to_string())
        .unwrap_or_else(|| "voice".into());
    let parent = cli.output.parent().unwrap_or(std::path::Path::new("."));
    let stage_path = |stage: &str| parent.join(format!("{stem}.{stage}.safetensors"));

    // Canonical 3-stage recipe from personal_voice_training_guide.md §4.
    // CLI overrides (--peak-lr, --epochs) take precedence per-stage.
    let stages = vec![
        CurriculumStage {
            name: "audiobook".into(),
            config: TrainingConfig {
                epochs: cli.epochs.unwrap_or(3),
                peak_lr: cli.peak_lr.unwrap_or(1e-4),
                end_lr: 1e-5,
                warmup_steps: 32,
                grad_clip: Some(cli.grad_clip),
                frames_per_step: cli.frames_per_step,
                seed: cli.seed,
            },
            checkpoint: Some(stage_path("audiobook")),
        },
        CurriculumStage {
            name: "podcast".into(),
            config: TrainingConfig {
                epochs: cli.epochs.unwrap_or(1),
                peak_lr: cli.peak_lr.unwrap_or(3e-5),
                end_lr: 1e-5,
                warmup_steps: 16,
                grad_clip: Some(cli.grad_clip),
                frames_per_step: cli.frames_per_step,
                seed: cli.seed.wrapping_add(1),
            },
            checkpoint: Some(stage_path("podcast")),
        },
        CurriculumStage {
            name: "va".into(),
            config: TrainingConfig {
                epochs: cli.epochs.unwrap_or(1),
                peak_lr: cli.peak_lr.unwrap_or(1e-5),
                end_lr: 1e-6,
                warmup_steps: 16,
                grad_clip: Some(cli.grad_clip),
                frames_per_step: cli.frames_per_step,
                seed: cli.seed.wrapping_add(2),
            },
            checkpoint: Some(stage_path("va")),
        },
    ];

    let mut runner = CurriculumTrainer::new(&mut generator, &vm, &dataset, stages);
    let traces = runner.run()?;

    for (i, losses) in traces.iter().enumerate() {
        if losses.is_empty() {
            continue;
        }
        let n = losses.len();
        let head: f32 = losses.iter().take(10).sum::<f32>() / 10.0_f32.min(n as f32);
        let tail: f32 = losses.iter().rev().take(10).sum::<f32>() / 10.0_f32.min(n as f32);
        println!(
            "stage {i}: {n} steps, first 10 avg = {head:.4}, last 10 avg = {tail:.4}, change = {:+.2}%",
            100.0 * (tail - head) / head
        );
    }

    let metadata = LoraAdapterMetadata::from_lora_config(&lora_cfg);
    save_lora_adapter_with_metadata(&vm, &cli.output, &metadata)?;
    println!(
        "\n✓ trained emotional LoRA adapter saved to {} (with embedded metadata)",
        cli.output.display()
    );
    println!(
        "Apply at inference: examples/generate --lora {} --text \"...\" \\\n\
         \t--emotion-hint \"[your_tag]\" --out /tmp/out.wav\n\
         (rank/alpha/extended auto-detected from embedded metadata)",
        cli.output.display(),
    );
    Ok(())
}