feat(am2): add file transcription support using WhisperKitPro (#2158)

* feat(am2): add file transcription support using WhisperKitPro

- Add Swift bindings for WhisperKitPro file transcription
- Add am2_transcribe_init, am2_transcribe_is_ready, am2_transcribe_file functions
- Add Rust wrapper module for transcription functionality
- Add test for transcription initialization

Co-Authored-By: yujonglee <[email protected]>

* refactor(am2): make transcribe::is_ready() consistent with vad::is_ready()

- Remove unnecessary FFI call in is_ready() to match vad module pattern
- Remove unused am2_transcribe_is_ready FFI declaration

Addresses code review feedback about design inconsistency.

Co-Authored-By: yujonglee <[email protected]>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: yujonglee <[email protected]>

Author: devin-ai-integration[bot]

crates/am2/src/lib.rs

@@ -6,6 +6,8 @@ swift!(fn initialize_am2_sdk(api_key: &SRString));
 swift!(fn am2_vad_init() -> bool);
 swift!(fn am2_vad_detect(samples_ptr: *const f32, samples_len: i64) -> SRObject<VadResultArray>);
 swift!(fn am2_vad_index_to_seconds(index: i64) -> f32);
+swift!(fn am2_transcribe_init(model: &SRString) -> bool);
+swift!(fn am2_transcribe_file(audio_path: &SRString) -> SRObject<TranscribeResultFFI>);
 
 static SDK_INITIALIZED: OnceLock<()> = OnceLock::new();
 
@@ -14,6 +16,12 @@ pub struct VadResultArray {
     pub data: SRArray<bool>,
 }
 
+#[repr(C)]
+pub struct TranscribeResultFFI {
+    pub text: SRString,
+    pub success: bool,
+}
+
 pub fn init() {
     SDK_INITIALIZED.get_or_init(|| {
         let api_key = std::env::var("AM_API_KEY").unwrap_or_default();
@@ -88,6 +96,40 @@ pub mod vad {
     }
 }
 
+pub mod transcribe {
+    use std::sync::OnceLock;
+
+    use super::*;
+
+    static TRANSCRIBE_INITIALIZED: OnceLock<bool> = OnceLock::new();
+
+    pub fn init(model: &str) -> bool {
+        *TRANSCRIBE_INITIALIZED.get_or_init(|| {
+            let model = SRString::from(model);
+            unsafe { am2_transcribe_init(&model) }
+        })
+    }
+
+    pub fn is_ready() -> bool {
+        TRANSCRIBE_INITIALIZED.get().copied().unwrap_or(false)
+    }
+
+    #[derive(Debug, Clone)]
+    pub struct TranscribeResult {
+        pub text: String,
+        pub success: bool,
+    }
+
+    pub fn transcribe_file(audio_path: &str) -> TranscribeResult {
+        let audio_path = SRString::from(audio_path);
+        let result = unsafe { am2_transcribe_file(&audio_path) };
+        TranscribeResult {
+            text: result.text.to_string(),
+            success: result.success,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -104,4 +146,11 @@ mod tests {
         assert!(vad::init());
         assert!(vad::is_ready());
     }
+
+    #[test]
+    fn test_am2_transcribe_init() {
+        init();
+        assert!(transcribe::init("large-v3-v20240930_626MB"));
+        assert!(transcribe::is_ready());
+    }
 }

crates/am2/swift-lib/src/lib.swift

@@ -3,6 +3,7 @@ import Foundation
 import SwiftRs
 
 private var vadInstance: VoiceActivityDetector?
+private var whisperKitProInstance: WhisperKitPro?
 
 @_cdecl("initialize_am2_sdk")
 public func initialize_am2_sdk(apiKey: SRString) {
@@ -68,3 +69,67 @@ public func am2_vad_index_to_seconds(index: Int) -> Float {
   }
   return vad.voiceActivityIndexToSeconds(index)
 }
+
+public class TranscribeResult: NSObject {
+  var text: SRString
+  var success: Bool
+
+  init(text: String, success: Bool) {
+    self.text = SRString(text)
+    self.success = success
+  }
+}
+
+@_cdecl("am2_transcribe_init")
+public func am2_transcribe_init(model: SRString) -> Bool {
+  let semaphore = DispatchSemaphore(value: 0)
+  var success = false
+
+  Task {
+    do {
+      let modelName = model.toString()
+      let config = WhisperKitProConfig(model: modelName)
+      whisperKitProInstance = try await WhisperKitPro(config)
+      success = true
+    } catch {
+      success = false
+    }
+    semaphore.signal()
+  }
+
+  semaphore.wait()
+  return success
+}
+
+@_cdecl("am2_transcribe_is_ready")
+public func am2_transcribe_is_ready() -> Bool {
+  return whisperKitProInstance != nil
+}
+
+@_cdecl("am2_transcribe_file")
+public func am2_transcribe_file(audioPath: SRString) -> TranscribeResult {
+  let semaphore = DispatchSemaphore(value: 0)
+  var resultText = ""
+  var success = false
+
+  Task {
+    guard let whisperKit = whisperKitProInstance else {
+      semaphore.signal()
+      return
+    }
+
+    do {
+      let path = audioPath.toString()
+      let results = try await whisperKit.transcribe(audioPath: path)
+      resultText = WhisperKitProUtils.mergeTranscriptionResults(results).text
+      success = true
+    } catch {
+      resultText = ""
+      success = false
+    }
+    semaphore.signal()
+  }
+
+  semaphore.wait()
+  return TranscribeResult(text: resultText, success: success)
+}