1use std::{
4 collections::BTreeMap,
5 time::{SystemTime, UNIX_EPOCH},
6};
7#[cfg(feature = "diarization")]
8use std::{collections::BTreeSet, path::Path};
9
10use audio_analysis_speakers::{
11 SpeakerAudio, SpeakerEmbedding, SpeakerEmbeddingExtractor, SpectralSpeakerEmbedder,
12};
13use audio_analysis_transcription::{LoadedAudio, TranscriptionPipelineResponse};
14use text_transcripts::TranscriptionContract;
15
16#[cfg(feature = "diarization")]
17use crate::config::is_pyannote_diarization_model;
18use crate::config::{
19 NativeWhisperxConfig, NativeWhisperxError, OutputConfig, SpeakerCorrectionReport,
20 SpeakerCorrectionRequest,
21};
22use crate::config_mapping::map_input_source;
23use crate::output::write_outputs;
24
25#[cfg(feature = "diarization")]
26pub(crate) fn save_draft_speakers_from_response(
27 response: &mut TranscriptionPipelineResponse,
28 config: &NativeWhisperxConfig,
29) -> Result<(), NativeWhisperxError> {
30 if config.asr.provider != crate::config::AsrProvider::Native
31 || !config.diarization.enabled
32 || config.diarization.disable_speaker_library
33 || is_pyannote_diarization_model(&config.diarization.model_id)
34 {
35 return Ok(());
36 }
37
38 if !config.diarization.save_draft_speakers {
39 response
40 .diagnostics
41 .push("speakerLibraryDraftProfilesSaved=0".to_string());
42 return Ok(());
43 }
44
45 let Some(diarization) = &response.diarization else {
46 response
47 .diagnostics
48 .push("speakerLibraryDraftProfilesSaved=0".to_string());
49 return Ok(());
50 };
51 let current_dir = std::env::current_dir()?;
52 let resolved =
53 crate::resolve_speaker_directory(&config.diarization.speaker_directory, ¤t_dir)?;
54 let mut library = crate::speaker_directory::load_speaker_library_if_present(&resolved.path)?
55 .unwrap_or_default();
56 let existing_labels = library
57 .profiles()
58 .map(|profile| profile.id().as_str().to_string())
59 .collect::<BTreeSet<_>>();
60 let labels = diarization
61 .segments
62 .iter()
63 .filter(|segment| is_transient_speaker_label(&segment.speaker))
64 .map(|segment| segment.speaker.clone())
65 .collect::<BTreeSet<_>>();
66 if labels.is_empty() {
67 response
68 .diagnostics
69 .push("speakerLibraryDraftProfilesSaved=0".to_string());
70 return Ok(());
71 }
72
73 let audio = LoadedAudio::mono_16khz_from_source(&map_input_source(&config.input))
74 .map_err(|error| NativeWhisperxError::Transcription(error.to_string()))?;
75 let mut saved = 0usize;
76 for label in labels {
77 if existing_labels.contains(&label) {
78 continue;
79 }
80 let ranges = diarization
81 .segments
82 .iter()
83 .filter(|segment| segment.speaker == label)
84 .map(|segment| crate::SpeakerCorrectionRange {
85 start_seconds: segment.start_seconds as f64,
86 end_seconds: segment.end_seconds as f64,
87 })
88 .collect::<Vec<_>>();
89 let embedding = speaker_embedding_for_ranges(&audio, &ranges)?;
90 let draft_id = format!(
91 "draft-{}-{}",
92 slug_speaker_id(&label),
93 current_unix_seconds()
94 );
95 let draft_label = format!("draft_{}", slug_speaker_id(&label));
96 let mut metadata = BTreeMap::new();
97 metadata.insert("status".to_string(), "draft".to_string());
98 metadata.insert("detectedLabel".to_string(), label);
99 let now = current_unix_seconds_string();
100 metadata.insert("createdAt".to_string(), now.clone());
101 metadata.insert("updatedAt".to_string(), now);
102 let (updated, _) = crate::speaker_directory::upsert_speaker_profile_embedding(
103 &library,
104 &draft_id,
105 &draft_label,
106 embedding,
107 metadata,
108 )?;
109 library = updated;
110 saved += 1;
111 }
112 if saved > 0 {
113 crate::speaker_directory::save_speaker_library(&resolved.path, &library)?;
114 }
115 response
116 .diagnostics
117 .push(format!("speakerLibraryDraftProfilesSaved={saved}"));
118 Ok(())
119}
120
121#[cfg(not(feature = "diarization"))]
122pub(crate) fn save_draft_speakers_from_response(
123 _response: &mut TranscriptionPipelineResponse,
124 _config: &NativeWhisperxConfig,
125) -> Result<(), NativeWhisperxError> {
126 Ok(())
127}
128
129#[cfg(feature = "diarization")]
130fn is_transient_speaker_label(label: &str) -> bool {
131 label
132 .strip_prefix("speaker_")
133 .is_some_and(|suffix| !suffix.is_empty() && suffix.chars().all(|ch| ch.is_ascii_digit()))
134}
135
136pub fn correct_speaker(
137 request: SpeakerCorrectionRequest,
138) -> Result<SpeakerCorrectionReport, NativeWhisperxError> {
139 let current_dir = std::env::current_dir()?;
140 let resolved = crate::resolve_speaker_directory(&request.speaker_directory, ¤t_dir)?;
141 let ranges =
142 speaker_correction_ranges(&request.transcript, &request.from_speaker, &request.ranges)?;
143 let audio = LoadedAudio::mono_16khz_from_source(&map_input_source(&request.audio))
144 .map_err(|error| NativeWhisperxError::Transcription(error.to_string()))?;
145 let embedding = speaker_embedding_for_ranges(&audio, &ranges)?;
146 let library = crate::speaker_directory::load_speaker_library_if_present(&resolved.path)?
147 .unwrap_or_default();
148 let profile_id = request
149 .speaker_id
150 .clone()
151 .unwrap_or_else(|| slug_speaker_id(&request.to_label));
152 let mut metadata = BTreeMap::new();
153 metadata.insert("status".to_string(), "confirmed".to_string());
154 metadata.insert("correctedFrom".to_string(), request.from_speaker.clone());
155 metadata.insert("updatedAt".to_string(), current_unix_seconds_string());
156 let (library, updated_existing_profile) =
157 crate::speaker_directory::upsert_speaker_profile_embedding(
158 &library,
159 &profile_id,
160 &request.to_label,
161 embedding,
162 metadata,
163 )?;
164 crate::speaker_directory::save_speaker_library(&resolved.path, &library)?;
165
166 let mut transcript = request.transcript;
167 replace_speaker_labels(
168 &mut transcript,
169 &request.from_speaker,
170 &request.to_label,
171 &request.ranges,
172 );
173 let response = speaker_correction_response(transcript.clone());
174 let output_files = write_outputs(&response, &request.output)?
175 .into_iter()
176 .map(|file| file.path)
177 .collect();
178
179 Ok(SpeakerCorrectionReport {
180 transcript,
181 speaker_directory_path: resolved.path,
182 profile_id,
183 label: request.to_label,
184 corrected_from: request.from_speaker,
185 enrolled_seconds: ranges
186 .iter()
187 .map(|range| range.end_seconds - range.start_seconds)
188 .sum(),
189 updated_existing_profile,
190 output_files,
191 })
192}
193
194fn speaker_correction_response(transcript: TranscriptionContract) -> TranscriptionPipelineResponse {
195 TranscriptionPipelineResponse {
196 accepted: true,
197 operation: "audio.transcription.correctSpeakers".to_string(),
198 provider: "native-whisperx".to_string(),
199 model_id: "speaker-correction".to_string(),
200 transcript,
201 vad_segments: Vec::new(),
202 alignment: None,
203 diarization: None,
204 artifacts: Vec::new(),
205 diagnostics: Vec::new(),
206 }
207}
208
209fn speaker_correction_ranges(
210 transcript: &TranscriptionContract,
211 from_speaker: &str,
212 filters: &[crate::SpeakerCorrectionRange],
213) -> Result<Vec<crate::SpeakerCorrectionRange>, NativeWhisperxError> {
214 let mut ranges = Vec::new();
215 for filter in filters {
216 validate_speaker_correction_range(*filter)?;
217 }
218 for segment in &transcript.segments {
219 if segment.speaker.as_deref() != Some(from_speaker) {
220 continue;
221 }
222 let Some((start_seconds, end_seconds)) = segment.start_seconds.zip(segment.end_seconds)
223 else {
224 continue;
225 };
226 let segment_range = crate::SpeakerCorrectionRange {
227 start_seconds,
228 end_seconds,
229 };
230 validate_speaker_correction_range(segment_range)?;
231 if filters.is_empty()
232 || filters
233 .iter()
234 .any(|filter| filter.overlaps(start_seconds, end_seconds))
235 {
236 ranges.push(segment_range);
237 }
238 }
239 if ranges.is_empty() {
240 return Err(NativeWhisperxError::InvalidConfig(format!(
241 "speaker correction found no timed transcript segments for speaker `{from_speaker}`"
242 )));
243 }
244 Ok(ranges)
245}
246
247fn validate_speaker_correction_range(
248 range: crate::SpeakerCorrectionRange,
249) -> Result<(), NativeWhisperxError> {
250 if !range.start_seconds.is_finite()
251 || !range.end_seconds.is_finite()
252 || range.start_seconds < 0.0
253 || range.end_seconds <= range.start_seconds
254 {
255 return Err(NativeWhisperxError::InvalidConfig(
256 "speaker correction ranges must be finite, non-negative, and have positive duration"
257 .to_string(),
258 ));
259 }
260 Ok(())
261}
262
263fn speaker_embedding_for_ranges(
264 audio: &LoadedAudio,
265 ranges: &[crate::SpeakerCorrectionRange],
266) -> Result<SpeakerEmbedding, NativeWhisperxError> {
267 let mut samples = Vec::new();
268 for range in ranges {
269 validate_speaker_correction_range(*range)?;
270 let start = ((range.start_seconds * audio.sample_rate as f64).floor() as usize)
271 .min(audio.samples.len());
272 let end = ((range.end_seconds * audio.sample_rate as f64).ceil() as usize)
273 .min(audio.samples.len());
274 if end > start {
275 samples.extend_from_slice(&audio.samples[start..end]);
276 }
277 }
278 if samples.is_empty() {
279 return Err(NativeWhisperxError::InvalidConfig(
280 "speaker correction did not select any audio samples".to_string(),
281 ));
282 }
283 let speaker_audio = SpeakerAudio::mono(&samples, audio.sample_rate).map_err(|error| {
284 NativeWhisperxError::Transcription(format!("speaker correction audio invalid: {error}"))
285 })?;
286 let mut embedder = SpectralSpeakerEmbedder::default();
287 embedder
288 .embed_speaker(&speaker_audio)
289 .map_err(|error| NativeWhisperxError::Transcription(error.to_string()))
290}
291
292fn replace_speaker_labels(
293 transcript: &mut TranscriptionContract,
294 from_speaker: &str,
295 to_label: &str,
296 filters: &[crate::SpeakerCorrectionRange],
297) {
298 for segment in &mut transcript.segments {
299 if segment.speaker.as_deref() != Some(from_speaker) {
300 continue;
301 }
302 let selected = if filters.is_empty() {
303 true
304 } else {
305 segment
306 .start_seconds
307 .zip(segment.end_seconds)
308 .is_some_and(|(start, end)| {
309 filters.iter().any(|filter| filter.overlaps(start, end))
310 })
311 };
312 if !selected {
313 continue;
314 }
315 segment.speaker = Some(to_label.to_string());
316 for word in &mut segment.words {
317 if word.speaker.as_deref() == Some(from_speaker) {
318 word.speaker = Some(to_label.to_string());
319 }
320 }
321 }
322}
323
324fn slug_speaker_id(value: &str) -> String {
325 let slug = value
326 .chars()
327 .map(|character| {
328 if character.is_ascii_alphanumeric() {
329 character.to_ascii_lowercase()
330 } else {
331 '-'
332 }
333 })
334 .collect::<String>()
335 .split('-')
336 .filter(|part| !part.is_empty())
337 .collect::<Vec<_>>()
338 .join("-");
339 if slug.is_empty() {
340 "speaker".to_string()
341 } else {
342 slug
343 }
344}
345
346fn current_unix_seconds() -> u64 {
347 SystemTime::now()
348 .duration_since(UNIX_EPOCH)
349 .map(|duration| duration.as_secs())
350 .unwrap_or_default()
351}
352
353fn current_unix_seconds_string() -> String {
354 current_unix_seconds().to_string()
355}
356
357#[cfg(feature = "diarization")]
358#[allow(dead_code)]
359pub(crate) fn read_whisperx_directory_state_for_ui(
360 path: &Path,
361) -> Result<crate::SpeakerDirectoryState, NativeWhisperxError> {
362 let resolved = crate::speaker_directory::read_speaker_directory_state(
363 &crate::speaker_directory::ResolvedSpeakerDirectory {
364 path: path.to_path_buf(),
365 scope: crate::speaker_directory::ResolvedSpeakerDirectoryScope::Local,
366 },
367 )?;
368 Ok(resolved)
369}
370
371#[cfg(feature = "diarization")]
372#[allow(dead_code)]
373pub(crate) fn ensure_output_format_dir(path: &Path) -> std::io::Result<()> {
374 std::fs::create_dir_all(path)
375}
376
377impl From<&NativeWhisperxConfig> for OutputConfig {
378 fn from(_value: &NativeWhisperxConfig) -> Self {
379 OutputConfig::default()
380 }
381}