wavo 0.1.0

// Copyright (c) vladkens | MIT License | https://github.com/vladkens/wavo
//! `wavo`: speech to text with models pulled by name into the Hugging Face cache.

mod audio;
mod hfs;
mod models;
mod output;
mod split;

use std::path::{Path, PathBuf};
use std::process::ExitCode;
use std::time::Instant;

use anyhow::{Context, Result, anyhow, bail};
use models::{MODELS, Source};

const USAGE: &str = "\
Usage:
  wavo run [MODEL] AUDIO [--json | --srt] [--segment SECS]
                                 transcribe AUDIO, with parakeet-v3 unless MODEL is given;
                                 --json adds tokens with start times, --srt prints subtitles;
                                 longer audio is split at pauses into segments of up to SECS
                                 (25 for GigaAM, 30 for Whisper, else 60; 0 for one pass)
  wavo pull MODEL                download a model
  wavo list                      list downloaded models
  wavo rm MODEL                  delete a downloaded model
  wavo bench MODEL AUDIO [-n N]  time the model load, the first call and N more (10)

MODEL is a name below, its full name or a path to a .gguf file. AUDIO is wav, mp3, m4a, flac
or ogg (Vorbis), downmixed to mono and resampled to 16 kHz.
Models live in the Hugging Face cache (HF_HUB_CACHE or HF_HOME), shared with the hf CLI;
only `wavo pull` uses the network.

Models:";

/// `wavo run` splits audio into segments of up to this for models without a window of their own.
const SEGMENT_MS: u32 = 60_000;

fn usage() -> String {
  let mut usage = USAGE.to_string();
  for m in MODELS {
    let full = if m.name == m.variant { String::new() } else { format!(" ({})", m.variant) };
    usage += &format!("\n  {:<18} {}  {}{full}", m.name, models::gb(m.size), m.about);
  }
  usage
}

fn main() -> ExitCode {
  let args: Vec<String> = std::env::args().skip(1).collect();
  match cli(&args.iter().map(String::as_str).collect::<Vec<_>>()) {
    Ok(()) => ExitCode::SUCCESS,
    Err(e) => {
      eprintln!("error: {e:#}");
      ExitCode::FAILURE
    }
  }
}

fn cli(args: &[&str]) -> Result<()> {
  let help = args.iter().any(|a| matches!(*a, "-h" | "--help"));
  match args {
    [] | ["help"] => println!("{}", usage()),
    _ if help => println!("{}", usage()),
    ["-V" | "--version"] => println!("wavo {}", env!("CARGO_PKG_VERSION")),
    ["run", args @ ..] => run(args)?,
    ["pull", name] => pull(name)?,
    ["list" | "ls"] => list()?,
    ["rm", name] => rm(name)?,
    ["bench", model, audio] => bench(model, audio, 10)?,
    ["bench", model, audio, "-n", n] => bench(model, audio, n.parse().context("-n")?)?,
    _ => bail!("unknown command: wavo {}\n\n{}", args.join(" "), usage()),
  }
  Ok(())
}

/// A downloaded model by name, or a `.gguf` path. Never touches the network.
fn model_path(arg: &str) -> Result<PathBuf> {
  match models::resolve(arg)? {
    Source::Path(path) if path.is_file() => Ok(path),
    Source::Path(path) => bail!("{}: no such file", path.display()),
    Source::Model(m) => {
      let cache = hfs::Cache::new();
      cache.find(&m.repo(), &m.file()).ok_or_else(|| {
        let (name, size, root) = (m.name, models::gb(m.size), cache.describe());
        anyhow!("{name} is not in {root}; run: wavo pull {name} ({size})")
      })
    }
  }
}

fn load(path: &Path) -> Result<wavo::Model> {
  wavo::Model::load(path).with_context(|| format!("loading {}", path.display()))
}

