Merge pull request #60 from pragmatrix/azure-diarization

Add support for speaker diarization in azure-transcribe

Author: pragmatrix

core/src/conversation.rs

@@ -227,11 +227,18 @@ impl ConversationOutput {
         self.post(Output::ClearAudio)
     }
 
-    pub fn text(&self, is_final: bool, text: String, language: Option<String>) -> Result<()> {
+    pub fn text(
+        &self,
+        is_final: bool,
+        text: String,
+        language: Option<String>,
+        speaker: Option<String>,
+    ) -> Result<()> {
         self.post(Output::Text {
             is_final,
             text,
             language,
+            speaker,
         })
     }
 
@@ -323,6 +330,7 @@ pub enum Output {
         is_final: bool,
         text: String,
         language: Option<String>,
+        speaker: Option<String>,
     },
     RequestCompleted {
         request_id: Option<RequestId>,

examples/azure-transcribe.rs

@@ -22,16 +22,26 @@ async fn main() -> Result<()> {
     dotenvy::dotenv_override()?;
     tracing_subscriber::fmt::init();
 
-    let mut args = env::args();
-    match args.len() {
-        1 => recognize_from_microphone().await?,
-        2 => recognize_from_wav(Path::new(&args.nth(1).unwrap())).await?,
-        _ => bail!("Invalid number of arguments, expect zero or one"),
+    let mut input_file: Option<String> = None;
+    let mut diarization = false;
+    for arg in env::args().skip(1) {
+        if arg == "--diarization" {
+            diarization = true;
+        } else if input_file.is_none() {
+            input_file = Some(arg);
+        } else {
+            bail!("Invalid arguments. Usage: azure-transcribe [--diarization] [wav-file]");
+        }
+    }
+
+    match input_file {
+        None => recognize_from_microphone(diarization).await?,
+        Some(path) => recognize_from_wav(Path::new(&path), diarization).await?,
     }
     Ok(())
 }
 
-async fn recognize_from_wav(file: &Path) -> Result<()> {
+async fn recognize_from_wav(file: &Path, diarization: bool) -> Result<()> {
     // For now we always convert to 16khz single channel (this is what we use internally for
     // testing).
     let format = AudioFormat {
@@ -50,10 +60,10 @@ async fn recognize_from_wav(file: &Path) -> Result<()> {
         producer.produce(frame)?;
     }
 
-    recognize(format, input_consumer).await
+    recognize(format, input_consumer, diarization).await
 }
 
-async fn recognize_from_microphone() -> Result<()> {
+async fn recognize_from_microphone(diarization: bool) -> Result<()> {
     // Keep an output sink alive so Bluetooth headsets (e.g. AirPods) can switch to a
     // bidirectional profile. Without this, some devices report an input stream of zeros.
     let _output_sink = match DeviceSinkBuilder::open_default_sink() {
@@ -105,17 +115,22 @@ async fn recognize_from_microphone() -> Result<()> {
 
     stream.play().expect("Failed to play stream");
 
-    recognize(format, input_consumer).await
+    recognize(format, input_consumer, diarization).await
 }
 
-async fn recognize(format: AudioFormat, mut input_consumer: AudioConsumer) -> Result<()> {
+async fn recognize(
+    format: AudioFormat,
+    mut input_consumer: AudioConsumer,
+    diarization: bool,
+) -> Result<()> {
     // TODO: clarify how to access configurations.
     let params = azure::transcribe::Params {
         host: None,
         region: Some(env::var("AZURE_REGION").expect("AZURE_REGION undefined")),
         subscription_key: env::var("AZURE_SUBSCRIPTION_KEY")
             .expect("AZURE_SUBSCRIPTION_KEY undefined"),
         language: LANGUAGE.into(),
+        diarization,
         speech_gate: false,
     };
 

examples/azure-translate.rs

@@ -182,8 +182,9 @@ async fn setup_audio_playback(
                     is_final,
                     text,
                     language,
+                    speaker,
                 } => {
-                    println!("Text: {text}, final: {is_final}, language: {language:?}")
+                    println!("Text ({is_final}, {language:?}, speaker: {speaker:?}): {text}");
                 }
                 Output::RequestCompleted { .. } => {}
                 Output::ClearAudio => {

examples/transcribe.rs

@@ -29,6 +29,8 @@ struct Args {
     model: Option<String>,
     #[arg(long)]
     region: Option<String>,
+    #[arg(long)]
+    diarization: bool,
 }
 
 #[derive(Debug, Clone, Copy, ValueEnum)]
@@ -52,10 +54,15 @@ async fn main() -> Result<()> {
     let languages = Languages::new(args.language)?;
     let model = args.model.as_deref();
     let region = args.region.as_deref();
+    let diarization = args.diarization;
 
     match args.input.as_deref() {
-        Some(path) => recognize_from_wav(args.provider, path, &languages, model, region).await?,
-        None => recognize_from_microphone(args.provider, &languages, model, region).await?,
+        Some(path) => {
+            recognize_from_wav(args.provider, path, &languages, model, region, diarization).await?
+        }
+        None => {
+            recognize_from_microphone(args.provider, &languages, model, region, diarization).await?
+        }
     }
 
     Ok(())
@@ -67,6 +74,7 @@ async fn recognize_from_wav(
     languages: &Languages,
     model: Option<&str>,
     region: Option<&str>,
+    diarization: bool,
 ) -> Result<()> {
     let format = AudioFormat {
         channels: 1,
@@ -83,14 +91,24 @@ async fn recognize_from_wav(
         producer.produce(frame)?;
     }
 
-    recognize(provider, format, input_consumer, languages, model, region).await
+    recognize(
+        provider,
+        format,
+        input_consumer,
+        languages,
+        model,
+        region,
+        diarization,
+    )
+    .await
 }
 
 async fn recognize_from_microphone(
     provider: Provider,
     languages: &Languages,
     model: Option<&str>,
     region: Option<&str>,
+    diarization: bool,
 ) -> Result<()> {
     // Keep an output sink alive so Bluetooth headsets can switch to a bidirectional profile.
     let _output_sink = match DeviceSinkBuilder::open_default_sink() {
@@ -131,7 +149,16 @@ async fn recognize_from_microphone(
 
     stream.play().expect("Failed to play stream");
 
-    recognize(provider, format, input_consumer, languages, model, region).await
+    recognize(
+        provider,
+        format,
+        input_consumer,
+        languages,
+        model,
+        region,
+        diarization,
+    )
+    .await
 }
 
 async fn recognize(
@@ -141,6 +168,7 @@ async fn recognize(
     languages: &Languages,
     model: Option<&str>,
     region: Option<&str>,
+    diarization: bool,
 ) -> Result<()> {
     let (output_producer, mut output_consumer) = unbounded_channel();
     let (conversation_input_producer, conversation_input_consumer) = channel(16_384);
@@ -150,6 +178,7 @@ async fn recognize(
         languages,
         model,
         region,
+        diarization,
         Conversation::new(
             InputModality::Audio { format },
             [OutputModality::Text, OutputModality::InterimText],
@@ -190,6 +219,7 @@ async fn start_conversation(
     languages: &Languages,
     model: Option<&str>,
     region: Option<&str>,
+    diarization: bool,
     conversation: Conversation,
 ) -> Result<()> {
     match provider {
@@ -203,11 +233,15 @@ async fn start_conversation(
                 subscription_key: env::var("AZURE_SUBSCRIPTION_KEY")
                     .expect("AZURE_SUBSCRIPTION_KEY undefined"),
                 language: languages.join_csv(),
+                diarization,
                 speech_gate: false,
             };
             AzureTranscribe.conversation(params, conversation).await
         }
         Provider::Elevenlabs => {
+            if diarization {
+                bail!("--diarization is only supported for the azure provider");
+            }
             if region.is_some() {
                 bail!("--region is only supported for the google provider");
             }
@@ -234,6 +268,9 @@ async fn start_conversation(
                 .await
         }
         Provider::Google => {
+            if diarization {
+                bail!("--diarization is only supported for the azure provider");
+            }
             let region = region
                 .map(str::to_owned)
                 .or_else(|| env::var("GOOGLE_TRANSCRIBE_REGION").ok());
@@ -259,6 +296,9 @@ async fn start_conversation(
             GoogleTranscribe.conversation(params, conversation).await
         }
         Provider::Aristech => {
+            if diarization {
+                bail!("--diarization is only supported for the azure provider");
+            }
             if region.is_some() {
                 bail!("--region is only supported for the google provider");
             }

external/azure-speech-sdk-rs

@@ -7,12 +7,18 @@ knife-release:
 transcribe-azure-de-en:
     cargo run --example transcribe -- azure --language de-DE,en-US
 
+transcribe-azure-diarization:
+    cargo run --example transcribe -- azure --diarization
+
+transcribe-azure-diarization-de:
+    cargo run --example transcribe -- azure --diarization --language de-DE
+
 transcribe-google-de-en:
     cargo run --example transcribe -- google --language de-DE,en-US --model chirp_3 --region eu
 
-
 transcribe-google-latest-short:
     cargo run --example transcribe -- google --language de-DE --model latest_short --region eu
 
 transcribe-google-latest-long:
     cargo run --example transcribe -- google --language de-DE --model latest_long --region eu
+

justfile

@@ -152,7 +152,7 @@ impl Service for AristechTranscribe {
 
                 // Instead of processing all alternatives, just take the first one
                 if let Some(alternative) = chunk.alternatives.into_iter().next() {
-                    output.text(is_final, alternative.text, None)?;
+                    output.text(is_final, alternative.text, None, None)?;
                 }
             }
         }

services/aristech/src/transcribe.rs

@@ -22,6 +22,8 @@ pub struct Params {
     pub subscription_key: String,
     pub language: String,
     #[serde(default)]
+    pub diarization: bool,
+    #[serde(default)]
     pub speech_gate: bool,
 }
 
@@ -55,6 +57,12 @@ impl Service for AzureTranscribe {
             // Disable profanity filter.
             .set_profanity(recognizer::Profanity::Raw);
 
+        let config = if params.diarization {
+            config.enable_recognize_speaker()
+        } else {
+            config
+        };
+
         let config = if languages.len() == 1 {
             config.set_language(recognizer::Language::Custom(languages.first().clone()))
         } else {
@@ -145,7 +153,7 @@ fn output_recognized_text(
     let recognizer::Recognized {
         text,
         primary_language,
-        ..
+        speaker_id: speaker,
     } = recognized;
 
     let language = if include_detected_language {
@@ -154,5 +162,9 @@ fn output_recognized_text(
         None
     };
 
-    output.text(is_final, text, language)
+    // Azure frequently reports speaker as "Unknown" during interim recognition,
+    // so we only emit speaker information for final text events.
+    let speaker = if is_final { speaker } else { None };
+
+    output.text(is_final, text, language, speaker)
 }

services/azure/src/transcribe.rs

@@ -127,13 +127,13 @@ impl Service for AzureTranslate {
                 Event::StartDetected(_, _) => {}
                 Event::EndDetected(_, _) => {}
                 Event::Translating(_, text, _, _, _) => {
-                    if output_modalities.text {
-                        output.text(false, text, None)?;
+                    if output_modalities.interim_text {
+                        output.text(false, text, None, None)?;
                     }
                 }
                 Event::Translated(_, text, _, _, _) => {
-                    if output_modalities.interim_text {
-                        output.text(true, text, None)?;
+                    if output_modalities.text {
+                        output.text(true, text, None, None)?;
                     }
                 }
                 Event::TranslationSynthesis(_, samples) => {

services/azure/src/translate.rs

@@ -447,15 +447,15 @@ fn process_server_json(
         }
         "partial_transcript" => {
             let event: PartialTranscript = serde_json::from_value(envelope.payload)?;
-            output.text(false, event.text, None)
+            output.text(false, event.text, None, None)
         }
         "committed_transcript" => {
             if include_language_detection {
                 // Ignoring committed_transcript because include_language_detection=true; expecting committed_transcript_with_timestamps
                 return Ok(());
             }
             let event: CommittedTranscript = serde_json::from_value(envelope.payload)?;
-            output.text(true, event.text, None)
+            output.text(true, event.text, None, None)
         }
         "committed_transcript_with_timestamps" => {
             let event: CommittedTranscriptWithTimestamps =
@@ -482,7 +482,7 @@ fn process_server_json(
                     }
                 });
 
-            output.text(true, text, language_code)
+            output.text(true, text, language_code, None)
         }
         // Not in the official documentation, but this happens when the language code is invalid.
         "invalid_request" => {