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

//! Phase 8.5 smoke test for the Moonshine conv stem.
//!
//! Verifies: weight loading from HF safetensors works, conv shapes
//! produce expected output dimensions, no panics on a realistic input.
//! Stops short of the encoder transformer block (Phase 8.6).
//!
//! Usage:
//! ```bash
//! cargo run -p rtx-csm --release --features metal --example moonshine_smoke
//! ```

use anyhow::{Context, Result};
use candle_core::{DType, Device, Tensor};
use hf_hub::api::sync::Api;
use rtx_csm::moonshine;

fn main() -> Result<()> {
    let device = if candle_core::utils::metal_is_available() {
        Device::new_metal(0)?
    } else {
        Device::Cpu
    };
    eprintln!("device: {device:?}");

    // Download weights (cached after first run).
    let api = Api::new()?;
    let weights = api
        .model("UsefulSensors/moonshine-tiny".to_string())
        .get("model.safetensors")?;
    eprintln!("weights: {}", weights.display());

    // Build the full encoder + decoder (Phase 8.5/8.6/8.7).
    let cfg = moonshine::MoonshineConfig::tiny();
    let (encoder, decoder) = moonshine::load_full(&weights, &device, &cfg).context("load full")?;
    eprintln!("encoder + decoder loaded");

    // Also keep the standalone conv stem for the conv-only check below.
    let stem = moonshine::load_conv_stem(&weights, &device).context("load conv stem")?;

    // Synthetic 10 s of 16 kHz audio (silence + a sine pulse).
    let sr = 16_000usize;
    let n = sr * 10;
    let mut samples = vec![0.0f32; n];
    let freq = 440.0;
    for (i, s) in samples.iter_mut().enumerate() {
        let t = i as f32 / sr as f32;
        // 1-3 s active sine, rest silence
        if (1.0..3.0).contains(&t) {
            *s = (2.0 * std::f32::consts::PI * freq * t).sin() * 0.3;
        }
    }
    let pcm = Tensor::from_vec(samples, (1, 1, n), &device).context("pcm tensor")?;
    eprintln!("input shape: {:?}", pcm.shape());

    // Forward.
    let out = stem.forward(&pcm).context("conv stem forward")?;
    eprintln!("output shape: {:?}", out.shape());

    // Expected: (B=1, T_seq, 288). T_seq ≈ n / 384.
    // Conv arithmetic (output_len = (input_len - kernel) / stride + 1):
    //   conv1: (160000 - 127) / 64 + 1 = 2498
    //   conv2: (2498 - 7) / 3 + 1     = 831
    //   conv3: (831 - 3) / 2 + 1      = 415
    let dims = out.dims();
    let expected_t_seq = ((((n - 127) / 64 + 1) - 7) / 3 + 1 - 3) / 2 + 1;
    eprintln!(
        "expected (B=1, T_seq={expected_t_seq}, hidden=288); got {:?}",
        dims
    );
    if dims == [1, expected_t_seq, 288] {
        println!("PASS: conv stem forward matches expected shape");
    } else {
        println!(
            "MISMATCH: expected [1, {expected_t_seq}, 288], got {:?}",
            dims
        );
        std::process::exit(2);
    }

    // Check the output isn't all zeros (sanity).
    let max = out.abs()?.max_keepdim(0)?.max_keepdim(1)?.max_keepdim(2)?;
    let max_val: f32 = max.flatten_all()?.to_vec1::<f32>()?[0];
    eprintln!("output max abs: {max_val:.4}");
    if max_val < 1e-6 {
        println!("WARN: output is all near-zero — conv weights may not be loading");
    } else {
        println!("output has signal — weight loading verified");
    }

    // Now run through the FULL encoder (stem + 6 transformer layers + LN).
    eprintln!();
    eprintln!("=== full encoder forward ===");
    let full_t = std::time::Instant::now();
    let enc_out = encoder.forward(&pcm).context("encoder forward")?;
    let full_ms = full_t.elapsed().as_millis();
    eprintln!("forward: {full_ms} ms");
    eprintln!("encoder output shape: {:?}", enc_out.shape());
    let enc_dims = enc_out.dims();
    if enc_dims == [1, expected_t_seq, 288] {
        println!("PASS: encoder output preserves (B, T_seq, 288) shape");
    } else {
        println!(
            "MISMATCH: expected [1, {expected_t_seq}, 288], got {:?}",
            enc_dims
        );
        std::process::exit(2);
    }
    let enc_max = enc_out
        .abs()?
        .max_keepdim(0)?
        .max_keepdim(1)?
        .max_keepdim(2)?;
    let enc_max_val: f32 = enc_max.flatten_all()?.to_vec1::<f32>()?[0];
    let enc_mean = enc_out.mean_all()?;
    let enc_mean_val: f32 = enc_mean.to_vec0::<f32>()?;
    eprintln!("encoder output max abs: {enc_max_val:.4}");
    eprintln!("encoder output mean    : {enc_mean_val:.4}");
    if enc_max_val < 1e-6 {
        println!("WARN: encoder output is all near-zero");
    } else {
        println!("encoder output has signal — full transformer pipeline verified");
    }

    // Phase 8.7: decoder forward step. Feed [bos] + encoder output,
    // expect logits over the 32 768 vocab. We don't care about the
    // exact predicted token yet (parity check is Phase 8.9) — just
    // shape + signal.
    eprintln!();
    eprintln!("=== decoder forward step ===");
    let bos = cfg.bos_token_id;
    let tokens = Tensor::from_vec(vec![bos], (1, 1), &device).context("bos tokens")?;
    let dec_t = std::time::Instant::now();
    let logits = decoder.forward(&tokens, &enc_out).context("decoder")?;
    let dec_ms = dec_t.elapsed().as_millis();
    eprintln!("forward: {dec_ms} ms");
    eprintln!("decoder logits shape: {:?}", logits.shape());
    let logit_dims = logits.dims();
    if logit_dims == [1, 1, cfg.vocab_size] {
        println!("PASS: decoder produces (B=1, T=1, vocab=32768) logits");
    } else {
        println!(
            "MISMATCH: expected [1, 1, {}], got {:?}",
            cfg.vocab_size, logit_dims
        );
        std::process::exit(2);
    }
    let logit_max = logits
        .abs()?
        .max_keepdim(0)?
        .max_keepdim(1)?
        .max_keepdim(2)?;
    let logit_max_val: f32 = logit_max.flatten_all()?.to_vec1::<f32>()?[0];
    eprintln!("logit max abs: {logit_max_val:.4}");
    // Greedy argmax over the vocab dim — what's the model's first
    // prediction given silence-mostly audio?
    let last_step = logits.narrow(1, 0, 1)?.squeeze(1)?; // (1, vocab)
    let argmax = last_step.argmax(1)?;
    let token_id: u32 = argmax.to_dtype(DType::U32)?.to_vec1::<u32>()?[0];
    eprintln!("argmax token_id: {token_id} (eos={})", cfg.eos_token_id);

    Ok(())
}