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

//! End-to-end raw-audio → training-manifest pipeline.
//!
//! Closes the data-prep loop entirely in-crate (no Python, no ort, no
//! whisper.cpp): given a single wav of raw recorded audio (e.g. a podcast,
//! audiobook chapter, conversation), emits a JSONL manifest ready for
//! `examples/lora_train_emotional`.
//!
//! Pipeline:
//!   1. Diarize the input (Phase 13.1: Silero V5 + WavLM-SV + clustering)
//!   2. For each segment, slice the audio + transcribe via Moonshine (English)
//!   3. Write per-segment clip wavs and a `manifest.jsonl` row each
//!
//! Usage:
//!   cargo run -p rtx-csm --release --features metal --example audio_to_manifest -- \
//!       --in /data/podcast.wav \
//!       --out-dir /tmp/podcast_clips \
//!       --wavlm-sv-weights /tmp/wavlm_sv.safetensors
//!
//! Optional knobs let you carry a static `--emotion-tag` and `--stage` label
//! across every emitted row (handy when the source is uniformly e.g. a
//! single audiobook). Per-row labels can be overridden with a post-process
//! script — the manifest is just JSONL.

use anyhow::{Context, Result};
use candle_core::{DType, Device, Tensor};
use clap::Parser;
use hf_hub::api::sync::Api;
use rtx_csm::Generator;
use rtx_csm::audio_io;
use rtx_csm::diarize::{DiarizationConfig, DiarizedSegment, Diarizer, SAMPLE_RATE};
use rtx_csm::moonshine;
use rtx_csm::silero_vad::SileroVad;
use rtx_csm::wavlm_sv::WavLmSv;
use std::path::PathBuf;

#[derive(Debug, Parser)]
struct Cli {
    /// Source wav. Any sample rate; resampled internally to 16 kHz.
    #[arg(long = "in")]
    input: PathBuf,

    /// Output directory. Will be created if missing. Per-segment clips
    /// land here as `<stem>.spk{N}.{idx:04}.wav` and the manifest as
    /// `manifest.jsonl`.
    #[arg(long)]
    out_dir: PathBuf,

    /// Converted WavLM-SV safetensors. See `examples/wavlm_sv_convert`.
    #[arg(long)]
    wavlm_sv_weights: PathBuf,

    /// Pre-computed segments JSON (output of `examples/diarize`). When set,
    /// skip the diarization step and use this instead. Useful for
    /// re-transcribing without re-clustering.
    #[arg(long)]
    segments_json: Option<PathBuf>,

    /// Static emotion tag added to every manifest row (e.g. `[neutral]`).
    /// Per-row override is expected via post-processing. When
    /// `--auto-emotion-tag` is also set, the auto label takes precedence.
    #[arg(long)]
    emotion_tag: Option<String>,

    /// Auto-tag every row with a per-segment emotion label. Default
    /// classifier is the Phase 13.3 prosody-rule placeholder (RMS + F0
    /// + voicing → 5 buckets); pair with `--use-emotion2vec` for the
    /// real emotion2vec_plus_base classifier (Phase 13.8).
    #[arg(long, default_value_t = false)]
    auto_emotion_tag: bool,

    /// When set together with `--auto-emotion-tag`, swap the prosody-rule
    /// classifier for the emotion2vec_plus_base candle port. ~93 M
    /// params, downloaded once from HF Hub (~1.1 GB). Per-segment cost
    /// is ~150 ms on Metal vs <1 ms for the prosody rule, but accuracy
    /// is comparable to a real SOTA SER model on conversational audio.
    #[arg(long, default_value_t = false)]
    use_emotion2vec: bool,

    /// Static curriculum stage label added to every manifest row
    /// (e.g. `audiobook` for the Phase 12.3 recipe).
    #[arg(long)]
    stage: Option<String>,

    /// Default speaker id when the diarizer outputs a single cluster
    /// — used to anchor the manifest to a specific CSM speaker slot.
    #[arg(long, default_value_t = 0)]
    speaker_offset: u32,

    /// Maximum decoded tokens per segment (Moonshine config caps at 194).
    #[arg(long, default_value_t = 100)]
    max_tokens: usize,

    /// Drop segments whose transcript is shorter than this many chars
    /// (filters out segments where ASR fails or returns punctuation only).
    #[arg(long, default_value_t = 4)]
    min_transcript_chars: usize,