fn run(args: &[&str]) -> Result<()> {
  const USAGE: &str = "usage: wavo run [MODEL] AUDIO [--json | --srt] [--segment SECS]";
  let (mut format, mut segment, mut rest) = (None, None, Vec::new());
  let mut args = args.iter();
  while let Some(&arg) = args.next() {
    match arg {
      "--json" | "--srt" if format.is_none() => format = Some(arg),
      "--segment" if segment.is_none() => {
        let secs = args.next().and_then(|s| s.parse::<u32>().ok());
        let secs = secs.filter(|&s| s == 0 || s >= 5).ok_or_else(|| {
          anyhow!("--segment takes whole seconds: 0 for one pass, else at least 5")
        })?;
        segment = Some(secs.saturating_mul(1000));
      }
      _ if arg.starts_with("--") => bail!(USAGE),
      _ => rest.push(arg),
    }
  }
  let (model, audio) = match *rest {
    [arg] if models::find(arg).is_some() && !Path::new(arg).exists() => {
      bail!("no audio file given: wavo run {arg} AUDIO")
    }
    [audio] => (models::DEFAULT, audio),
    [model, audio] => (model, audio),
    _ => bail!(USAGE),
  };
  let path = model_path(model)?;
  let pcm = audio::read(audio)?;
  let model = load(&path)?;
  let max_ms = segment.or(model.max_audio_ms()).unwrap_or(SEGMENT_MS);
  let max = if max_ms == 0 { usize::MAX } else { max_ms as usize * 16 };
  let ranges = split::segments(&pcm, max);
  let parts = ranges.into_iter().map(|r| Ok((r.start, model.transcribe(&pcm[r])?)));
  let (transcript, starts) = split::join(parts.collect::<Result<Vec<_>>>()?);
  match format {
    Some("--json") => println!("{}", output::json(&transcript)),
    // Whisper times segments, not tokens.
    Some(_) if models::architecture(&path).as_deref() == Some("whisper") => {
      print!("{}", output::srt_segments(&transcript.tokens, (pcm.len() / 16) as u32))
    }
    Some(_) => print!("{}", output::srt(&transcript.tokens, &starts, (pcm.len() / 16) as u32)),
    None => println!("{}", transcript.text),
  }
  Ok(())
}

/// A model by name, for the commands that manage the cache.
fn named(arg: &str) -> Result<&'static models::Model> {
  match models::resolve(arg)? {
    Source::Model(m) => Ok(m),
    Source::Path(_) => bail!("expected a model name, not a path: {arg}"),
  }
}

fn pull(arg: &str) -> Result<()> {
  let m = named(arg)?;
  let path = hfs::Cache::new().pull(&m.repo(), &m.file());
  println!("{}", path.with_context(|| format!("pulling {}", m.name))?.display());
  Ok(())
}

fn list() -> Result<()> {
  let cache = hfs::Cache::new();
  let found: Vec<_> =
    MODELS.iter().filter_map(|m| cache.find(&m.repo(), &m.file()).map(|path| (m, path))).collect();
  if found.is_empty() {
    eprintln!("no models downloaded; try: wavo pull {}", models::DEFAULT);
  } else {
    println!("{:<18} {:>7}  PATH", "NAME", "SIZE");
  }
  for (m, path) in found {
    println!("{:<18} {}  {}", m.name, models::gb(std::fs::metadata(&path)?.len()), path.display());
  }
  Ok(())
}

fn rm(arg: &str) -> Result<()> {
  let m = named(arg)?;
  let freed = hfs::Cache::new().remove(&m.repo(), &m.file());
  match freed.with_context(|| format!("removing {}", m.name))? {
    Some(freed) => eprintln!("removed {} ({} freed)", m.name, models::gb(freed)),
    None => bail!("{} is not downloaded", m.name),
  }
  Ok(())
}

/// Load time, first call and warm median/min over `n` calls, in milliseconds.
fn bench(model: &str, audio: &str, n: usize) -> Result<()> {
  if n == 0 {
    bail!("-n must be at least 1");
  }
  let pcm = audio::read(audio)?;
  let path = model_path(model)?;
  let ms = |t: Instant| t.elapsed().as_secs_f64() * 1e3;
  let t = Instant::now();
  let model = load(&path)?;
  let load = ms(t);
  let t = Instant::now();
  model.transcribe(&pcm)?;
  let first = ms(t);
  let mut warm = (0..n)
    .map(|_| {
      let t = Instant::now();
      model.transcribe(&pcm).map(|_| ms(t))
    })
    .collect::<Result<Vec<_>, _>>()?;
  warm.sort_by(f64::total_cmp);
  let median = (warm[(n - 1) / 2] + warm[n / 2]) / 2.0;
  println!(
    "load {load:.1} ms, first {first:.1} ms, warm median {median:.1} ms, min {:.1} ms",
    warm[0]
  );
  Ok(())
}