feat: add support for speech to text (#68)

Reviewed-on: #68
This commit was merged in pull request #68.
This commit is contained in:
2026-10-27 06:25:19 +00:00
parent b0121f8df2
commit 77b2a5f3cd
28 changed files with 2238 additions and 230 deletions
+30
View File
@@ -0,0 +1,30 @@
[package]
name = "speech"
version.workspace = true
edition.workspace = true
publish.workspace = true
[dependencies]
anyhow.workspace = true
log.workspace = true
gpui.workspace = true
gpui-kit = { workspace = true, features = ["speech"] }
flume.workspace = true
instant.workspace = true
[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
common = { path = "../common" }
parakeet-rs = "=0.3.8"
ort = "=2.0.0-rc.13"
reqwest = { version = "0.12", default-features = false, features = [
"blocking",
"rustls-tls",
] }
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
[target.'cfg(target_arch = "wasm32")'.dependencies]
js-sys = "0.3"
wasm-bindgen = "0.2"
web-sys = { version = "0.3", features = ["Window"] }
+53
View File
@@ -0,0 +1,53 @@
#[cfg(not(target_arch = "wasm32"))]
pub mod model;
#[cfg(not(target_arch = "wasm32"))]
mod parakeet;
#[cfg(target_arch = "wasm32")]
mod web;
use std::time::Duration;
use gpui::{
Animation, AnimationExt as _, AnyElement, App, IntoElement as _, ParentElement as _,
Styled as _, div,
};
use gpui_kit::component::ActiveTheme as _;
use gpui_kit::component::speech::SpeechRecognizer;
/// A ticking `M:SS` counter for the dictation in progress.
pub fn elapsed(started: instant::Instant, cx: &App) -> AnyElement {
div()
.text_xs()
.text_color(cx.theme().muted_foreground)
.with_animation(
"dictation-elapsed",
Animation::new(Duration::from_secs(1))
.repeat()
.with_max_fps(1.),
move |counter, _| counter.child(format_elapsed(started)),
)
.into_any_element()
}
fn format_elapsed(started: instant::Instant) -> String {
let seconds = started.elapsed().as_secs();
format!("{}:{:02}", seconds / 60, seconds % 60)
}
/// The recognizer backing the composer's microphone button.
#[cfg(not(target_arch = "wasm32"))]
pub fn recognizer() -> impl SpeechRecognizer {
parakeet::ParakeetRecognizer
}
/// The recognizer backing the composer's microphone button.
#[cfg(target_arch = "wasm32")]
pub fn recognizer() -> impl SpeechRecognizer {
web::WebSpeechRecognizer
}
/// The audio source backing the composer's microphone button on the web.
#[cfg(target_arch = "wasm32")]
pub fn input() -> impl gpui_kit::component::speech::AudioInput {
web::SelfCapturedInput
}
+23
View File
@@ -0,0 +1,23 @@
{
"repository": "istupakov/parakeet-tdt-0.6b-v3-onnx",
"revision": "8f23f0c03c8761650bdb5b40aaf3e40d2c15f1ce",
"source": "nvidia/parakeet-tdt-0.6b-v3",
"license": "CC-BY-4.0",
"files": [
{
"name": "decoder_joint-model.int8.onnx",
"size": 18202004,
"sha256": "eea7483ee3d1a30375daedc8ed83e3960c91b098812127a0d99d1c8977667a70"
},
{
"name": "encoder-model.int8.onnx",
"size": 652183999,
"sha256": "6139d2fa7e1b086097b277c7149725edbab89cc7c7ae64b23c741be4055aff09"
},
{
"name": "vocab.txt",
"size": 93939,
"sha256": "d58544679ea4bc6ac563d1f545eb7d474bd6cfa467f0a6e2c1dc1c7d37e3c35d"
}
]
}
+295
View File
@@ -0,0 +1,295 @@
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::{fs, io};
use anyhow::{Result, bail};
use serde::Deserialize;
/// How the model is currently available.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Status {
NotInstalled,
Downloading { downloaded: u64, total: u64 },
Installed,
}
impl Status {
pub fn is_downloading(self) -> bool {
matches!(self, Status::Downloading { .. })
}
}
#[derive(Debug, Deserialize)]
struct Manifest {
repository: String,
revision: String,
files: Vec<Artifact>,
}
#[derive(Debug, Deserialize)]
struct Artifact {
name: String,
size: u64,
sha256: String,
}
fn manifest() -> &'static Manifest {
static MANIFEST: OnceLock<Manifest> = OnceLock::new();
MANIFEST.get_or_init(|| {
serde_json::from_str(include_str!("model.json")).expect("pinned model manifest")
})
}
/// The directory the model artifacts are downloaded into.
pub fn dir() -> PathBuf {
common::config_dir()
.join("speech")
.join("parakeet-tdt-0.6b-v3")
}
/// Total download size in bytes.
pub fn total_size() -> u64 {
manifest().files.iter().map(|file| file.size).sum()
}
/// The revision of the installed model, if its artifacts are complete.
fn installed_revision() -> Option<String> {
let manifest = manifest();
let dir = dir();
let complete = manifest.files.iter().all(|file| {
fs::metadata(dir.join(&file.name)).is_ok_and(|metadata| metadata.len() == file.size)
});
if complete
&& fs::read_to_string(dir.join("verified"))
.is_ok_and(|revision| revision == manifest.revision)
{
Some(manifest.revision.clone())
} else {
None
}
}
/// Whether the model is ready to load.
pub fn installed() -> bool {
status() == Status::Installed
}
/// The current status, reconciling the download state with the disk.
pub fn status() -> Status {
let state = state().lock().unwrap_or_else(|error| error.into_inner());
if let Download::Running {
downloaded, total, ..
} = state.download
{
return Status::Downloading { downloaded, total };
}
if installed_revision().is_some() {
Status::Installed
} else {
Status::NotInstalled
}
}
/// Take the last download failure, if it has not been reported yet.
pub fn take_error() -> Option<String> {
let mut state = state().lock().unwrap_or_else(|error| error.into_inner());
state.error.take()
}
/// Start downloading the model, unless a download is already running.
pub fn start_download() {
let mut state = state().lock().unwrap_or_else(|error| error.into_inner());
if matches!(state.download, Download::Running { .. }) {
return;
}
let cancel = Arc::new(AtomicBool::new(false));
state.error = None;
state.download = Download::Running {
downloaded: 0,
total: total_size(),
cancel: cancel.clone(),
};
if let Err(error) = std::thread::Builder::new()
.name("speech-model-download".into())
.spawn(move || download(cancel))
{
state.download = Download::Idle;
state.error = Some(error.to_string());
}
}
/// Request cancellation of the running download.
pub fn cancel_download() {
let state = state().lock().unwrap_or_else(|error| error.into_inner());
if let Download::Running { cancel, .. } = &state.download {
cancel.store(true, Ordering::Release);
}
}
/// Delete the downloaded model and any partial download.
pub fn delete() {
super::parakeet::unload();
cancel_download();
if let Err(error) = fs::remove_dir_all(dir())
&& error.kind() != io::ErrorKind::NotFound
{
log::warn!("failed to remove the speech model: {error}");
}
}
fn download(cancel: Arc<AtomicBool>) {
let result = download_inner(&cancel);
let mut state = state().lock().unwrap_or_else(|error| error.into_inner());
state.download = Download::Idle;
// A cancelled download must not report a spurious failure.
if !cancelled(&cancel)
&& let Err(error) = result
{
state.error = Some(format!("{error:#}"));
}
}
fn download_inner(cancel: &AtomicBool) -> Result<()> {
let manifest = manifest();
let dir = dir();
let client = reqwest::blocking::Client::builder()
.connect_timeout(std::time::Duration::from_secs(20))
.build()?;
fs::create_dir_all(&dir)?;
for file in &manifest.files {
if cancelled(cancel) {
bail!("Cancelled");
}
let staging = dir.join(format!("{}.part", file.name));
let result = download_file(&client, manifest, file, &staging, cancel);
if result.is_err() || cancelled(cancel) {
let _unused = fs::remove_file(&staging);
result?;
bail!("Cancelled");
}
}
fs::write(dir.join("verified"), &manifest.revision)?;
Ok(())
}
fn download_file(
client: &reqwest::blocking::Client,
manifest: &Manifest,
file: &Artifact,
staging: &Path,
cancel: &AtomicBool,
) -> Result<()> {
use std::io::{Read as _, Write as _};
use sha2::Digest as _;
let url = format!(
"https://huggingface.co/{}/resolve/{}/{}",
manifest.repository, manifest.revision, file.name
);
let mut response = client.get(url).send()?.error_for_status()?;
let mut out = fs::File::create(staging)?;
let mut hash = sha2::Sha256::new();
let mut size = 0u64;
loop {
if cancelled(cancel) {
bail!("Cancelled");
}
let mut chunk = [0u8; 64 * 1024];
let read = response.read(&mut chunk)?;
if read == 0 {
break;
}
size += read as u64;
if size > file.size {
bail!("Unexpected model size");
}
hash.update(&chunk[..read]);
out.write_all(&chunk[..read])?;
let mut state = state().lock().unwrap_or_else(|error| error.into_inner());
if let Download::Running { downloaded, .. } = &mut state.download {
*downloaded += read as u64;
}
}
if size != file.size || format!("{:x}", hash.finalize()) != file.sha256 {
bail!("Model checksum verification failed");
}
out.sync_all()?;
fs::rename(staging, dir().join(&file.name))?;
Ok(())
}
fn cancelled(cancel: &AtomicBool) -> bool {
cancel.load(Ordering::Acquire)
}
enum Download {
Idle,
Running {
downloaded: u64,
total: u64,
cancel: Arc<AtomicBool>,
},
}
struct State {
download: Download,
error: Option<String>,
}
fn state() -> &'static Mutex<State> {
static STATE: Mutex<State> = Mutex::new(State {
download: Download::Idle,
error: None,
});
&STATE
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn manifest_pins_v3_artifacts_with_exact_sizes() {
let manifest = manifest();
assert_eq!(manifest.repository, "istupakov/parakeet-tdt-0.6b-v3-onnx");
assert_eq!(manifest.revision.len(), 40);
assert_eq!(manifest.files.len(), 3);
assert_eq!(total_size(), 670_479_942);
assert!(
manifest
.files
.iter()
.all(|file| file.sha256.len() == 64 && !file.name.contains('/'))
);
}
}
+214
View File
@@ -0,0 +1,214 @@
use std::path::PathBuf;
use std::sync::OnceLock;
use std::time::Duration;
use anyhow::{Context as _, Result, anyhow};
use gpui::{App, Task};
use gpui_kit::component::speech::{
AudioFormat, RecognitionSession, SpeechError, SpeechRecognizer, SpeechSink,
};
use parakeet_rs::{ParakeetTDT, Transcriber as _};
const MODEL_RATE: u32 = 16_000;
const MAX_SECONDS: usize = 120;
/// The on-device recognizer.
pub struct ParakeetRecognizer;
impl SpeechRecognizer for ParakeetRecognizer {
fn audio_format(&self) -> AudioFormat {
AudioFormat::new(MODEL_RATE, 1)
}
fn is_available(&self, _cx: &App) -> bool {
true
}
fn start(
&self,
sink: SpeechSink,
cx: &mut App,
) -> Result<Box<dyn RecognitionSession>, SpeechError> {
if super::model::installed() {
worker().send_or_log(Command::Warm(super::model::dir()));
} else if super::model::status().is_downloading() {
return Err(SpeechError::recognizer(anyhow!(
"The speech model is still downloading. Try again when it finishes."
)));
} else {
return Err(SpeechError::recognizer(anyhow!(
"No speech model is installed. Download it from Settings."
)));
}
// The session buffers everything it is pushed,
// so audio is accepted right away.
sink.ready(cx);
Ok(Box::new(ParakeetSession {
sink,
samples: Vec::new(),
_result: None,
}))
}
}
struct ParakeetSession {
sink: SpeechSink,
samples: Vec<f32>,
_result: Option<Task<()>>,
}
impl RecognitionSession for ParakeetSession {
fn push_audio(&mut self, samples: &[i16], _cx: &mut App) {
if self.samples.len() < MAX_SECONDS * MODEL_RATE as usize {
self.samples
.extend(samples.iter().map(|sample| f32::from(*sample) / 32_768.0));
}
}
fn finish(&mut self, cx: &mut App) {
let samples = std::mem::take(&mut self.samples);
let sink = self.sink.clone();
// Speech too short to transcribe ends the session at once.
let silent = samples.iter().all(|sample| sample.abs() < 0.0001);
if samples.len() < MODEL_RATE as usize / 5 || silent {
sink.finish(cx);
return;
}
let (sender, receiver) = flume::unbounded();
worker().send_or_log(Command::Transcribe { samples, sender });
self._result = Some(cx.spawn(async move |cx| {
let result = receiver.recv_async().await;
cx.update(|cx| match result {
Ok(Ok(text)) if !text.is_empty() => {
sink.phrase(text, cx);
sink.finish(cx);
}
Ok(Ok(_)) => sink.finish(cx),
Ok(Err(error)) => sink.error(SpeechError::recognizer(error), cx),
Err(_) => sink.error(
SpeechError::recognizer(anyhow!("The speech worker is unavailable")),
cx,
),
});
}));
}
}
enum Command {
Warm(PathBuf),
Transcribe {
samples: Vec<f32>,
sender: flume::Sender<Result<String, anyhow::Error>>,
},
Unload,
}
/// A handle to the single inference worker.
fn worker() -> &'static flume::Sender<Command> {
static WORKER: OnceLock<flume::Sender<Command>> = OnceLock::new();
WORKER.get_or_init(|| {
let (sender, receiver) = flume::unbounded::<Command>();
let spawned = std::thread::Builder::new()
.name("speech-parakeet".into())
.spawn(move || run_worker(receiver));
if let Err(error) = spawned {
log::error!("failed to spawn the speech worker: {error}");
}
sender
})
}
fn run_worker(receiver: flume::Receiver<Command>) {
let mut cached: Option<(PathBuf, ParakeetTDT)> = None;
while let Ok(command) = receiver.recv_timeout(Duration::from_secs(30)) {
match command {
Command::Unload => {
cached = None;
}
Command::Warm(dir) => {
let loaded = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
load(&mut cached, &dir).map(|_| ())
}));
match flatten_panic(loaded) {
Ok(Ok(())) => {}
Ok(Err(error)) | Err(error) => {
log::warn!("failed to load the speech model: {error:#}")
}
}
}
Command::Transcribe { samples, sender } => {
let transcribed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
transcribe(&mut cached, samples)
}));
let result = match flatten_panic(transcribed) {
Ok(result) => result,
Err(_) => {
cached = None;
Err(anyhow!("Dictation stopped unexpectedly. Try again."))
}
};
if sender.send(result).is_err() {
log::debug!("speech session was cancelled before its result arrived");
}
}
}
}
}
fn load<'a>(
cached: &'a mut Option<(PathBuf, ParakeetTDT)>,
dir: &PathBuf,
) -> Result<&'a mut ParakeetTDT> {
if cached.as_ref().is_none_or(|(path, _)| path != dir) {
let model = ParakeetTDT::from_pretrained(dir, None)
.map_err(|error| anyhow!("Could not load the speech model: {error}"))?;
*cached = Some((dir.clone(), model));
}
let (_, model) = cached.as_mut().context("speech model was unloaded")?;
Ok(model)
}
fn transcribe(cached: &mut Option<(PathBuf, ParakeetTDT)>, samples: Vec<f32>) -> Result<String> {
let model = load(cached, &super::model::dir())?;
model
.transcribe_samples(samples, MODEL_RATE, 1, None)
.map(|result| result.text)
.map_err(|error| anyhow!("Could not transcribe this recording: {error}"))
}
trait SendOrLog {
fn send_or_log(&self, command: Command);
}
impl SendOrLog for flume::Sender<Command> {
fn send_or_log(&self, command: Command) {
if self.send(command).is_err() {
log::debug!("speech worker is gone; command dropped");
}
}
}
/// Retire cached weights, e.g. after the model was deleted.
pub(crate) fn unload() {
worker().send_or_log(Command::Unload);
}
fn flatten_panic<T>(result: std::thread::Result<T>) -> Result<T> {
result.map_err(|_| anyhow!("The speech worker panicked"))
}
+256
View File
@@ -0,0 +1,256 @@
use std::cell::RefCell;
use anyhow::anyhow;
use gpui::App;
use gpui_kit::Subscription;
use gpui_kit::component::speech::{
AudioFormat, AudioInput, AudioSink, RecognitionSession, SpeechError, SpeechRecognizer,
SpeechSink,
};
use js_sys::{Array, Function, JsString, Reflect};
use wasm_bindgen::closure::Closure;
use wasm_bindgen::{JsCast as _, JsValue};
use web_sys::window;
/// The recognizer backed by `SpeechRecognition` (or its `webkit` prefix).
pub struct WebSpeechRecognizer;
impl SpeechRecognizer for WebSpeechRecognizer {
fn audio_format(&self) -> AudioFormat {
AudioFormat::default()
}
fn is_available(&self, _cx: &App) -> bool {
true
}
fn start(
&self,
sink: SpeechSink,
cx: &mut App,
) -> Result<Box<dyn RecognitionSession>, SpeechError> {
let recognition = Recognition::open(sink, cx).map_err(SpeechError::recognizer)?;
Ok(Box::new(recognition))
}
}
pub struct SelfCapturedInput;
impl AudioInput for SelfCapturedInput {
fn start(
&self,
_format: AudioFormat,
_sink: AudioSink,
_cx: &mut App,
) -> Result<Subscription, SpeechError> {
Ok(Subscription::new(|| {}))
}
}
enum Event {
Hypothesis(String),
Phrase(String),
Error(String),
End,
}
struct Recognition {
recognition: JsValue,
on_result: Option<Closure<dyn FnMut(JsValue)>>,
on_error: Option<Closure<dyn FnMut(JsValue)>>,
on_end: Option<Closure<dyn FnMut(JsValue)>>,
_task: gpui::Task<()>,
}
impl Recognition {
fn open(sink: SpeechSink, cx: &mut App) -> anyhow::Result<Self> {
let window = window().ok_or_else(|| anyhow!("No browser window"))?;
let constructor = Reflect::get(&window, &JsString::from("SpeechRecognition"))
.ok()
.filter(|value| value.is_function())
.or_else(|| {
Reflect::get(&window, &JsString::from("webkitSpeechRecognition"))
.ok()
.filter(|value| value.is_function())
})
.ok_or_else(|| anyhow!("This browser does not support speech recognition"))?;
let recognition =
Reflect::construct(constructor.unchecked_ref::<Function>(), &Array::new())
.map_err(|error| js_error("constructing the recognizer", error))?;
let (sender, receiver) = flume::unbounded::<Event>();
let error_sender = sender.clone();
let end_sender = sender.clone();
let finalized = RefCell::new(0usize);
let on_result = Closure::wrap(Box::new(move |event: JsValue| {
let read = |value: &JsValue, property: &str| -> Option<JsValue> {
Reflect::get(value, &JsString::from(property)).ok()
};
let Some(results) = read(&event, "results") else {
return;
};
let length = read(&results, "length")
.and_then(|length| length.as_f64())
.unwrap_or(0.0) as usize;
let mut finalized = finalized.borrow_mut();
let mut interim = String::new();
for index in *finalized..length {
let Some(result) = read(&results, &index.to_string()) else {
continue;
};
let Some(alternative) = read(&result, "0") else {
continue;
};
let Some(text) = read(&alternative, "transcript").and_then(|text| text.as_string())
else {
continue;
};
let is_final = read(&result, "isFinal")
.and_then(|final_value| final_value.as_bool())
.unwrap_or(false);
if is_final {
*finalized += 1;
if sender.send(Event::Phrase(text)).is_err() {
log::debug!("speech session was cancelled mid-utterance");
}
} else {
interim.push_str(&text);
}
}
if !interim.is_empty() && sender.send(Event::Hypothesis(interim)).is_err() {
log::debug!("speech session was cancelled mid-utterance");
}
}) as Box<dyn FnMut(JsValue)>);
let on_error = Closure::wrap(Box::new(move |event: JsValue| {
let message = Reflect::get(&event, &JsString::from("error"))
.ok()
.and_then(|error| error.as_string())
.unwrap_or_else(|| "unknown speech error".into());
if error_sender.send(Event::Error(message)).is_err() {
log::debug!("speech session was cancelled before its error arrived");
}
}) as Box<dyn FnMut(JsValue)>);
let on_end = Closure::wrap(Box::new(move |_event: JsValue| {
if end_sender.send(Event::End).is_err() {
log::debug!("speech session was cancelled before its end arrived");
}
}) as Box<dyn FnMut(JsValue)>);
set_callback(&recognition, "onresult", on_result.as_ref().unchecked_ref())?;
set_callback(&recognition, "onerror", on_error.as_ref().unchecked_ref())?;
set_callback(&recognition, "onend", on_end.as_ref().unchecked_ref())?;
set_bool(&recognition, "continuous", true)?;
set_bool(&recognition, "interimResults", true)?;
call_method(&recognition, "start")?;
// The browser consumes audio as soon as it has permission
sink.ready(cx);
let task = cx.spawn(async move |cx| {
while let Ok(event) = receiver.recv_async().await {
let done = matches!(event, Event::End);
// The speech state ignores sink calls once its session ended,
// so a late event after a cancel is harmless.
cx.update(|cx| match event {
Event::Hypothesis(text) => sink.hypothesis(text, cx),
Event::Phrase(text) => sink.phrase(text, cx),
Event::Error(message) => {
sink.error(SpeechError::recognizer(anyhow!("{message}")), cx)
}
Event::End => sink.finish(cx),
});
if done {
break;
}
}
});
Ok(Self {
recognition,
on_result: Some(on_result),
on_error: Some(on_error),
on_end: Some(on_end),
_task: task,
})
}
fn stop(&mut self) {
self.on_result = None;
self.on_error = None;
self.on_end = None;
let _ignored = call_method(&self.recognition, "abort");
}
}
impl Drop for Recognition {
fn drop(&mut self) {
self.stop();
}
}
impl RecognitionSession for Recognition {
fn push_audio(&mut self, _samples: &[i16], _cx: &mut App) {
// The browser recognizer captures the microphone itself.
// the state still pushes audio through its input, which is a no-op here.
}
fn finish(&mut self, _cx: &mut App) {
// The end handler reports completion, so no result is fabricated here.
let _ignored = call_method(&self.recognition, "stop");
}
}
fn set_callback(target: &JsValue, property: &str, callback: &Function) -> anyhow::Result<()> {
Reflect::set(target, &JsString::from(property), callback)
.map_err(|error| js_error(property, error))?;
Ok(())
}
fn set_bool(target: &JsValue, property: &str, value: bool) -> anyhow::Result<()> {
Reflect::set(target, &JsString::from(property), &JsValue::from(value))
.map_err(|error| js_error(property, error))?;
Ok(())
}
fn call_method(target: &JsValue, method: &str) -> anyhow::Result<()> {
let function = Reflect::get(target, &JsString::from(method))
.map_err(|error| js_error(method, error))?
.unchecked_into::<Function>();
function
.call0(target)
.map(|_| ())
.map_err(|error| js_error(method, error))
}
fn js_error(what: &str, error: JsValue) -> anyhow::Error {
let description = error
.as_string()
.or_else(|| {
Reflect::get(&error, &JsString::from("message"))
.ok()
.and_then(|message| message.as_string())
})
.unwrap_or_else(|| "unknown error".into());
anyhow!("{what} failed: {description}")
}