    /// Diarization knobs (forwarded to Diarizer; ignored when
    /// `--segments-json` is set).
    #[arg(long, default_value_t = 0.5)]
    vad_threshold: f32,
    #[arg(long, default_value_t = 0.5)]
    cluster_threshold: f32,
    #[arg(long, default_value_t = 2.0)]
    window_s: f32,
    #[arg(long, default_value_t = 1.0)]
    hop_s: f32,
    #[arg(long, default_value_t = 0.5)]
    min_segment_s: f32,
    #[arg(long)]
    n_speakers: Option<usize>,

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

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

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

    std::fs::create_dir_all(&cli.out_dir).context("create out_dir")?;

    // 1. Load 16 kHz audio
    let audio = audio_io::load_mono_at_rate(&cli.input, SAMPLE_RATE)?;
    let total_secs = audio.len() as f32 / SAMPLE_RATE as f32;
    println!(
        "loaded {} samples ({:.2}s at {} Hz)",
        audio.len(),
        total_secs,
        SAMPLE_RATE,
    );

    // 2. Get segments — either from a precomputed JSON or by running diarization
    let segments: Vec<DiarizedSegment> = if let Some(seg_path) = cli.segments_json.as_ref() {
        let body = std::fs::read_to_string(seg_path).context("read segments_json")?;
        serde_json::from_str(&body).context("parse segments_json")?
    } else {
        let vad_path = rtx_csm::silero_vad::ensure_default_weights()?;
        let vad = SileroVad::load_from_file(&vad_path, &device)?;
        let vb = unsafe {
            candle_nn::VarBuilder::from_mmaped_safetensors(
                &[&cli.wavlm_sv_weights],
                DType::F32,
                &device,
            )
        }
        .context("opening WavLM-SV safetensors")?;
        let sv = WavLmSv::new(vb).context("WavLmSv::new")?;
        let mut diarizer = Diarizer::new(vad, sv, device.clone());
        let cfg = DiarizationConfig {
            vad_threshold: cli.vad_threshold,
            cluster_threshold: cli.cluster_threshold,
            window_s: cli.window_s,
            hop_s: cli.hop_s,
            n_speakers: cli.n_speakers,
            min_segment_s: cli.min_segment_s,
            ..DiarizationConfig::default()
        };
        let t0 = std::time::Instant::now();
        let segs = diarizer.diarize(&audio, &cfg)?;
        println!(
            "diarized {} segments in {} ms",
            segs.len(),
            t0.elapsed().as_millis()
        );
        segs
    };
    if segments.is_empty() {
        anyhow::bail!("diarization produced no segments — check VAD threshold or input audio");
    }
    let n_speakers: std::collections::HashSet<_> = segments.iter().map(|s| s.speaker).collect();
    println!(
        "{} speaker(s), {} segment(s)",
        n_speakers.len(),
        segments.len()
    );

    // 3. Load Moonshine
    let api = Api::new()?;
    let repo = api.model("UsefulSensors/moonshine-tiny".to_string());
    let weights = repo.get("model.safetensors")?;
    let tok_path = repo.get("tokenizer.json")?;
    let cfg = moonshine::MoonshineConfig::tiny();
    let (encoder, decoder) = moonshine::load_full(&weights, &device, &cfg)?;
    let tok =
        moonshine::load_tokenizer(&tok_path).map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
    println!("Moonshine loaded");

    // Optional: per-segment auto-labeling. Prosody-rule by default; the
    // emotion2vec_plus_base candle port is opt-in via --use-emotion2vec.
    // Boxed as `dyn EmotionDetector` so both paths share the call site.
    let ser: Option<Box<dyn rtx_csm::ser::EmotionDetector>> = if cli.auto_emotion_tag {
        if cli.use_emotion2vec {
            let api = hf_hub::api::sync::Api::new()?;
            let path = api
                .model("emotion2vec/emotion2vec_plus_base".into())
                .get("model.pt")
                .context("download emotion2vec_plus_base/model.pt")?;
            let model = rtx_csm::emotion2vec::Emotion2Vec::load_from_pickle(&path, &device)?;
            tracing::info!("loaded emotion2vec_plus_base from {}", path.display());
            Some(Box::new(model))
        } else {
            Some(Box::new(rtx_csm::ser::ProsodyDetector::default()))
        }
    } else {
        None
    };

