use burn::backend::Wgpu;
use quiet_crab::{
    inference::Transcriber,
    model::{
        ModelSize, WhisperModel,
        config::{WhisperConfig, WhisperModelParams},
    },
    tokenizer::whisper_tokenizer::WhisperTokenizer,
};

type Backend = Wgpu;

fn main() -> anyhow::Result<()> {
    let device = Default::default();

    // Specify Options
    let config = WhisperConfig {
        model_params: WhisperModelParams::from_size(ModelSize::Medium),
        no_timestamps: false,
        language: Some(String::from("en")),
    };

    println!("Loading model weights...");
    let model = WhisperModel::<Backend>::from_safetensors(
        &config.model_params,
        "whisper-medium.safetensors",
        &device,
    )?;

    println!("Loading tokenizer...");
    let tokenizer = WhisperTokenizer::from_file("/home/jesse/code/rust/quiet_crab/tokenizer.json")?;

    let transcriber = Transcriber::new(model, tokenizer, config, device);

    println!("Transcribing...");
    let text = transcriber.transcribe("sample.mp3")?;

    println!("{text}");
    Ok(())
}