diff --git a/crates/ai-formatting/src/pipeline.rs b/crates/ai-formatting/src/pipeline.rs index 8f24878..304f78f 100644 --- a/crates/ai-formatting/src/pipeline.rs +++ b/crates/ai-formatting/src/pipeline.rs @@ -9,6 +9,9 @@ pub struct PostProcessOptions { pub british_english: bool, pub anti_hallucination: bool, pub format_mode: FormatMode, + /// Custom vocabulary terms loaded from the user's dictionary. Injected + /// into the LLM cleanup prompt so the model knows how to spell them. + pub dictionary_terms: Vec, } /// How aggressively to format the transcript text. @@ -82,6 +85,19 @@ mod tests { ] } + #[test] + fn dictionary_terms_stored_on_options() { + let options = PostProcessOptions { + remove_fillers: false, + british_english: false, + anti_hallucination: false, + format_mode: FormatMode::Raw, + dictionary_terms: vec!["Wren".to_string(), "CORBEL".to_string()], + }; + assert_eq!(options.dictionary_terms.len(), 2); + assert_eq!(options.dictionary_terms[0], "Wren"); + } + #[test] fn post_process_applies_all_filters() { let mut segments = make_segments(); @@ -90,6 +106,7 @@ mod tests { british_english: true, anti_hallucination: true, format_mode: FormatMode::Clean, + dictionary_terms: vec![], }; post_process_segments(&mut segments, &options); @@ -110,6 +127,7 @@ mod tests { british_english: false, anti_hallucination: false, format_mode: FormatMode::Smart, + dictionary_terms: vec![], }; post_process_segments(&mut segments, &options); diff --git a/src-tauri/src/commands/live.rs b/src-tauri/src/commands/live.rs index 220bd9f..47cc720 100644 --- a/src-tauri/src/commands/live.rs +++ b/src-tauri/src/commands/live.rs @@ -166,11 +166,20 @@ pub async fn start_live_transcription_session( let worker_status = status_channel.clone(); let worker_results = result_channel.clone(); + let dictionary_terms: Vec = + kon_storage::database::list_dictionary(&state.db) + .await + .unwrap_or_default() + .into_iter() + .map(|e| e.term) + .collect(); + let handle = tokio::task::spawn_blocking(move || { run_live_session( session_id, engine, config, + dictionary_terms, worker_results, worker_status, worker_stop, @@ -250,6 +259,7 @@ fn run_live_session( session_id: u64, engine: Arc, config: StartLiveTranscriptionConfig, + dictionary_terms: Vec, result_channel: Channel, status_channel: Channel, stop_flag: Arc, @@ -284,6 +294,7 @@ fn run_live_session( &mut inflight, session_id, &config, + &dictionary_terms, &result_channel, &status_channel, )? {} @@ -386,6 +397,7 @@ fn run_live_session( &mut inflight, session_id, &config, + &dictionary_terms, &result_channel, &status_channel, )?; @@ -512,6 +524,7 @@ fn poll_inference( inflight: &mut Option, session_id: u64, config: &StartLiveTranscriptionConfig, + dictionary_terms: &[String], result_channel: &Channel, status_channel: &Channel, ) -> Result, String> { @@ -531,6 +544,7 @@ fn poll_inference( british_english: config.british_english, anti_hallucination: config.anti_hallucination, format_mode: FormatMode::parse(&config.format_mode), + dictionary_terms: dictionary_terms.to_vec(), }, ); let segment_count = segments.len(); diff --git a/src-tauri/src/commands/transcription.rs b/src-tauri/src/commands/transcription.rs index 04729f2..458b51b 100644 --- a/src-tauri/src/commands/transcription.rs +++ b/src-tauri/src/commands/transcription.rs @@ -50,6 +50,13 @@ pub async fn transcribe_pcm( .await .map_err(|e| e.to_string())??; + let dictionary_terms: Vec = kon_storage::database::list_dictionary(&state.db) + .await + .unwrap_or_default() + .into_iter() + .map(|e| e.term) + .collect(); + let mut segments: Vec = timed.transcript.segments().to_vec(); post_process_segments( &mut segments, @@ -58,6 +65,7 @@ pub async fn transcribe_pcm( british_english, anti_hallucination, format_mode: FormatMode::parse(&format_mode), + dictionary_terms, }, ); @@ -114,6 +122,13 @@ pub async fn transcribe_file( .await .map_err(|e| e.to_string())??; + let dictionary_terms: Vec = kon_storage::database::list_dictionary(&state.db) + .await + .unwrap_or_default() + .into_iter() + .map(|e| e.term) + .collect(); + let mut segments: Vec = timed.transcript.segments().to_vec(); post_process_segments( &mut segments, @@ -122,6 +137,7 @@ pub async fn transcribe_file( british_english, anti_hallucination, format_mode: FormatMode::parse(&format_mode), + dictionary_terms, }, ); @@ -157,6 +173,13 @@ pub async fn transcribe_pcm_parakeet( .await .map_err(|e| e.to_string())??; + let dictionary_terms: Vec = kon_storage::database::list_dictionary(&state.db) + .await + .unwrap_or_default() + .into_iter() + .map(|e| e.term) + .collect(); + let mut segments: Vec = timed.transcript.segments().to_vec(); post_process_segments( &mut segments, @@ -165,6 +188,7 @@ pub async fn transcribe_pcm_parakeet( british_english, anti_hallucination, format_mode: FormatMode::parse(&format_mode), + dictionary_terms, }, );