    // 4. Per-segment transcribe + write
    let stem = cli
        .input
        .file_stem()
        .map(|s| s.to_string_lossy().to_string())
        .unwrap_or_else(|| "audio".into());
    let manifest_path = cli.out_dir.join("manifest.jsonl");
    let mut manifest_buf = String::new();
    let mut kept = 0usize;
    let mut dropped = 0usize;
    for (i, seg) in segments.iter().enumerate() {
        let s_idx = (seg.start_s * SAMPLE_RATE as f32) as usize;
        let e_idx = ((seg.end_s * SAMPLE_RATE as f32) as usize).min(audio.len());
        if e_idx <= s_idx {
            dropped += 1;
            continue;
        }
        let slice = &audio[s_idx..e_idx];

        // Moonshine encode + decode
        let pcm_t = Tensor::from_vec(slice.to_vec(), (1, 1, slice.len()), &device)?;
        let enc = encoder.forward(&pcm_t)?;
        let token_ids = decoder.generate_cached(&enc, &cfg, cli.max_tokens)?;
        let transcript = tok
            .decode(&token_ids, true)
            .map_err(|e| anyhow::anyhow!("detok: {e}"))?;
        let transcript = transcript.trim().to_string();
        if transcript.chars().count() < cli.min_transcript_chars {
            dropped += 1;
            continue;
        }

        // Write the per-segment clip at the manifest's expected sample rate
        // (CSM training pipeline ingests at 24 kHz; load_from_manifest uses
        // load_mono_24k internally, which resamples). We can store the slice
        // at 16 kHz and the trainer will resample on load — same path the
        // Phase 12.3 trainer takes for any non-24k file.
        let clip_name = format!("{}.spk{}.{:04}.wav", stem, seg.speaker, i);
        let clip_path = cli.out_dir.join(&clip_name);
        audio_io::write_wav_mono(&clip_path, slice, SAMPLE_RATE)?;

        // ManifestRow shape (matches src/training.rs)
        #[derive(serde::Serialize)]
        struct Row<'a> {
            wav: &'a str,
            transcript: &'a str,
            #[serde(skip_serializing_if = "Option::is_none")]
            emotion_tag: Option<&'a str>,
            #[serde(skip_serializing_if = "Option::is_none")]
            stage: Option<&'a str>,
            speaker: u32,
        }
        // Auto-label takes precedence over the static --emotion-tag when
        // both are set. Both being None leaves the manifest row's tag null
        // (Phase 12.3 trainer treats that as no-emotion-conditioning).
        let auto_tag_owned: Option<String> = ser.as_ref().map(|d| {
            // d: &Box<dyn EmotionDetector>; method-style call
            // auto-derefs through the box to the trait impl.
            let label = d
                .classify(slice)
                .unwrap_or(rtx_csm::ser::EmotionLabel::Neutral);
            label.as_tag().to_string()
        });
        let resolved_tag = auto_tag_owned.as_deref().or(cli.emotion_tag.as_deref());

        let row = Row {
            // Manifest paths are resolved relative to the manifest's
            // directory by load_from_manifest, so a bare filename suffices.
            wav: &clip_name,
            transcript: &transcript,
            emotion_tag: resolved_tag,
            stage: cli.stage.as_deref(),
            speaker: cli.speaker_offset + seg.speaker as u32,
        };
        manifest_buf.push_str(&serde_json::to_string(&row)?);
        manifest_buf.push('\n');
        kept += 1;
        if kept <= 5 || kept % 25 == 0 {
            println!(
                "  [{kept}] spk{} {:.2}s..{:.2}s : {}",
                seg.speaker,
                seg.start_s,
                seg.end_s,
                if transcript.len() > 60 {
                    format!("{}…", &transcript[..60])
                } else {
                    transcript.clone()
                }
            );
        }
    }
    std::fs::write(&manifest_path, manifest_buf)?;
    println!(
        "\n✓ wrote {kept} rows ({dropped} dropped) → {}",
        manifest_path.display()
    );
    println!(
        "Train: examples/lora_train_emotional --manifest {} --output /tmp/voice.safetensors --extended-lora",
        manifest_path.display()
    );
    Ok(())
}