1use std::sync::Arc;
4
5use anyhow::{Context, ensure};
6use chrono::{DateTime, Utc};
7use kcode_audio_speaker_review::{
8 CLASSIFIER_MODEL, CLASSIFIER_PROMPT_VERSION, CLASSIFIER_PROVIDER, CLASSIFIER_SCHEMA_VERSION,
9 CandidateMapping, ConfirmationState, ObservationKey, ParsedSpeaker, training_object_id,
10};
11pub(crate) use kcode_audio_speaker_review::{
12 CorrectionChunk, CorrectionObservation, CorrectionPacket, ParsedChunk,
13};
14use kcode_speaker_extract::ExtractionOutcome;
15use kcode_speaker_system::{Cohort, IdentifyEvidence, SpeechClassifier};
16use uuid::Uuid;
17
18const READ_ONLY_IDENTIFY_THRESHOLD: f64 = 1e308;
19
20#[derive(Clone)]
21pub(crate) struct ClassificationContext {
22 pub(crate) recording_id: Uuid,
23 pub(crate) user_id: String,
24 pub(crate) sha256: String,
25 pub(crate) original_filename: String,
26 pub(crate) size_bytes: u64,
27 pub(crate) recorded_at: DateTime<Utc>,
28 pub(crate) classifier: Arc<SpeechClassifier>,
29}
30
31pub(crate) fn parsed_chunk_from_extraction(
32 extraction: &ExtractionOutcome,
33) -> anyhow::Result<ParsedChunk> {
34 let (count, clip_valid, clip_validity_reason) = match extraction {
35 ExtractionOutcome::Scored(scored) if scored.additional_speakers.is_empty() => {
36 (scored.speakers.len(), true, None)
37 }
38 ExtractionOutcome::Scored(scored) => (
39 scored.speakers.len() + scored.additional_speakers.len(),
40 false,
41 Some("One or more speakers lacked a complete feature profile.".to_owned()),
42 ),
43 ExtractionOutcome::Unscorable {
44 reason,
45 additional_speakers,
46 } => (additional_speakers.len(), false, Some(reason.clone())),
47 };
48 let speakers = (0..count)
49 .map(|ordinal| {
50 let profile = match extraction {
51 ExtractionOutcome::Scored(scored) => scored
52 .speakers
53 .iter()
54 .find(|profile| usize::from(profile.speaker_ordinal) == ordinal),
55 ExtractionOutcome::Unscorable { .. } => None,
56 };
57 ParsedSpeaker {
58 local_label: format!("Speaker {}", ordinal + 1),
59 primary_language: profile.map(|value| value.primary_language.as_ref().to_owned()),
60 feature_row: profile.map(|value| value.features),
61 }
62 })
63 .collect();
64 Ok(ParsedChunk {
65 clip_valid,
66 clip_validity_reason,
67 speakers,
68 })
69}
70
71pub(crate) fn classify_speakers(
72 context: &ClassificationContext,
73 chunk_index: usize,
74 parsed: &ParsedChunk,
75) -> anyhow::Result<Vec<CorrectionObservation>> {
76 let mut observations = Vec::with_capacity(parsed.speakers.len());
77 for (ordinal, speaker) in parsed.speakers.iter().enumerate() {
78 let speaker_ordinal = u32::try_from(ordinal)
79 .context("chunk has more speakers than the key schema supports")?;
80 let candidate = match (
81 speaker.primary_language.as_deref(),
82 speaker.feature_row.as_ref(),
83 ) {
84 (Some(primary_language), Some(feature_row)) => {
85 let probe_key = ObservationKey {
86 object_id: probe_object_id(context.recording_id, chunk_index),
87 piece_index: speaker_ordinal,
88 };
89 let outcome = context
90 .classifier
91 .identify(
92 probe_key.clone(),
93 cohort(primary_language),
94 *feature_row,
95 READ_ONLY_IDENTIFY_THRESHOLD,
96 )
97 .with_context(|| {
98 format!(
99 "read-only identity scoring failed for chunk {chunk_index} speaker {}",
100 speaker.local_label
101 )
102 })?;
103 if outcome.speaker_id.is_some() {
104 context
105 .classifier
106 .delete(probe_key)
107 .context("removing an unexpectedly accepted read-only probe")?;
108 }
109 outcome.evidence.as_ref().map(candidate_mapping)
110 }
111 (None, None) => None,
112 _ => anyhow::bail!("speaker profile is only partially present"),
113 };
114 observations.push(CorrectionObservation {
115 local_label: speaker.local_label.clone(),
116 speaker_ordinal,
117 observation_key: ObservationKey {
118 object_id: training_object_id(context.recording_id, chunk_index),
119 piece_index: speaker_ordinal,
120 },
121 candidate,
122 resolution: None,
123 });
124 }
125 Ok(observations)
126}
127
128pub(crate) fn unclassified_observations(
129 recording_id: Uuid,
130 chunk_index: usize,
131 parsed: &ParsedChunk,
132) -> anyhow::Result<Vec<CorrectionObservation>> {
133 parsed
134 .speakers
135 .iter()
136 .enumerate()
137 .map(|(ordinal, speaker)| {
138 let speaker_ordinal = u32::try_from(ordinal)
139 .context("chunk has more speakers than the key schema supports")?;
140 Ok(CorrectionObservation {
141 local_label: speaker.local_label.clone(),
142 speaker_ordinal,
143 observation_key: ObservationKey {
144 object_id: training_object_id(recording_id, chunk_index),
145 piece_index: speaker_ordinal,
146 },
147 candidate: None,
148 resolution: None,
149 })
150 })
151 .collect()
152}
153
154pub(crate) fn build_packet(
155 context: &ClassificationContext,
156 chunks: Vec<CorrectionChunk>,
157) -> anyhow::Result<CorrectionPacket> {
158 ensure!(!chunks.is_empty(), "correction packet has no chunks");
159 let chunk_count = chunks.len();
160 ensure!(
161 chunks.iter().enumerate().all(|(index, chunk)| {
162 chunk.chunk_index == index
163 && chunk.chunk_count == chunk_count
164 && chunk.audio_end_ms > chunk.audio_start_ms
165 && chunk.observations.len() == chunk.parsed.speakers.len()
166 && !chunk.signed_off
167 }),
168 "correction packet chunks are not one complete chronological plan"
169 );
170 Ok(CorrectionPacket {
171 recording_id: context.recording_id,
172 user_id: context.user_id.clone(),
173 sha256: context.sha256.clone(),
174 original_filename: context.original_filename.clone(),
175 size_bytes: context.size_bytes,
176 recorded_at: context.recorded_at,
177 chunk_count,
178 chunks,
179 confirmation_state: ConfirmationState::Unconfirmed,
180 })
181}
182
183fn probe_object_id(recording_id: Uuid, chunk_index: usize) -> String {
184 format!("kcode-audio-ingress/probe/{recording_id}/chunk/{chunk_index}")
185}
186
187fn cohort(primary_language: &str) -> Cohort {
188 Cohort {
189 provider: CLASSIFIER_PROVIDER.into(),
190 model: CLASSIFIER_MODEL.into(),
191 prompt_version: CLASSIFIER_PROMPT_VERSION.into(),
192 schema_version: CLASSIFIER_SCHEMA_VERSION.into(),
193 primary_language: primary_language.into(),
194 }
195}
196
197fn candidate_mapping(evidence: &IdentifyEvidence) -> CandidateMapping {
198 CandidateMapping {
199 full_name: evidence.best.speaker_id.clone(),
200 score: -evidence.best.cost,
201 runner_up_score: evidence.runner_up.as_ref().map(|candidate| -candidate.cost),
202 }
203}