tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Careless whisper (#6877)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Douwe Osinga
2026-02-03 12:54:29 +01:00
committed by GitHub
parent 3d0bb3c670
commit 1373d9c5f9
25 changed files with 118541 additions and 485 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/Cargo.toml
+2
View File
@@ -45,6 +45,8 @@ socket2 = "0.6.1"
fs2 = "0.4.3"
rustls = { version = "0.23", features = ["ring"] }
uuid = { version = "1.19.0", features = ["v4"] }
once_cell = "1.20.2"
dirs = "5.0"
[target.'cfg(windows)'.dependencies]
winreg = { version = "0.55.0" }
crates/goose-server/src/openapi.rs
+10 -1
View File
@@ -4,6 +4,7 @@ use goose::agents::ExtensionConfig;
use goose::config::permission::PermissionLevel;
use goose::config::ExtensionEntry;
use goose::conversation::Conversation;
use goose::dictation::download_manager::{DownloadProgress, DownloadStatus};
use goose::model::ModelConfig;
use goose::permission::permission_confirmation::PrincipalType;
use goose::providers::base::{ConfigKey, ModelInfo, ProviderMetadata, ProviderType};
@@ -414,6 +415,11 @@ derive_utoipa!(Icon as IconSchema);
super::routes::telemetry::send_telemetry_event,
super::routes::dictation::transcribe_dictation,
super::routes::dictation::get_dictation_config,
super::routes::dictation::list_models,
super::routes::dictation::download_model,
super::routes::dictation::get_download_progress,
super::routes::dictation::cancel_download,
super::routes::dictation::delete_model,
),
components(schemas(
super::routes::config_management::UpsertConfigQuery,
@@ -574,8 +580,11 @@ derive_utoipa!(Icon as IconSchema);
goose::goose_apps::ResourceMetadata,
super::routes::dictation::TranscribeRequest,
super::routes::dictation::TranscribeResponse,
super::routes::dictation::DictationProvider,
goose::dictation::providers::DictationProvider,
super::routes::dictation::DictationProviderStatus,
super::routes::dictation::WhisperModelResponse,
DownloadProgress,
DownloadStatus,
))
)]
pub struct ApiDoc;
crates/goose-server/src/routes/dictation.rs
+202 -240
View File
@@ -1,76 +1,31 @@
use crate::routes::errors::ErrorResponse;
use crate::state::AppState;
use axum::{
extract::{DefaultBodyLimit, Path},
http::StatusCode,
routing::{get, post},
routing::{delete, get, post},
Json, Router,
};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use goose::providers::api_client::{ApiClient, AuthMethod};
use goose::dictation::download_manager::{get_download_manager, DownloadProgress};
use goose::dictation::providers::{
is_configured, transcribe_local, transcribe_with_provider, DictationProvider, PROVIDERS,
};
use goose::dictation::whisper;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use utoipa::ToSchema;
const MAX_AUDIO_SIZE_BYTES: usize = 25 * 1024 * 1024;
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_AUDIO_SIZE_BYTES: usize = 50 * 1024 * 1024;
// DictationProvider definitions
struct DictationProviderDef {
config_key: &'static str,
default_url: &'static str,
host_key: Option<&'static str>,
description: &'static str,
uses_provider_config: bool,
settings_path: Option<&'static str>,
}
const PROVIDERS: &[(&str, DictationProviderDef)] = &[
(
"openai",
DictationProviderDef {
config_key: "OPENAI_API_KEY",
default_url: "https://api.openai.com/v1/audio/transcriptions",
host_key: Some("OPENAI_HOST"),
description: "Uses OpenAI Whisper API for high-quality transcription.",
uses_provider_config: true,
settings_path: Some("Settings > Models"),
},
),
(
"elevenlabs",
DictationProviderDef {
config_key: "ELEVENLABS_API_KEY",
default_url: "https://api.elevenlabs.io/v1/speech-to-text",
host_key: None,
description: "Uses ElevenLabs speech-to-text API for advanced voice processing.",
uses_provider_config: false,
settings_path: None,
},
),
];
fn get_provider_def(name: &str) -> Option<&'static DictationProviderDef> {
PROVIDERS
.iter()
.find_map(|(n, def)| if *n == name { Some(def) } else { None })
}
#[derive(Debug, Deserialize, ToSchema)]
#[serde(rename_all = "lowercase")]
pub enum DictationProvider {
OpenAI,
ElevenLabs,
}
impl DictationProvider {
fn as_str(&self) -> &'static str {
match self {
DictationProvider::OpenAI => "openai",
DictationProvider::ElevenLabs => "elevenlabs",
}
}
#[derive(Debug, Serialize, ToSchema)]
pub struct WhisperModelResponse {
#[serde(flatten)]
#[schema(inline)]
model: &'static whisper::WhisperModel,
downloaded: bool,
recommended: bool,
}
#[derive(Debug, Deserialize, ToSchema)]
@@ -113,17 +68,6 @@ fn validate_audio(audio: &str, mime_type: &str) -> Result<(Vec<u8>, &'static str
.decode(audio)
.map_err(|_| ErrorResponse::bad_request("Invalid base64 audio data"))?;
if audio_bytes.len() > MAX_AUDIO_SIZE_BYTES {
return Err(ErrorResponse {
message: format!(
"Audio file too large: {} bytes (max: {} bytes)",
audio_bytes.len(),
MAX_AUDIO_SIZE_BYTES
),
status: StatusCode::PAYLOAD_TOO_LARGE,
});
}
let extension = match mime_type {
"audio/webm" | "audio/webm;codecs=opus" => "webm",
"audio/mp4" => "mp4",
@@ -141,164 +85,37 @@ fn validate_audio(audio: &str, mime_type: &str) -> Result<(Vec<u8>, &'static str
Ok((audio_bytes, extension))
}
async fn handle_response_error(response: reqwest::Response) -> ErrorResponse {
let status = response.status();
let error_text = response.text().await.unwrap_or_default();
fn convert_error(e: anyhow::Error) -> ErrorResponse {
let error_msg = e.to_string();
ErrorResponse {
message: if status == 401
|| error_text.contains("Invalid API key")
|| error_text.contains("Unauthorized")
{
"Invalid API key".to_string()
} else if status == 429 || error_text.contains("quota") || error_text.contains("limit") {
"Rate limit exceeded".to_string()
} else {
format!("API error: {}", error_text)
},
status: if status.is_client_error() {
status
} else {
StatusCode::BAD_GATEWAY
},
}
}
fn build_api_client(provider: &str) -> Result<ApiClient, ErrorResponse> {
let config = goose::config::Config::global();
let def = get_provider_def(provider)
.ok_or_else(|| ErrorResponse::bad_request(format!("Unknown provider: {}", provider)))?;
let api_key = config
.get_secret(def.config_key)
.map_err(|_| ErrorResponse {
message: format!("{} not configured", def.config_key),
status: StatusCode::PRECONDITION_FAILED,
})?;
let url = if let Some(host_key) = def.host_key {
config
.get(host_key, false)
.ok()
.and_then(|v| v.as_str().map(|s| s.to_string()))
.map(|custom_host| {
let path = def
.default_url
.splitn(4, '/')
.nth(3)
.map(|p| format!("/{}", p))
.unwrap_or_default();
format!("{}{}", custom_host.trim_end_matches('/'), path)
})
.unwrap_or_else(|| def.default_url.to_string())
} else {
def.default_url.to_string()
};
let auth = match provider {
"openai" => AuthMethod::BearerToken(api_key),
"elevenlabs" => AuthMethod::ApiKey {
header_name: "xi-api-key".to_string(),
key: api_key,
},
_ => {
return Err(ErrorResponse::bad_request(format!(
"Unknown provider: {}",
provider
)))
if error_msg.contains("Invalid API key") {
ErrorResponse {
message: error_msg,
status: StatusCode::UNAUTHORIZED,
}
};
ApiClient::with_timeout(url, auth, REQUEST_TIMEOUT)
.map_err(|e| ErrorResponse::internal(format!("Failed to create client: {}", e)))
}
async fn transcribe_openai(
audio_bytes: Vec<u8>,
extension: &str,
mime_type: &str,
client: &ApiClient,
) -> Result<String, ErrorResponse> {
let part = reqwest::multipart::Part::bytes(audio_bytes)
.file_name(format!("audio.{}", extension))
.mime_str(mime_type)
.map_err(|e| ErrorResponse::internal(format!("Failed to create multipart: {}", e)))?;
let form = reqwest::multipart::Form::new()
.part("file", part)
.text("model", "whisper-1");
let response = client
.request(None, "")
.multipart_post(form)
.await
.map_err(|e| ErrorResponse {
message: if e.to_string().contains("timeout") {
"Request timed out".to_string()
} else {
format!("Request failed: {}", e)
},
status: if e.to_string().contains("timeout") {
StatusCode::GATEWAY_TIMEOUT
} else {
StatusCode::SERVICE_UNAVAILABLE
},
})?;
if !response.status().is_success() {
return Err(handle_response_error(response).await);
} else if error_msg.contains("Rate limit exceeded") || error_msg.contains("quota") {
ErrorResponse {
message: error_msg,
status: StatusCode::TOO_MANY_REQUESTS,
}
} else if error_msg.contains("not configured") {
ErrorResponse {
message: error_msg,
status: StatusCode::PRECONDITION_FAILED,
}
} else if error_msg.contains("timeout") {
ErrorResponse {
message: error_msg,
status: StatusCode::GATEWAY_TIMEOUT,
}
} else if error_msg.contains("API error") {
ErrorResponse {
message: error_msg,
status: StatusCode::BAD_GATEWAY,
}
} else {
ErrorResponse::internal(error_msg)
}
let data: TranscribeResponse = response
.json()
.await
.map_err(|e| ErrorResponse::internal(format!("Failed to parse response: {}", e)))?;
Ok(data.text)
}
async fn transcribe_elevenlabs(
audio_bytes: Vec<u8>,
extension: &str,
mime_type: &str,
client: &ApiClient,
) -> Result<String, ErrorResponse> {
let part = reqwest::multipart::Part::bytes(audio_bytes)
.file_name(format!("audio.{}", extension))
.mime_str(mime_type)
.map_err(|_| ErrorResponse::internal("Failed to create multipart"))?;
let form = reqwest::multipart::Form::new()
.part("file", part)
.text("model_id", "scribe_v1");
let response = client
.request(None, "")
.multipart_post(form)
.await
.map_err(|e| ErrorResponse {
message: if e.to_string().contains("timeout") {
"Request timed out".to_string()
} else {
format!("Request failed: {}", e)
},
status: if e.to_string().contains("timeout") {
StatusCode::GATEWAY_TIMEOUT
} else {
StatusCode::SERVICE_UNAVAILABLE
},
})?;
if !response.status().is_success() {
return Err(handle_response_error(response).await);
}
let data: TranscribeResponse = response
.json()
.await
.map_err(|e| ErrorResponse::internal(format!("Failed to parse response: {}", e)))?;
Ok(data.text)
}
#[utoipa::path(
@@ -309,11 +126,11 @@ async fn transcribe_elevenlabs(
(status = 200, description = "Audio transcribed successfully", body = TranscribeResponse),
(status = 400, description = "Invalid request (bad base64 or unsupported format)"),
(status = 401, description = "Invalid API key"),
(status = 412, description = "DictationProvider not configured"),
(status = 412, description = "Provider not configured"),
(status = 413, description = "Audio file too large (max 25MB)"),
(status = 429, description = "Rate limit exceeded"),
(status = 500, description = "Internal server error"),
(status = 502, description = "DictationProvider API error"),
(status = 502, description = "Provider API error"),
(status = 503, description = "Service unavailable"),
(status = 504, description = "Request timeout")
)
@@ -322,16 +139,39 @@ pub async fn transcribe_dictation(
Json(request): Json<TranscribeRequest>,
) -> Result<Json<TranscribeResponse>, ErrorResponse> {
let (audio_bytes, extension) = validate_audio(&request.audio, &request.mime_type)?;
let provider_name = request.provider.as_str();
let client = build_api_client(provider_name)?;
let text = match request.provider {
DictationProvider::OpenAI => {
transcribe_openai(audio_bytes, extension, &request.mime_type, &client).await?
}
DictationProvider::ElevenLabs => {
transcribe_elevenlabs(audio_bytes, extension, &request.mime_type, &client).await?
}
DictationProvider::OpenAI => transcribe_with_provider(
DictationProvider::OpenAI,
"model".to_string(),
"whisper-1".to_string(),
audio_bytes,
extension,
&request.mime_type,
)
.await
.map_err(convert_error)?,
DictationProvider::Groq => transcribe_with_provider(
DictationProvider::Groq,
"model".to_string(),
"whisper-large-v3-turbo".to_string(),
audio_bytes,
extension,
&request.mime_type,
)
.await
.map_err(convert_error)?,
DictationProvider::ElevenLabs => transcribe_with_provider(
DictationProvider::ElevenLabs,
"model_id".to_string(),
"scribe_v1".to_string(),
audio_bytes,
extension,
&request.mime_type,
)
.await
.map_err(convert_error)?,
DictationProvider::Local => transcribe_local(audio_bytes).await.map_err(convert_error)?,
};
Ok(Json(TranscribeResponse { text }))
@@ -345,12 +185,13 @@ pub async fn transcribe_dictation(
)
)]
pub async fn get_dictation_config(
) -> Result<Json<HashMap<String, DictationProviderStatus>>, ErrorResponse> {
) -> Result<Json<HashMap<DictationProvider, DictationProviderStatus>>, ErrorResponse> {
let config = goose::config::Config::global();
let mut providers = HashMap::new();
for (name, def) in PROVIDERS.iter() {
let configured = config.get_secret::<String>(def.config_key).is_ok();
for def in PROVIDERS {
let provider = def.provider;
let configured = is_configured(provider);
let host = if let Some(host_key) = def.host_key {
config
@@ -362,7 +203,7 @@ pub async fn get_dictation_config(
};
providers.insert(
name.to_string(),
provider,
DictationProviderStatus {
configured,
host,
@@ -381,9 +222,130 @@ pub async fn get_dictation_config(
Ok(Json(providers))
}
#[utoipa::path(
get,
path = "/dictation/models",
responses(
(status = 200, description = "List of available Whisper models", body = Vec<WhisperModelResponse>)
)
)]
pub async fn list_models() -> Result<Json<Vec<WhisperModelResponse>>, ErrorResponse> {
let recommended_id = whisper::recommend_model();
let models = whisper::available_models()
.iter()
.map(|m| WhisperModelResponse {
model: m,
downloaded: m.is_downloaded(),
recommended: m.id == recommended_id,
})
.collect();
Ok(Json(models))
}
#[utoipa::path(
post,
path = "/dictation/models/{model_id}/download",
responses(
(status = 202, description = "Download started"),
(status = 400, description = "Download already in progress"),
(status = 500, description = "Internal server error")
)
)]
pub async fn download_model(Path(model_id): Path<String>) -> Result<StatusCode, ErrorResponse> {
let model = whisper::get_model(&model_id)
.ok_or_else(|| ErrorResponse::bad_request("Model not found"))?;
let manager = get_download_manager();
manager
.download_model(
model.id.to_string(),
model.url.to_string(),
model.local_path(),
)
.await
.map_err(convert_error)?;
Ok(StatusCode::ACCEPTED)
}
#[utoipa::path(
get,
path = "/dictation/models/{model_id}/download",
responses(
(status = 200, description = "Download progress", body = DownloadProgress),
(status = 404, description = "Download not found")
)
)]
pub async fn get_download_progress(
Path(model_id): Path<String>,
) -> Result<Json<DownloadProgress>, ErrorResponse> {
let manager = get_download_manager();
let progress = manager
.get_progress(&model_id)
.ok_or_else(|| ErrorResponse::bad_request("Download not found"))?;
Ok(Json(progress))
}
#[utoipa::path(
delete,
path = "/dictation/models/{model_id}/download",
responses(
(status = 200, description = "Download cancelled"),
(status = 404, description = "Download not found")
)
)]
pub async fn cancel_download(Path(model_id): Path<String>) -> Result<StatusCode, ErrorResponse> {
let manager = get_download_manager();
manager.cancel_download(&model_id).map_err(convert_error)?;
Ok(StatusCode::OK)
}
#[utoipa::path(
delete,
path = "/dictation/models/{model_id}",
responses(
(status = 200, description = "Model deleted"),
(status = 404, description = "Model not found or not downloaded"),
(status = 500, description = "Failed to delete model")
)
)]
pub async fn delete_model(Path(model_id): Path<String>) -> Result<StatusCode, ErrorResponse> {
let model = whisper::get_model(&model_id)
.ok_or_else(|| ErrorResponse::bad_request("Model not found"))?;
let path = model.local_path();
if !path.exists() {
return Err(ErrorResponse::bad_request("Model not downloaded"));
}
tokio::fs::remove_file(&path)
.await
.map_err(|e| ErrorResponse::internal(format!("Failed to delete model: {}", e)))?;
Ok(StatusCode::OK)
}
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/dictation/transcribe", post(transcribe_dictation))
.route("/dictation/config", get(get_dictation_config))
.route("/dictation/models", get(list_models))
.route(
"/dictation/models/{model_id}/download",
post(download_model),
)
.route(
"/dictation/models/{model_id}/download",
get(get_download_progress),
)
.route(
"/dictation/models/{model_id}/download",
delete(cancel_download),
)
.route("/dictation/models/{model_id}", delete(delete_model))
.layer(DefaultBodyLimit::max(MAX_AUDIO_SIZE_BYTES))
.with_state(state)
}
crates/goose/Cargo.toml
+19
View File
@@ -7,6 +7,9 @@ license.workspace = true
repository.workspace = true
description.workspace = true
[features]
default = []
[lints]
workspace = true
@@ -87,6 +90,17 @@ dashmap = "6.1"
ahash = "0.8"
tokio-util = { version = "0.7.15", features = ["compat"] }
unicode-normalization = "0.1"
goose-mcp = { path = "../goose-mcp" }
# For local Whisper transcription
candle-core = { version = "0.8.4" }
candle-nn = { version = "0.8.4" }
candle-transformers = { version = "0.8.4" }
byteorder = "1.5.0"
tokenizers = "0.21.0"
hf-hub = { version = "0.4.3", default-features = false, features = ["tokio"] }
symphonia = { version = "0.5", features = ["all"] }
rubato = "0.16"
zip = "0.6"
sys-info = "0.9"
@@ -104,6 +118,11 @@ unbinder = "0.1.7"
[target.'cfg(target_os = "windows")'.dependencies]
winapi = { version = "0.3", features = ["wincred"] }
# Platform-specific GPU acceleration for Whisper
[target.'cfg(target_os = "macos")'.dependencies]
candle-core = { version = "0.8.4", features = ["metal"] }
candle-nn = { version = "0.8.4", features = ["metal"] }
[dev-dependencies]
serial_test = { workspace = true }
mockall = "0.13.1"
crates/goose/examples/test_whisper.rs
+31
View File
@@ -0,0 +1,31 @@
use goose::dictation::whisper::{get_model, WhisperTranscriber};
const WHISPER_TOKENIZER_JSON: &str = include_str!("../src/dictation/whisper_data/tokens.json");
fn main() -> anyhow::Result<()> {
// Initialize logging
tracing_subscriber::fmt::init();
let audio_path = "/tmp/whisper_audio_16k.wav";
let model_id = "tiny";
let model =
get_model(model_id).ok_or_else(|| anyhow::anyhow!("Model {} not found", model_id))?;
let model_path = model.local_path();
println!("Loading model from: {}", model_path.display());
let mut transcriber =
WhisperTranscriber::new_with_tokenizer(model_id, &model_path, WHISPER_TOKENIZER_JSON)?;
println!("Reading audio from: {}", audio_path);
let audio_data = std::fs::read(audio_path)?;
println!("Transcribing...");
let text = transcriber.transcribe(&audio_data)?;
println!("\n========== FINAL TRANSCRIPTION ==========");
println!("{}", text);
println!("=========================================\n");
Ok(())
}
crates/goose/src/context_mgmt/mod.rs
+2 -13
View File
@@ -11,8 +11,8 @@ use rmcp::model::Role;
use serde::Serialize;
use std::sync::Arc;
use tokio::task::JoinHandle;
use tracing::info;
use tracing::log::warn;
use tracing::{debug, info};
pub const DEFAULT_COMPACTION_THRESHOLD: f64 = 0.8;
@@ -188,7 +188,7 @@ pub async fn check_if_compaction_needed(
let context_limit = provider.get_model_config().context_limit();
let (current_tokens, token_source) = match session.total_tokens {
let (current_tokens, _token_source) = match session.total_tokens {
Some(tokens) => (tokens as usize, "session metadata"),
None => {
let token_counter = create_token_counter()
@@ -212,17 +212,6 @@ pub async fn check_if_compaction_needed(
} else {
usage_ratio > threshold
};
debug!(
"Compaction check: {} / {} tokens ({:.1}%), threshold: {:.1}%, needs compaction: {}, source: {}",
current_tokens,
context_limit,
usage_ratio * 100.0,
threshold * 100.0,
needs_compaction,
token_source
);
Ok(needs_compaction)
}
crates/goose/src/dictation/download_manager.rs
+251
View File
@@ -0,0 +1,251 @@
use crate::dictation::whisper::LOCAL_WHISPER_MODEL_CONFIG_KEY;
use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use tokio::io::AsyncWriteExt;
use utoipa::ToSchema;
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct DownloadProgress {
/// Model ID being downloaded
pub model_id: String,
/// Download status
pub status: DownloadStatus,
/// Bytes downloaded so far
pub bytes_downloaded: u64,
/// Total bytes to download
pub total_bytes: u64,
/// Download progress percentage (0-100)
pub progress_percent: f32,
/// Download speed in bytes per second
pub speed_bps: Option<u64>,
/// Estimated time remaining in seconds
pub eta_seconds: Option<u64>,
/// Error message if failed
pub error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum DownloadStatus {
Downloading,
Completed,
Failed,
Cancelled,
}
type DownloadMap = Arc<Mutex<HashMap<String, DownloadProgress>>>;
pub struct DownloadManager {
downloads: DownloadMap,
}
impl Default for DownloadManager {
fn default() -> Self {
Self::new()
}
}
impl DownloadManager {
pub fn new() -> Self {
Self {
downloads: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn get_progress(&self, model_id: &str) -> Option<DownloadProgress> {
self.downloads.lock().ok()?.get(model_id).cloned()
}
pub fn cancel_download(&self, model_id: &str) -> Result<()> {
let mut downloads = self
.downloads
.lock()
.map_err(|_| anyhow::anyhow!("Failed to acquire lock"))?;
if let Some(progress) = downloads.get_mut(model_id) {
progress.status = DownloadStatus::Cancelled;
Ok(())
} else {
anyhow::bail!("Download not found")
}
}
pub async fn download_model(
&self,
model_id: String,
url: String,
destination: PathBuf,
) -> Result<()> {
// Initialize progress
{
let mut downloads = self
.downloads
.lock()
.map_err(|_| anyhow::anyhow!("Failed to acquire lock"))?;
if downloads.contains_key(&model_id) {
anyhow::bail!("Download already in progress");
}
downloads.insert(
model_id.clone(),
DownloadProgress {
model_id: model_id.clone(),
status: DownloadStatus::Downloading,
bytes_downloaded: 0,
total_bytes: 0,
progress_percent: 0.0,
speed_bps: None,
eta_seconds: None,
error: None,
},
);
}
// Create parent directory if it doesn't exist
if let Some(parent) = destination.parent() {
tokio::fs::create_dir_all(parent)
.await
.map_err(|e| anyhow::anyhow!("Failed to create directory: {}", e))?;
}
let downloads = self.downloads.clone();
let model_id_clone = model_id.clone();
// Download in background task
tokio::spawn(async move {
match Self::download_file(&url, &destination, &downloads, &model_id_clone).await {
Ok(_) => {
if let Ok(mut downloads) = downloads.lock() {
if let Some(progress) = downloads.get_mut(&model_id_clone) {
progress.status = DownloadStatus::Completed;
progress.progress_percent = 100.0;
}
}
let _ = crate::config::Config::global()
.set_param(LOCAL_WHISPER_MODEL_CONFIG_KEY, model_id_clone.clone());
}
Err(e) => {
if let Ok(mut downloads) = downloads.lock() {
if let Some(progress) = downloads.get_mut(&model_id_clone) {
progress.status = DownloadStatus::Failed;
progress.error = Some(e.to_string());
}
}
}
}
});
Ok(())
}
async fn download_file(
url: &str,
destination: &PathBuf,
downloads: &DownloadMap,
model_id: &str,
) -> Result<(), anyhow::Error> {
let client = reqwest::Client::new();
let mut response = client.get(url).send().await?;
if !response.status().is_success() {
anyhow::bail!("Failed to download: HTTP {}", response.status());
}
let total_bytes = response.content_length().unwrap_or(0);
{
if let Ok(mut downloads) = downloads.lock() {
if let Some(progress) = downloads.get_mut(model_id) {
progress.total_bytes = total_bytes;
}
}
}
let mut file = tokio::fs::File::create(destination).await?;
let mut bytes_downloaded = 0u64;
let start_time = std::time::Instant::now();
while let Some(chunk) = response.chunk().await? {
// Check if cancelled
let should_cancel = {
if let Ok(downloads) = downloads.lock() {
if let Some(progress) = downloads.get(model_id) {
progress.status == DownloadStatus::Cancelled
} else {
false
}
} else {
false
}
};
if should_cancel {
// Clean up partial download
let _ = tokio::fs::remove_file(destination).await;
return Ok(());
}
file.write_all(&chunk).await?;
bytes_downloaded += chunk.len() as u64;
// Update progress
let elapsed = start_time.elapsed().as_secs_f64();
let speed_bps = if elapsed > 0.0 {
Some((bytes_downloaded as f64 / elapsed) as u64)
} else {
None
};
let eta_seconds = if let Some(speed) = speed_bps {
if speed > 0 && total_bytes > 0 {
Some((total_bytes - bytes_downloaded) / speed)
} else {
None
}
} else {
None
};
if let Ok(mut downloads) = downloads.lock() {
if let Some(progress) = downloads.get_mut(model_id) {
progress.bytes_downloaded = bytes_downloaded;
progress.progress_percent = if total_bytes > 0 {
(bytes_downloaded as f64 / total_bytes as f64 * 100.0) as f32
} else {
0.0
};
progress.speed_bps = speed_bps;
progress.eta_seconds = eta_seconds;
}
}
}
file.flush().await?;
Ok(())
}
pub fn clear_completed(&self, model_id: &str) {
if let Ok(mut downloads) = self.downloads.lock() {
if let Some(progress) = downloads.get(model_id) {
if progress.status == DownloadStatus::Completed
|| progress.status == DownloadStatus::Failed
|| progress.status == DownloadStatus::Cancelled
{
downloads.remove(model_id);
}
}
}
}
}
static DOWNLOAD_MANAGER: once_cell::sync::Lazy<DownloadManager> =
once_cell::sync::Lazy::new(DownloadManager::new);
pub fn get_download_manager() -> &'static DownloadManager {
&DOWNLOAD_MANAGER
}
crates/goose/src/dictation/mod.rs
+3
View File
@@ -0,0 +1,3 @@
pub mod download_manager;
pub mod providers;
pub mod whisper;
crates/goose/src/dictation/providers.rs
+238
View File
@@ -0,0 +1,238 @@
use crate::config::Config;
use crate::dictation::whisper::LOCAL_WHISPER_MODEL_CONFIG_KEY;
use crate::providers::api_client::{ApiClient, AuthMethod};
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::sync::Mutex;
use std::time::Duration;
use utoipa::ToSchema;
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
// Global lazy-initialized transcriber to reuse the loaded model
// Stores (model_path, transcriber) to detect when model changes
static LOCAL_TRANSCRIBER: once_cell::sync::Lazy<
Mutex<Option<(String, super::whisper::WhisperTranscriber)>>,
> = once_cell::sync::Lazy::new(|| Mutex::new(None));
// Bundled tokenizer JSON (2.4MB)
const WHISPER_TOKENIZER_JSON: &str = include_str!("whisper_data/tokens.json");
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Deserialize, Serialize, ToSchema)]
#[serde(rename_all = "lowercase")]
pub enum DictationProvider {
OpenAI,
ElevenLabs,
Groq,
Local,
}
pub struct DictationProviderDef {
pub provider: DictationProvider,
pub config_key: &'static str,
pub default_base_url: &'static str,
pub endpoint_path: &'static str,
pub host_key: Option<&'static str>,
pub description: &'static str,
pub uses_provider_config: bool,
pub settings_path: Option<&'static str>,
}
pub const PROVIDERS: &[DictationProviderDef] = &[
DictationProviderDef {
provider: DictationProvider::OpenAI,
config_key: "OPENAI_API_KEY",
default_base_url: "https://api.openai.com",
endpoint_path: "v1/audio/transcriptions",
host_key: Some("OPENAI_HOST"),
description: "Uses OpenAI Whisper API for high-quality transcription.",
uses_provider_config: true,
settings_path: Some("Settings > Models"),
},
DictationProviderDef {
provider: DictationProvider::Groq,
config_key: "GROQ_API_KEY",
default_base_url: "https://api.groq.com/openai/v1",
endpoint_path: "audio/transcriptions",
host_key: None,
description: "Uses Groq's ultra-fast Whisper implementation with LPU acceleration.",
uses_provider_config: false,
settings_path: None,
},
DictationProviderDef {
provider: DictationProvider::ElevenLabs,
config_key: "ELEVENLABS_API_KEY",
default_base_url: "https://api.elevenlabs.io",
endpoint_path: "v1/speech-to-text",
host_key: None,
description: "Uses ElevenLabs speech-to-text API for advanced voice processing.",
uses_provider_config: false,
settings_path: None,
},
DictationProviderDef {
provider: DictationProvider::Local,
config_key: LOCAL_WHISPER_MODEL_CONFIG_KEY,
default_base_url: "",
endpoint_path: "",
host_key: None,
description: "Uses local Whisper model for transcription. No API key needed.",
uses_provider_config: false,
settings_path: None,
},
];
pub fn get_provider_def(provider: DictationProvider) -> &'static DictationProviderDef {
PROVIDERS
.iter()
.find(|def| def.provider == provider)
.unwrap() // Safe because all enum variants are in PROVIDERS
}
pub fn is_configured(provider: DictationProvider) -> bool {
let config = Config::global();
match provider {
DictationProvider::Local => config
.get(LOCAL_WHISPER_MODEL_CONFIG_KEY, false)
.ok()
.and_then(|v| v.as_str().map(|s| s.to_string()))
.and_then(|id| super::whisper::get_model(&id))
.is_some_and(|m| m.is_downloaded()),
_ => {
let def = get_provider_def(provider);
config.get_secret::<String>(def.config_key).is_ok()
}
}
}
pub async fn transcribe_local(audio_bytes: Vec<u8>) -> Result<String> {
// Run transcription in a blocking task to avoid blocking the async runtime
tokio::task::spawn_blocking(move || {
// Get model ID from config
let config = Config::global();
let model_id = config
.get(LOCAL_WHISPER_MODEL_CONFIG_KEY, false)
.ok()
.and_then(|v| v.as_str().map(|s| s.to_string()))
.ok_or_else(|| anyhow::anyhow!("Local Whisper model not configured"))?;
// Convert model ID to full path
let model = super::whisper::get_model(&model_id)
.ok_or_else(|| anyhow::anyhow!("Unknown model: {}", model_id))?;
let model_path = model.local_path();
// Get or initialize the transcriber
let mut transcriber_lock = LOCAL_TRANSCRIBER
.lock()
.map_err(|e| anyhow::anyhow!("Failed to lock transcriber: {}", e))?;
// Check if we need to load/reload the transcriber
let model_path_str = model_path.to_string_lossy().to_string();
let needs_reload = match transcriber_lock.as_ref() {
None => true,
Some((cached_path, _)) => cached_path != &model_path_str,
};
if needs_reload {
tracing::info!("Loading Whisper model from: {}", model_path.display());
let transcriber = super::whisper::WhisperTranscriber::new_with_tokenizer(
&model_id,
&model_path,
WHISPER_TOKENIZER_JSON,
)?;
*transcriber_lock = Some((model_path_str, transcriber));
}
// Transcribe the audio
let (_, transcriber) = transcriber_lock.as_mut().unwrap();
let text = transcriber
.transcribe(&audio_bytes)
.context("Transcription failed")?;
Ok(text)
})
.await
.context("Transcription task failed")?
}
fn build_api_client(provider: DictationProvider) -> Result<ApiClient> {
let config = Config::global();
let def = get_provider_def(provider);
let api_key = config
.get_secret(def.config_key)
.context(format!("{} not configured", def.config_key))?;
let base_url = if let Some(host_key) = def.host_key {
config
.get(host_key, false)
.ok()
.and_then(|v| v.as_str().map(|s| s.to_string()))
.unwrap_or_else(|| def.default_base_url.to_string())
} else {
def.default_base_url.to_string()
};
let auth = match provider {
DictationProvider::OpenAI => AuthMethod::BearerToken(api_key),
DictationProvider::Groq => AuthMethod::BearerToken(api_key),
DictationProvider::ElevenLabs => AuthMethod::ApiKey {
header_name: "xi-api-key".to_string(),
key: api_key,
},
DictationProvider::Local => anyhow::bail!("Local provider should not use API client"),
};
ApiClient::with_timeout(base_url, auth, REQUEST_TIMEOUT).context("Failed to create API client")
}
pub async fn transcribe_with_provider(
provider: DictationProvider,
model_param: String,
model_value: String,
audio_bytes: Vec<u8>,
extension: &str,
mime_type: &str,
) -> Result<String> {
let client = build_api_client(provider)?;
let def = get_provider_def(provider);
let part = reqwest::multipart::Part::bytes(audio_bytes)
.file_name(format!("audio.{}", extension))
.mime_str(mime_type)
.context("Failed to create multipart")?;
let form = reqwest::multipart::Form::new()
.part("file", part)
.text(model_param, model_value);
let response = client
.request(None, def.endpoint_path)
.multipart_post(form)
.await
.context("Request failed")?;
if !response.status().is_success() {
let status = response.status();
let error_text = response.text().await.unwrap_or_default();
if status == 401 || error_text.contains("Invalid API key") {
anyhow::bail!("Invalid API key");
} else if status == 429 || error_text.contains("quota") {
anyhow::bail!("Rate limit exceeded");
} else {
anyhow::bail!("API error: {}", error_text);
}
}
let data: serde_json::Value = response.json().await.context("Failed to parse response")?;
let text = data["text"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("Missing 'text' field in response"))?
.to_string();
Ok(text)
}
crates/goose/src/dictation/whisper.rs
+757
View File
@@ -0,0 +1,757 @@
//! Local Whisper transcription using Candle
//!
//! This module provides local audio transcription using OpenAI's Whisper model
//! via the Candle ML framework. It supports loading GGUF quantized models for
//! efficient CPU inference.
//! Heavily "inspired" by the Candle Whisper example:
//! https://github.com/huggingface/candle/tree/main/candle-examples/whisper
use crate::config::paths::Paths;
pub const LOCAL_WHISPER_MODEL_CONFIG_KEY: &str = "LOCAL_WHISPER_MODEL";
use anyhow::{Context, Result};
use candle_core::{Device, IndexOp, Tensor};
use candle_nn::ops::log_softmax;
use candle_transformers::models::whisper::{self as m, audio, Config, N_FRAMES};
use serde::{Deserialize, Serialize};
use std::io::Cursor;
use std::path::{Path, PathBuf};
use symphonia::core::audio::{AudioBufferRef, Layout, Signal};
use symphonia::core::codecs::DecoderOptions;
use symphonia::core::formats::FormatOptions;
use symphonia::core::io::MediaSourceStream;
use symphonia::core::meta::MetadataOptions;
use symphonia::core::probe::Hint;
use tokenizers::Tokenizer;
use utoipa::ToSchema;
// Common suppress tokens for all Whisper models
const SUPPRESS_TOKENS: &[u32] = &[
1, 2, 7, 8, 9, 10, 14, 25, 26, 27, 28, 29, 31, 58, 59, 60, 61, 62, 63, 90, 91, 92, 93, 359,
503, 522, 542, 873, 893, 902, 918, 922, 931, 1350, 1853, 1982, 2460, 2627, 3246, 3253, 3268,
3536, 3846, 3961, 4183, 4667, 6585, 6647, 7273, 9061, 9383, 10428, 10929, 11938, 12033, 12331,
12562, 13793, 14157, 14635, 15265, 15618, 16553, 16604, 18362, 18956, 20075, 21675, 22520,
26130, 26161, 26435, 28279, 29464, 31650, 32302, 32470, 36865, 42863, 47425, 49870, 50254,
50258, 50360, 50362,
];
// Special token IDs
const SOT_TOKEN: u32 = 50258;
const TRANSCRIBE_TOKEN: u32 = 50359;
const EOT_TOKEN: u32 = 50257;
const TIMESTAMP_BEGIN: u32 = 50364;
const SAMPLE_BEGIN: usize = 3;
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct WhisperModel {
/// Model identifier (e.g., "tiny", "base", "small")
pub id: &'static str,
/// Model file size in MB
pub size_mb: u32,
/// Download URL from HuggingFace
pub url: &'static str,
/// Description
pub description: &'static str,
}
const MODELS: &[WhisperModel] = &[
WhisperModel {
id: "tiny",
size_mb: 40,
url: "https://huggingface.co/oxide-lab/whisper-tiny-GGUF/resolve/main/model-tiny-q80.gguf",
description: "Fastest, ~2-3x realtime on CPU (5-10x with GPU)",
},
WhisperModel {
id: "base",
size_mb: 78,
url: "https://huggingface.co/oxide-lab/whisper-base-GGUF/resolve/main/whisper-base-q8_0.gguf",
description: "Good balance, ~1.5-2x realtime on CPU (4-8x with GPU)",
},
WhisperModel {
id: "small",
size_mb: 247,
url: "https://huggingface.co/oxide-lab/whisper-small-GGUF/resolve/main/whisper-small-q8_0.gguf",
description: "High accuracy, ~0.8-1x realtime on CPU (3-5x with GPU)",
},
WhisperModel {
id: "medium",
size_mb: 777,
url: "https://huggingface.co/oxide-lab/whisper-medium-GGUF/resolve/main/whisper-medium-q8_0.gguf",
description: "Highest accuracy, ~0.5x realtime on CPU (2-4x with GPU)",
},
];
impl WhisperModel {
pub fn local_path(&self) -> PathBuf {
let filename = self.url.rsplit('/').next().unwrap_or("");
Paths::in_data_dir("models").join(filename)
}
pub fn is_downloaded(&self) -> bool {
self.local_path().exists()
}
pub fn config(&self) -> Config {
match self.id {
"tiny" => Config {
num_mel_bins: 80,
max_source_positions: 1500,
d_model: 384,
encoder_attention_heads: 6,
encoder_layers: 4,
decoder_attention_heads: 6,
decoder_layers: 4,
vocab_size: 51865,
suppress_tokens: SUPPRESS_TOKENS.to_vec(),
max_target_positions: 448,
},
"base" => Config {
num_mel_bins: 80,
max_source_positions: 1500,
d_model: 512,
encoder_attention_heads: 8,
encoder_layers: 6,
decoder_attention_heads: 8,
decoder_layers: 6,
vocab_size: 51865,
suppress_tokens: SUPPRESS_TOKENS.to_vec(),
max_target_positions: 448,
},
"small" => Config {
num_mel_bins: 80,
max_source_positions: 1500,
d_model: 768,
encoder_attention_heads: 12,
encoder_layers: 12,
decoder_attention_heads: 12,
decoder_layers: 12,
vocab_size: 51865,
suppress_tokens: SUPPRESS_TOKENS.to_vec(),
max_target_positions: 448,
},
"medium" => Config {
num_mel_bins: 80,
max_source_positions: 1500,
d_model: 1024,
encoder_attention_heads: 16,
encoder_layers: 24,
decoder_attention_heads: 16,
decoder_layers: 24,
vocab_size: 51865,
suppress_tokens: SUPPRESS_TOKENS.to_vec(),
max_target_positions: 448,
},
_ => {
tracing::warn!("Unknown model '{}', falling back to tiny config", self.id);
Config {
num_mel_bins: 80,
max_source_positions: 1500,
d_model: 384,
encoder_attention_heads: 6,
encoder_layers: 4,
decoder_attention_heads: 6,
decoder_layers: 4,
vocab_size: 51865,
suppress_tokens: SUPPRESS_TOKENS.to_vec(),
max_target_positions: 448,
}
}
}
}
}
pub fn available_models() -> &'static [WhisperModel] {
MODELS
}
pub fn get_model(id: &str) -> Option<&'static WhisperModel> {
MODELS.iter().find(|m| m.id == id)
}
pub fn recommend_model() -> &'static str {
let has_gpu_or_metal = Device::new_cuda(0).is_ok() || Device::new_metal(0).is_ok();
if has_gpu_or_metal {
"small"
} else {
let cpu_count = sys_info::cpu_num().unwrap_or(1) as u64;
let cpu_speed_mhz = sys_info::cpu_speed().unwrap_or(0);
if cpu_count * cpu_speed_mhz >= 16_000 {
"base"
} else {
"tiny"
}
}
}
pub struct WhisperTranscriber {
model: m::quantized_model::Whisper,
config: Config,
device: Device,
mel_filters: Vec<f32>,
tokenizer: Tokenizer,
eot_token: u32,
no_timestamps_token: u32,
language_token: u32,
max_initial_timestamp_index: u32,
}
impl WhisperTranscriber {
pub fn new_with_tokenizer<P: AsRef<Path>>(
model_id: &str,
model_path: P,
bundled_tokenizer: &str,
) -> Result<Self> {
let device = if let Ok(device) = Device::new_cuda(0) {
device
} else if let Ok(device) = Device::new_metal(0) {
device
} else {
Device::Cpu
};
let model_path_ref = model_path.as_ref();
if !model_path_ref.exists() {
anyhow::bail!("Model file not found: {}", model_path_ref.display());
}
let model =
get_model(model_id).ok_or_else(|| anyhow::anyhow!("Unknown model: {}", model_id))?;
let config = model.config();
let mel_bytes = match config.num_mel_bins {
80 => include_bytes!("whisper_data/melfilters.bytes").as_slice(),
128 => include_bytes!("whisper_data/melfilters128.bytes").as_slice(),
nmel => anyhow::bail!("unexpected num_mel_bins {nmel}"),
};
let mut mel_filters = vec![0f32; mel_bytes.len() / 4];
byteorder::ReadBytesExt::read_f32_into::<byteorder::LittleEndian>(
&mut &mel_bytes[..],
&mut mel_filters,
)?;
let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(
model_path_ref,
&device,
)?;
let model = m::quantized_model::Whisper::load(&vb, config.clone())?;
let tokenizer = Self::load_tokenizer(model_path_ref, Some(bundled_tokenizer))?;
Ok(Self {
model,
config,
device,
mel_filters,
tokenizer,
eot_token: 50257,
no_timestamps_token: 50363,
language_token: 50259,
max_initial_timestamp_index: 50,
})
}
fn load_tokenizer(model_dir: &Path, bundled_tokenizer: Option<&str>) -> Result<Tokenizer> {
let tokenizer_path = model_dir
.parent()
.unwrap_or(model_dir)
.join("tokenizer.json");
if tokenizer_path.exists() {
return Tokenizer::from_file(tokenizer_path)
.map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e));
}
if let Some(tokenizer_json) = bundled_tokenizer {
if let Some(parent) = tokenizer_path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(&tokenizer_path, tokenizer_json)?;
return Tokenizer::from_file(tokenizer_path)
.map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e));
}
anyhow::bail!(
"Tokenizer not found at {} and no bundled tokenizer provided",
tokenizer_path.display()
)
}
pub fn transcribe(&mut self, audio_data: &[u8]) -> Result<String> {
let mel_tensor = self.prepare_audio_input(audio_data)?;
let (_, _, content_frames) = mel_tensor.dims3()?;
let num_segments = content_frames.div_ceil(N_FRAMES);
let mut all_text_tokens = Vec::new();
let mut seek = 0;
let mut segment_num = 0;
while seek < content_frames {
segment_num += 1;
let segment_size = usize::min(content_frames - seek, N_FRAMES);
let segment_text_tokens =
self.process_segment(&mel_tensor, seek, segment_size, segment_num, num_segments)?;
all_text_tokens.extend(segment_text_tokens);
seek += segment_size;
}
self.decode_tokens(&all_text_tokens)
}
fn prepare_audio_input(&self, audio_data: &[u8]) -> Result<Tensor> {
let pcm_data = decode_audio_simple(audio_data)?;
let mel = audio::pcm_to_mel(&self.config, &pcm_data, &self.mel_filters);
let mel_len = mel.len();
let mel_tensor = Tensor::from_vec(
mel,
(
1,
self.config.num_mel_bins,
mel_len / self.config.num_mel_bins,
),
&self.device,
)?;
Ok(mel_tensor)
}
fn process_segment(
&mut self,
mel_tensor: &Tensor,
seek: usize,
segment_size: usize,
_segment_num: usize,
_num_segments: usize,
) -> Result<Vec<u32>> {
let _time_offset = (seek * 160) as f32 / 16000.0; // HOP_LENGTH = 160
let _segment_duration = (segment_size * 160) as f32 / 16000.0;
let mel_segment = mel_tensor.narrow(2, seek, segment_size)?;
self.model.decoder.reset_kv_cache();
let audio_features = self.model.encoder.forward(&mel_segment, true)?;
let suppress_tokens = {
let mut suppress = vec![0f32; self.config.vocab_size];
for &token_id in &self.config.suppress_tokens {
if (token_id as usize) < suppress.len() {
suppress[token_id as usize] = f32::NEG_INFINITY;
}
}
suppress[self.no_timestamps_token as usize] = f32::NEG_INFINITY;
Tensor::from_vec(suppress, self.config.vocab_size, &self.device)?
};
let mut tokens = vec![SOT_TOKEN, self.language_token, TRANSCRIBE_TOKEN];
let sample_len = self.config.max_target_positions / 2;
for i in 0..sample_len {
let tokens_tensor = Tensor::new(tokens.as_slice(), &self.device)?.unsqueeze(0)?;
let ys = self
.model
.decoder
.forward(&tokens_tensor, &audio_features, i == 0)?;
let (_, seq_len, _) = ys.dims3()?;
let mut logits = self
.model
.decoder
.final_linear(&ys.i((..1, seq_len - 1..))?)?
.i(0)?
.i(0)?;
logits = self.apply_timestamp_rules(&logits, &tokens)?;
let logits = logits.broadcast_add(&suppress_tokens)?;
let logits_v: Vec<f32> = logits.to_vec1()?;
let next_token = logits_v
.iter()
.enumerate()
.max_by(|(_, u), (_, v)| u.total_cmp(v))
.map(|(i, _)| i as u32)
.unwrap();
tokens.push(next_token);
if next_token == EOT_TOKEN || tokens.len() > self.config.max_target_positions {
break;
}
}
let segment_text_tokens: Vec<u32> = tokens[3..]
.iter()
.filter(|&&t| t != EOT_TOKEN && t < TIMESTAMP_BEGIN)
.copied()
.collect();
Ok(segment_text_tokens)
}
fn apply_timestamp_rules(&self, input_logits: &Tensor, tokens: &[u32]) -> Result<Tensor> {
let device = input_logits.device().clone();
let vocab_size = self.model.config.vocab_size as u32;
let sampled_tokens = if tokens.len() > SAMPLE_BEGIN {
&tokens[SAMPLE_BEGIN..]
} else {
&[]
};
let mut masks = Vec::new();
let mut mask_buffer = vec![0.0f32; vocab_size as usize];
self.apply_timestamp_pairing_rule(
sampled_tokens,
vocab_size,
&mut masks,
&mut mask_buffer,
&device,
)?;
self.apply_initial_timestamp_rule(
tokens.len(),
vocab_size,
&mut masks,
&mut mask_buffer,
&device,
)?;
let mut logits = input_logits.clone();
for mask in masks {
logits = logits.broadcast_add(&mask)?;
}
logits =
self.apply_timestamp_probability_rule(&logits, vocab_size, &mut mask_buffer, &device)?;
Ok(logits)
}
fn apply_timestamp_pairing_rule(
&self,
sampled_tokens: &[u32],
vocab_size: u32,
masks: &mut Vec<Tensor>,
mask_buffer: &mut [f32],
device: &Device,
) -> Result<()> {
if sampled_tokens.is_empty() {
return Ok(());
}
let last_was_timestamp = sampled_tokens
.last()
.map(|&t| t >= TIMESTAMP_BEGIN)
.unwrap_or(false);
let penultimate_was_timestamp = if sampled_tokens.len() >= 2 {
sampled_tokens[sampled_tokens.len() - 2] >= TIMESTAMP_BEGIN
} else {
false
};
if last_was_timestamp {
if penultimate_was_timestamp {
for i in 0..vocab_size {
mask_buffer[i as usize] = if i >= TIMESTAMP_BEGIN {
f32::NEG_INFINITY
} else {
0.0
};
}
masks.push(Tensor::new(mask_buffer as &[f32], device)?);
} else {
for i in 0..vocab_size {
mask_buffer[i as usize] = if i < self.eot_token {
f32::NEG_INFINITY
} else {
0.0
};
}
masks.push(Tensor::new(mask_buffer as &[f32], device)?);
}
}
let timestamp_tokens: Vec<u32> = sampled_tokens
.iter()
.filter(|&&t| t >= TIMESTAMP_BEGIN)
.cloned()
.collect();
if !timestamp_tokens.is_empty() {
let timestamp_last = if last_was_timestamp && !penultimate_was_timestamp {
*timestamp_tokens.last().unwrap()
} else {
timestamp_tokens.last().unwrap() + 1
};
for i in 0..vocab_size {
mask_buffer[i as usize] = if i >= TIMESTAMP_BEGIN && i < timestamp_last {
f32::NEG_INFINITY
} else {
0.0
};
}
masks.push(Tensor::new(mask_buffer as &[f32], device)?);
}
Ok(())
}
fn apply_initial_timestamp_rule(
&self,
tokens_len: usize,
vocab_size: u32,
masks: &mut Vec<Tensor>,
mask_buffer: &mut [f32],
device: &Device,
) -> Result<()> {
if tokens_len != SAMPLE_BEGIN {
return Ok(());
}
for i in 0..vocab_size {
mask_buffer[i as usize] = if i < TIMESTAMP_BEGIN {
f32::NEG_INFINITY
} else {
0.0
};
}
masks.push(Tensor::new(mask_buffer as &[f32], device)?);
let last_allowed = TIMESTAMP_BEGIN + self.max_initial_timestamp_index;
if last_allowed < vocab_size {
for i in 0..vocab_size {
mask_buffer[i as usize] = if i > last_allowed {
f32::NEG_INFINITY
} else {
0.0
};
}
masks.push(Tensor::new(mask_buffer as &[f32], device)?);
}
Ok(())
}
fn apply_timestamp_probability_rule(
&self,
logits: &Tensor,
vocab_size: u32,
mask_buffer: &mut [f32],
device: &Device,
) -> Result<Tensor> {
let log_probs = log_softmax(logits, 0)?;
let timestamp_log_probs = log_probs.narrow(
0,
TIMESTAMP_BEGIN as usize,
vocab_size as usize - TIMESTAMP_BEGIN as usize,
)?;
let text_log_probs = log_probs.narrow(0, 0, TIMESTAMP_BEGIN as usize)?;
let timestamp_logprob = {
let max_val = timestamp_log_probs.max(0)?;
let shifted = timestamp_log_probs.broadcast_sub(&max_val)?;
let exp_shifted = shifted.exp()?;
let sum_exp = exp_shifted.sum(0)?;
let log_sum = sum_exp.log()?;
max_val.broadcast_add(&log_sum)?.to_scalar::<f32>()?
};
let max_text_token_logprob: f32 = text_log_probs.max(0)?.to_scalar::<f32>()?;
if timestamp_logprob > max_text_token_logprob {
for i in 0..vocab_size {
mask_buffer[i as usize] = if i < TIMESTAMP_BEGIN {
f32::NEG_INFINITY
} else {
0.0
};
}
let mask_tensor = Tensor::new(mask_buffer as &[f32], device)?;
return logits.broadcast_add(&mask_tensor).map_err(Into::into);
}
Ok(logits.clone())
}
fn decode_tokens(&self, tokens: &[u32]) -> Result<String> {
self.tokenizer
.decode(tokens, true)
.map_err(|e| anyhow::anyhow!("Failed to decode tokens: {}", e))
}
}
fn decode_audio_simple(audio_data: &[u8]) -> Result<Vec<f32>> {
let audio_vec = audio_data.to_vec();
let cursor = Cursor::new(audio_vec);
let mss = MediaSourceStream::new(Box::new(cursor), Default::default());
let hint = Hint::new();
let probed = symphonia::default::get_probe()
.format(
&hint,
mss,
&FormatOptions::default(),
&MetadataOptions::default(),
)
.context("Failed to probe audio format - unsupported format")?;
let mut format = probed.format;
let track = format
.default_track()
.context("No default audio track found")?;
let sample_rate = track
.codec_params
.sample_rate
.context("No sample rate in audio track")?;
let channels = if let Some(ch) = track.codec_params.channels {
ch.count()
} else if let Some(layout) = track.codec_params.channel_layout {
match layout {
Layout::Mono => 1,
Layout::Stereo => 2,
_ => 1,
}
} else {
anyhow::bail!("No channel information in audio track (neither channels nor channel_layout)")
};
let mut decoder = symphonia::default::get_codecs()
.make(&track.codec_params, &DecoderOptions::default())
.context("Failed to create audio decoder - please ensure browser sends WAV format audio")?;
let mut pcm_data = Vec::new();
loop {
let packet = match format.next_packet() {
Ok(packet) => packet,
Err(symphonia::core::errors::Error::IoError(e))
if e.kind() == std::io::ErrorKind::UnexpectedEof =>
{
break;
}
Err(e) => return Err(e).context("Failed to read audio packet")?,
};
match decoder.decode(&packet) {
Ok(decoded) => {
pcm_data.extend(audio_buffer_to_f32(&decoded));
}
Err(symphonia::core::errors::Error::DecodeError(_)) => {
continue;
}
Err(e) => return Err(e).context("Failed to decode audio packet")?,
}
}
let mono_data = if channels > 1 {
convert_to_mono(&pcm_data, channels)
} else {
pcm_data
};
let resampled = if sample_rate != 16000 {
resample_audio(&mono_data, sample_rate, 16000)?
} else {
mono_data
};
Ok(resampled)
}
fn audio_buffer_to_f32(buffer: &AudioBufferRef) -> Vec<f32> {
let num_channels = buffer.spec().channels.count();
let num_frames = buffer.frames();
let mut samples = Vec::with_capacity(num_frames * num_channels);
match buffer {
AudioBufferRef::F32(buf) => {
for frame_idx in 0..num_frames {
for ch_idx in 0..num_channels {
samples.push(buf.chan(ch_idx)[frame_idx]);
}
}
}
AudioBufferRef::S16(buf) => {
for frame_idx in 0..num_frames {
for ch_idx in 0..num_channels {
samples.push(buf.chan(ch_idx)[frame_idx] as f32 / 32768.0);
}
}
}
AudioBufferRef::S32(buf) => {
for frame_idx in 0..num_frames {
for ch_idx in 0..num_channels {
samples.push(buf.chan(ch_idx)[frame_idx] as f32 / 2147483648.0);
}
}
}
AudioBufferRef::F64(buf) => {
for frame_idx in 0..num_frames {
for ch_idx in 0..num_channels {
samples.push(buf.chan(ch_idx)[frame_idx] as f32);
}
}
}
_ => {
tracing::warn!("Unsupported audio buffer format, returning silence");
}
}
samples
}
fn convert_to_mono(data: &[f32], channels: usize) -> Vec<f32> {
if channels == 1 {
return data.to_vec();
}
let frames = data.len() / channels;
let mut mono = Vec::with_capacity(frames);
for frame_idx in 0..frames {
let mut sum = 0.0;
for ch in 0..channels {
sum += data[frame_idx * channels + ch];
}
mono.push(sum / channels as f32);
}
mono
}
fn resample_audio(data: &[f32], from_rate: u32, to_rate: u32) -> Result<Vec<f32>> {
use rubato::{
Resampler, SincFixedIn, SincInterpolationParameters, SincInterpolationType, WindowFunction,
};
if from_rate == to_rate {
return Ok(data.to_vec());
}
let params = SincInterpolationParameters {
sinc_len: 256,
f_cutoff: 0.95,
interpolation: SincInterpolationType::Linear,
oversampling_factor: 256,
window: WindowFunction::BlackmanHarris2,
};
let mut resampler = SincFixedIn::<f32>::new(
to_rate as f64 / from_rate as f64,
2.0,
params,
data.len(),
1,
)?;
let waves_in = vec![data.to_vec()];
let waves_out = resampler.process(&waves_in, None)?;
Ok(waves_out[0].clone())
}
crates/goose/src/dictation/whisper_data/melfilters.bytes
BIN
View File
Binary file not shown.
crates/goose/src/dictation/whisper_data/melfilters128.bytes
BIN
View File
Binary file not shown.
crates/goose/src/dictation/whisper_data/tokens.json
+114853
View File
File diff suppressed because it is too large Load Diff
crates/goose/src/lib.rs
+1
View File
@@ -4,6 +4,7 @@ pub mod builtin_extension;
pub mod config;
pub mod context_mgmt;
pub mod conversation;
pub mod dictation;
pub mod execution;
pub mod goose_apps;
pub mod hints;