swarm-stt: the speech-to-text core (hash cache, shared runs, ffmpeg and whisper-cli engine)

swarm-stt/src/engine.rs

@@ -0,0 +1,172 @@
+//! The engine that turns audio bytes into text.
+
+use std::path::PathBuf;
+use std::process::Stdio;
+
+use async_trait::async_trait;
+use tokio::io::AsyncWriteExt as _;
+use tokio::process::Command;
+
+/// Turns the bytes of an audio file into its transcript.
+#[async_trait]
+pub trait Engine: Send + Sync {
+    /// The text spoken in `audio`, or why it could not be told.
+    async fn transcribe(&self, audio: &[u8]) -> Result<String, String>;
+}
+
+/// `ffmpeg` decodes any container to 16 kHz mono WAV, then `whisper-cli` reads
+/// it with the configured model.
+#[derive(Debug, Clone)]
+pub struct Whisper {
+    /// The `ffmpeg` binary.
+    pub ffmpeg: PathBuf,
+    /// The `whisper-cli` binary (whisper.cpp).
+    pub whisper: PathBuf,
+    /// The ggml model file, such as `ggml-small.bin`.
+    pub model: PathBuf,
+    /// Where scratch files go; the caller creates it.
+    pub scratch: PathBuf,
+    /// Seconds of audio to decode; the rest is cut.
+    pub max_secs: u64,
+}
+
+#[async_trait]
+impl Engine for Whisper {
+    async fn transcribe(&self, audio: &[u8]) -> Result<String, String> {
+        let dir = tempfile_in(&self.scratch)
+            .await
+            .map_err(|e| format!("making a scratch directory: {e}"))?;
+        let result = self.run(&dir, audio).await;
+        let _ = tokio::fs::remove_dir_all(&dir).await;
+        result
+    }
+}
+
+impl Whisper {
+    async fn run(&self, dir: &std::path::Path, audio: &[u8]) -> Result<String, String> {
+        let wav = dir.join("audio.wav");
+        let mut decode = Command::new(&self.ffmpeg)
+            .args(["-nostdin", "-v", "error", "-i", "pipe:0", "-t"])
+            .arg(self.max_secs.to_string())
+            .args(["-ar", "16000", "-ac", "1", "-y"])
+            .arg(&wav)
+            .stdin(Stdio::piped())
+            .stdout(Stdio::null())
+            .stderr(Stdio::piped())
+            .kill_on_drop(true)
+            .spawn()
+            .map_err(|e| format!("starting ffmpeg: {e}"))?;
+        if let Some(mut stdin) = decode.stdin.take() {
+            // ffmpeg may stop reading early (a cut at `-t`): a broken pipe is fine.
+            let _ = stdin.write_all(audio).await;
+        }
+        let decoded = decode
+            .wait_with_output()
+            .await
+            .map_err(|e| format!("waiting for ffmpeg: {e}"))?;
+        if !decoded.status.success() {
+            return Err(format!(
+                "ffmpeg could not decode the audio: {}",
+                String::from_utf8_lossy(&decoded.stderr).trim()
+            ));
+        }
+        let heard = Command::new(&self.whisper)
+            .arg("-m")
+            .arg(&self.model)
+            .arg("-f")
+            .arg(&wav)
+            .args(["-l", "auto", "-nt", "-np"])
+            .stdin(Stdio::null())
+            .kill_on_drop(true)
+            .output()
+            .await
+            .map_err(|e| format!("starting whisper-cli: {e}"))?;
+        if !heard.status.success() {
+            return Err(format!(
+                "whisper-cli failed: {}",
+                String::from_utf8_lossy(&heard.stderr).trim()
+            ));
+        }
+        Ok(String::from_utf8_lossy(&heard.stdout).trim().to_owned())
+    }
+}
+
+/// A fresh empty directory under `parent`, named for this process and a counter.
+async fn tempfile_in(parent: &std::path::Path) -> std::io::Result<PathBuf> {
+    use std::sync::atomic::{AtomicU64, Ordering};
+    static NEXT: AtomicU64 = AtomicU64::new(0);
+    let dir = parent.join(format!(
+        "run-{}-{}",
+        std::process::id(),
+        NEXT.fetch_add(1, Ordering::Relaxed)
+    ));
+    tokio::fs::create_dir_all(&dir).await?;
+    Ok(dir)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{Engine as _, Whisper};
+    use std::os::unix::fs::PermissionsExt as _;
+    use std::path::Path;
+
+    fn script(path: &Path, body: &str) {
+        std::fs::write(path, format!("#!/bin/sh\n{body}\n")).unwrap();
+        std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o755)).unwrap();
+    }
+
+    fn whisper(dir: &Path, ffmpeg: &str, whisper: &str) -> Whisper {
+        script(&dir.join("ffmpeg"), ffmpeg);
+        script(&dir.join("whisper-cli"), whisper);
+        Whisper {
+            ffmpeg: dir.join("ffmpeg"),
+            whisper: dir.join("whisper-cli"),
+            model: dir.join("model.bin"),
+            scratch: dir.join("scratch"),
+            max_secs: 300,
+        }
+    }
+
+    // The fake ffmpeg writes the last argument (the wav path), the fake
+    // whisper-cli prints what it finds in the file named by `-f`.
+    const FFMPEG_OK: &str = r#"for last; do :; done; cat > "$last""#;
+    const WHISPER_ECHO: &str =
+        r#"while [ $# -gt 0 ]; do [ "$1" = -f ] && f=$2; shift; done; cat "$f""#;
+
+    #[tokio::test]
+    async fn audio_goes_through_ffmpeg_then_whisper_and_the_text_is_trimmed() {
+        let dir = tempfile::tempdir().unwrap();
+        let engine = whisper(dir.path(), FFMPEG_OK, WHISPER_ECHO);
+        let text = engine.transcribe(b"  hello there \n").await.unwrap();
+        assert_eq!(text, "hello there");
+        assert_eq!(
+            std::fs::read_dir(dir.path().join("scratch"))
+                .unwrap()
+                .count(),
+            0,
+            "the scratch directory is removed"
+        );
+    }
+
+    #[tokio::test]
+    async fn an_ffmpeg_failure_is_reported_with_its_stderr() {
+        let dir = tempfile::tempdir().unwrap();
+        let engine = whisper(dir.path(), "echo 'bad data' >&2; exit 1", WHISPER_ECHO);
+        let err = engine.transcribe(b"x").await.unwrap_err();
+        assert!(
+            err.contains("could not decode") && err.contains("bad data"),
+            "{err}"
+        );
+    }
+
+    #[tokio::test]
+    async fn a_whisper_failure_is_reported_with_its_stderr() {
+        let dir = tempfile::tempdir().unwrap();
+        let engine = whisper(dir.path(), FFMPEG_OK, "echo 'no model' >&2; exit 2");
+        let err = engine.transcribe(b"x").await.unwrap_err();
+        assert!(
+            err.contains("whisper-cli failed") && err.contains("no model"),
+            "{err}"
+        );
+    }
+}