Merge pull request #63 from pragmatrix/google-diarization

Support speaker diarization for Google

Author: pragmatrix

examples/transcribe.rs

@@ -119,7 +119,8 @@ async fn recognize_from_microphone(
 ) -> Result<()> {
     // Keep an output sink alive so Bluetooth headsets can switch to a bidirectional profile.
     let _output_sink = match DeviceSinkBuilder::open_default_sink() {
-        Ok(sink) => {
+        Ok(mut sink) => {
+            sink.log_on_drop(false);
             println!("Opened default output sink for headset profile");
             Some(sink)
         }
@@ -279,9 +280,6 @@ 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());
@@ -297,11 +295,15 @@ async fn start_conversation(
                 None => google_transcribe::transcribe::Region::default(),
             };
 
+            // Check model/language/region feature support (including diarization):
+            // https://docs.cloud.google.com/speech-to-text/docs/speech-to-text-supported-languages
+
             let params = google_transcribe::transcribe::Params {
                 model: model.map(str::to_owned).unwrap_or_else(|| {
                     env::var("GOOGLE_TRANSCRIBE_MODEL").unwrap_or_else(|_| "latest_long".to_owned())
                 }),
                 language: languages.join_csv(),
+                diarization,
                 region,
             };
             GoogleTranscribe.conversation(params, conversation).await

justfile

@@ -22,3 +22,6 @@ transcribe-google-latest-short:
 transcribe-google-latest-long:
     cargo run --example transcribe -- google --language de-DE --model latest_long --region eu
 
+transcribe-google-diarization:
+    cargo run --example transcribe -- google --diarization --language de-DE --model chirp_3 --region eu
+

services/google-transcribe/src/client.rs

@@ -14,8 +14,9 @@ use google_cloud_token::TokenSource;
 use googleapis_tonic_google_cloud_speech_v2::google::cloud::speech::v2::recognition_config::DecodingConfig;
 use googleapis_tonic_google_cloud_speech_v2::google::cloud::speech::v2::speech_client::SpeechClient;
 use googleapis_tonic_google_cloud_speech_v2::google::cloud::speech::v2::{
-    ExplicitDecodingConfig, RecognitionConfig, StreamingRecognitionConfig,
+    ExplicitDecodingConfig, RecognitionConfig, RecognitionFeatures, StreamingRecognitionConfig,
     StreamingRecognitionFeatures, StreamingRecognizeRequest, StreamingRecognizeResponse,
+    SpeakerDiarizationConfig,
     explicit_decoding_config,
 };
 use googleapis_tonic_google_cloud_speech_v2::google::cloud::speech::v2::streaming_recognize_request::StreamingRequest;
@@ -165,6 +166,7 @@ impl TranscribeClient {
         &mut self,
         model: &str,
         language_codes: &[String],
+        diarization: bool,
         interim_results: bool,
         audio_format: AudioFormat,
         mut audio_receiver: UnboundedReceiver<Vec<i16>>,
@@ -180,7 +182,13 @@ impl TranscribeClient {
             // TODO: configure
             model: model.into(),
             language_codes: language_codes.to_vec(),
-            features: None,
+            features: diarization.then_some(RecognitionFeatures {
+                diarization_config: Some(SpeakerDiarizationConfig {
+                    min_speaker_count: 0,
+                    max_speaker_count: 0,
+                }),
+                ..Default::default()
+            }),
             adaptation: None,
             transcript_normalization: None,
             denoiser_config: None,
@@ -206,6 +214,7 @@ impl TranscribeClient {
             recognizer = %recognizer,
             model = %model,
             language_codes = ?language_codes,
+            diarization,
             interim_results,
             "Starting Google streaming_recognize"
         );

services/google-transcribe/src/transcribe.rs

@@ -2,9 +2,10 @@ use anyhow::{Context, Result};
 use async_trait::async_trait;
 use futures::{Stream, StreamExt};
 use googleapis_tonic_google_cloud_speech_v2::google::cloud::speech::v2::{
-    StreamingRecognizeResponse, streaming_recognize_response::SpeechEventType,
+    StreamingRecognizeResponse, WordInfo, streaming_recognize_response::SpeechEventType,
 };
 use serde::Deserialize;
+use std::collections::HashMap;
 use tokio::sync::mpsc::UnboundedReceiver;
 use tonic::Code;
 
@@ -23,6 +24,8 @@ pub struct Params {
     pub model: String,
     pub language: String,
     #[serde(default)]
+    pub diarization: bool,
+    #[serde(default)]
     pub region: Region,
 }
 
@@ -154,6 +157,7 @@ async fn transcribe_and_process_stream_session(
         .transcribe(
             &params.model,
             languages,
+            params.diarization,
             interim_results,
             audio_format,
             audio_receiver,
@@ -221,7 +225,12 @@ where
                 let language = include_detected_language
                     .then(|| one.language_code.trim().to_owned())
                     .filter(|x| !x.is_empty());
-                text_output.final_text(alternative.transcript.trim().to_owned(), language)?;
+                let speaker = speaker_with_max_assigned_characters(&alternative.words);
+                text_output.final_text(
+                    alternative.transcript.trim().to_owned(),
+                    language,
+                    speaker,
+                )?;
             }
             [_, ..] => {
                 let interim_text = response
@@ -321,6 +330,25 @@ fn should_restart_for_stream_limit(code: Code, message: &str) -> bool {
     code == Code::Aborted && message.contains("max duration of 5 minutes reached for stream")
 }
 
+fn speaker_with_max_assigned_characters(words: &[WordInfo]) -> Option<String> {
+    let mut char_count_by_speaker = HashMap::<&str, usize>::new();
+
+    for word in words {
+        let speaker = word.speaker_label.trim();
+        if speaker.is_empty() {
+            continue;
+        }
+
+        let spoken_chars = word.word.chars().count();
+        *char_count_by_speaker.entry(speaker).or_default() += spoken_chars;
+    }
+
+    char_count_by_speaker
+        .into_iter()
+        .max_by_key(|(_, char_count)| *char_count)
+        .map(|(speaker, _)| speaker.to_owned())
+}
+
 // Keeps interim/final emission behavior in one place across all exit paths.
 // If a session ends without any final result, we promote the last interim text
 // to final in Drop so callers still receive a terminal text event.
@@ -338,8 +366,13 @@ impl<'a> StreamTextOutput<'a> {
         }
     }
 
-    fn final_text(&mut self, text: String, language: Option<String>) -> Result<()> {
-        self.output.text(true, text, language, None)?;
+    fn final_text(
+        &mut self,
+        text: String,
+        language: Option<String>,
+        speaker: Option<String>,
+    ) -> Result<()> {
+        self.output.text(true, text, language, speaker)?;
         self.pending_interim_text = None;
         Ok(())
     }