rustytorch/crates/models/rtx-multimodal/tests/audio_tests.rs

use approx::assert_abs_diff_eq;
use rtx_multimodal::audio::conformer::*;
use rtx_multimodal::audio::whisper::*;
use rtx_multimodal::error::Result;
use rtx_tensor::{Device, Tensor};

#[test]
fn test_mel_spectrogram() -> Result<()> {
    let device = Device::default();
    let config = MelSpectrogramConfig {
        n_fft: 512,
        hop_length: 160,
        n_mels: 80,
        sample_rate: 16000,
        f_min: 0.0,
        f_max: 8000.0,
    };

    let mel_spec = MelSpectrogram::new(&config, &device)?;

    // Input audio: batch_size=2, samples=16000 (1 second at 16kHz)
    let audio = Tensor::randn(&[2, 16000], &device)?;
    let spectrogram = mel_spec.forward(&audio)?;

    // Expected: [batch, n_mels, time_frames]
    // time_frames = (16000 - 512) / 160 + 1 = 97
    assert_eq!(spectrogram.shape(), &[2, 80, 97]);

    Ok(())
}

#[test]
fn test_conformer_convolution_module() -> Result<()> {
    let device = Device::default();
    let config = ConformerConfig {
        d_model: 512,
        num_heads: 8,
        feed_forward_expansion_factor: 4,
        conv_kernel_size: 31,
        dropout: 0.1,
    };

    let conv_module = ConformerConvolutionModule::new(&config, &device)?;

    let input = Tensor::randn(&[2, 100, 512], &device)?;
    let output = conv_module.forward(&input)?;

    assert_eq!(output.shape(), &[2, 100, 512]);

    Ok(())
}

#[test]
fn test_conformer_feed_forward() -> Result<()> {
    let device = Device::default();
    let config = ConformerConfig {
        d_model: 512,
        num_heads: 8,
        feed_forward_expansion_factor: 4,
        conv_kernel_size: 31,
        dropout: 0.1,
    };

    let ff_module = ConformerFeedForward::new(&config, &device)?;

    let input = Tensor::randn(&[2, 100, 512], &device)?;
    let output = ff_module.forward(&input)?;

    assert_eq!(output.shape(), &[2, 100, 512]);

    Ok(())
}

#[test]
fn test_conformer_block() -> Result<()> {
    let device = Device::default();
    let config = ConformerConfig {
        d_model: 512,
        num_heads: 8,
        feed_forward_expansion_factor: 4,
        conv_kernel_size: 31,
        dropout: 0.1,
    };

    let block = ConformerBlock::new(&config, &device)?;

    let input = Tensor::randn(&[2, 100, 512], &device)?;
    let output = block.forward(&input)?;

    assert_eq!(output.shape(), &[2, 100, 512]);

    Ok(())
}

#[test]
fn test_conformer_encoder() -> Result<()> {
    let device = Device::default();
    let config = ConformerEncoderConfig {
        input_dim: 80, // Mel features
        d_model: 512,
        num_layers: 6,
        num_heads: 8,
        feed_forward_expansion_factor: 4,
        conv_kernel_size: 31,
        dropout: 0.1,
    };

    let encoder = ConformerEncoder::new(&config, &device)?;

    // Input: mel spectrogram features [batch, time, mel_features]
    let input = Tensor::randn(&[2, 100, 80], &device)?;
    let output = encoder.forward(&input)?;

    assert_eq!(output.shape(), &[2, 100, 512]);

    Ok(())
}

#[test]
fn test_whisper_encoder() -> Result<()> {
    let device = Device::default();
    let config = WhisperEncoderConfig {
        n_mels: 80,
        n_audio_ctx: 1500,
        n_audio_state: 512,
        n_audio_head: 8,
        n_audio_layer: 6,
    };

    let encoder = WhisperEncoder::new(&config, &device)?;

    // Input: mel spectrogram [batch, n_mels, time]
    let input = Tensor::randn(&[2, 80, 1500], &device)?;
    let output = encoder.forward(&input)?;

    // Output should be [batch, n_audio_ctx, n_audio_state]
    assert_eq!(output.shape(), &[2, 1500, 512]);

    Ok(())
}

#[test]
fn test_whisper_decoder() -> Result<()> {
    let device = Device::default();
    let encoder_config = WhisperEncoderConfig {
        n_mels: 80,
        n_audio_ctx: 1500,
        n_audio_state: 512,
        n_audio_head: 8,
        n_audio_layer: 6,
    };

    let decoder_config = WhisperDecoderConfig {
        n_vocab: 51865,
        n_text_ctx: 448,
        n_text_state: 512,
        n_text_head: 8,
        n_text_layer: 6,
    };

    let decoder = WhisperDecoder::new(&decoder_config, &device)?;

    // Encoder output (key-value for cross attention)
    let encoder_output = Tensor::randn(&[2, 1500, 512], &device)?;

    // Decoder input (token IDs)
    let tokens = Tensor::randint(0, 51865, &[2, 100], &device)?;

    let logits = decoder.forward(&tokens, &encoder_output)?;

    // Should output logits over vocabulary
    assert_eq!(logits.shape(), &[2, 100, 51865]);

    Ok(())
}

#[test]
fn test_whisper_full_model() -> Result<()> {
    let device = Device::default();
    let config = WhisperConfig {
        encoder: WhisperEncoderConfig {
            n_mels: 80,
            n_audio_ctx: 1500,
            n_audio_state: 512,
            n_audio_head: 8,
            n_audio_layer: 6,
        },
        decoder: WhisperDecoderConfig {
            n_vocab: 51865,
            n_text_ctx: 448,
            n_text_state: 512,
            n_text_head: 8,
            n_text_layer: 6,
        },
    };

    let model = WhisperModel::new(&config, &device)?;

    let audio_features = Tensor::randn(&[1, 80, 1500], &device)?;
    let tokens = Tensor::randint(0, 51865, &[1, 50], &device)?;

    let logits = model.forward(&audio_features, &tokens)?;
    assert_eq!(logits.shape(), &[1, 50, 51865]);

    Ok(())
}

#[test]
fn test_conformer_positional_encoding() -> Result<()> {
    let device = Device::default();
    let pos_encoding = ConformerPositionalEncoding::new(512, 5000, &device)?;

    let input = Tensor::randn(&[2, 100, 512], &device)?;
    let output = pos_encoding.forward(&input)?;

    assert_eq!(output.shape(), &[2, 100, 512]);

    // Test different sequence lengths
    let shorter_input = Tensor::randn(&[1, 50, 512], &device)?;
    let shorter_output = pos_encoding.forward(&shorter_input)?;
    assert_eq!(shorter_output.shape(), &[1, 50, 512]);

    Ok(())
}

#[test]
fn test_audio_preprocessing() -> Result<()> {
    let device = Device::default();
    let preprocessor = AudioPreprocessor::new(16000, &device)?;

    // Simulate raw audio samples
    let raw_audio = vec![0.1, 0.2, -0.1, 0.15, -0.05]; // Small sample
    let processed = preprocessor.preprocess(&raw_audio)?;

    assert_eq!(processed.shape()[0], 1); // Batch dimension
    assert!(processed.shape()[1] > 0); // Should have some samples

    Ok(())
}