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

//! Held-out cross-entropy evaluator for LoRA adapters.
//!
//! Runs teacher-forced `forward_loss` over a held-out manifest (same JSONL
//! format as `lora_train_emotional`) and writes a JSON report. Run it twice
//! for an A/B:
//!
//! ```bash
//! # baseline (no adapter)
//! lora_eval --eval-manifest test.jsonl --report /tmp/base.json
//!
//! # with adapter
//! lora_eval --eval-manifest test.jsonl --lora /tmp/voice.safetensors \
//!     --report /tmp/lora.json
//!
//! # compare summaries
//! diff <(jq .summary /tmp/base.json) <(jq .summary /tmp/lora.json)
//! ```
//!
//! Lower mean / median / p90 loss = adapter is fitting the held-out audio
//! better. If LoRA loss is higher than base, your adapter has overfit the
//! training set.

use anyhow::Result;
use clap::Parser;
use rtx_csm::Generator;
use rtx_csm::training::{EvalSummary, TrainingDataset, evaluate_held_out};
use std::path::PathBuf;

#[derive(Debug, Parser)]
struct Cli {
    /// JSONL manifest of held-out (wav, transcript, [emotion_tag]) rows.
    /// Same format as `lora_train_emotional`.
    #[arg(long)]
    eval_manifest: PathBuf,

    /// Output JSON report path.
    #[arg(long)]
    report: PathBuf,

    /// Optional LoRA adapter to apply before scoring. When omitted, the
    /// base CSM-1B is scored — that's the report you compare against.
    #[arg(long)]
    lora: Option<PathBuf>,

    /// Force LoRA hyperparameters (overrides metadata). See `apply_lora_adapter`.
    #[arg(long)]
    lora_rank: Option<usize>,
    #[arg(long)]
    lora_alpha: Option<f32>,
    #[arg(long, default_value_t = false)]
    extended_lora: bool,

    /// Optional quantized GGUF (Q8/Q4_K_M). Combines with `--lora` per Phase 12.6.
    #[arg(long)]
    quantized_gguf: Option<PathBuf>,

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

    /// Frames sampled per example. 0 = score every frame (slow).
    #[arg(long, default_value_t = 8)]
    frames_per_example: usize,

    /// RNG seed for frame sampling. Two runs with the same seed score the
    /// same frames in the same order — required for a fair A/B.
    #[arg(long, default_value_t = 42)]
    seed: u64,

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

#[derive(Debug, serde::Serialize)]
struct Report {
    /// What was evaluated: base or path-to-adapter (or path with quantized base).
    label: String,
    n_examples: usize,
    frames_per_example_cap: usize,
    seed: u64,
    summary: EvalSummary,
    rows: Vec<rtx_csm::training::EvalRow>,
}

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 = if let Some(gguf) = cli.quantized_gguf.as_ref() {
        Generator::load_csm_1b_quantized(gguf, &device, /* enable_cfg */ false)?
    } else {
        Generator::load_csm_1b(&device)?
    };
    println!("model loaded");

    let label = if let Some(lora_path) = cli.lora.as_ref() {
        let extended_override = if cli.extended_lora { Some(true) } else { None };
        rtx_csm::training::apply_lora_adapter(
            &mut generator,
            lora_path,
            cli.lora_rank,
            cli.lora_alpha,
            extended_override,
            &device,
        )?;
        format!("lora:{}", lora_path.display())
    } else {
        "base".to_string()
    };

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

    let rows = evaluate_held_out(&mut generator, &dataset, cli.frames_per_example, cli.seed)?;
    let summary = EvalSummary::from_rows(&rows);
    println!(
        "summary: n={} frames={} mean={:.4} median={:.4} p90={:.4} min={:.4} max={:.4}",
        summary.n_examples,
        summary.n_frames_total,
        summary.mean_loss,
        summary.median_loss,
        summary.p90_loss,
        summary.min_loss,
        summary.max_loss,
    );

    let report = Report {
        label,
        n_examples: rows.len(),
        frames_per_example_cap: cli.frames_per_example,
        seed: cli.seed,
        summary,
        rows,
    };
    let json = serde_json::to_string_pretty(&report)?;
    std::fs::write(&cli.report, json)?;
    println!("\n✓ wrote report to {}", cli.report.display());
    Ok(())
}