Make Whisper runnable during unit tests

Telosnex · Feb 8, 2024 · ac30c6f · ac30c6f
1 parent 6dd8156
commit ac30c6f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/lib/models/whisper/whisper_isolate.dart b/lib/models/whisper/whisper_isolate.dart
@@ -12,13 +12,15 @@ class WhisperIsolateMessage {
   final SendPort replyPort;
   final String modelPath;
   final String? ortDylibPathOverride;
+  final String? ortExtensionsDylibPathOverride;
   final List<int> audioBytes;
 
   WhisperIsolateMessage({
     required this.replyPort,
     required this.modelPath,
     required this.audioBytes,
     this.ortDylibPathOverride,
+    this.ortExtensionsDylibPathOverride,
   });
 }
 
@@ -36,6 +38,10 @@ void whisperIsolateEntryPoint(SendPort mainSendPort) {
         if (message.ortDylibPathOverride != null) {
           fonnxOrtDylibPathOverride = message.ortDylibPathOverride;
         }
+        if (message.ortExtensionsDylibPathOverride != null) {
+          fonnxOrtExtensionsDylibPathOverride =
+              message.ortExtensionsDylibPathOverride;
+        }
         // Lazily create the Ort session if it's not already done.
         ortSessionObjects ??=
             createOrtSession(message.modelPath, includeOnnxExtensionsOps: true);
@@ -109,6 +115,7 @@ class WhisperIsolateManager {
     String modelPath,
     List<int> audioBytes, {
     String? ortDylibPathOverride,
+    String? ortExtensionsDylibPathOverride,
   }) async {
     await start();
     final response = ReceivePort();
@@ -117,6 +124,7 @@ class WhisperIsolateManager {
       modelPath: modelPath,
       audioBytes: audioBytes,
       ortDylibPathOverride: ortDylibPathOverride,
+      ortExtensionsDylibPathOverride: ortExtensionsDylibPathOverride,
     );
 
     _sendPort!.send(message);

diff --git a/lib/models/whisper/whisper_native.dart b/lib/models/whisper/whisper_native.dart
@@ -20,7 +20,12 @@ class WhisperNative implements Whisper {
   Future<String> doInference(List<int> bytes) async {
     await _whisperIsolateManager.start();
     if (!kIsWeb && Platform.environment['FLUTTER_TEST'] == 'true') {
-      return _whisperIsolateManager.sendInference(modelPath, bytes);
+      return _whisperIsolateManager.sendInference(
+        modelPath,
+        bytes,
+        ortDylibPathOverride: fonnxOrtDylibPathOverride,
+        ortExtensionsDylibPathOverride: fonnxOrtExtensionsDylibPathOverride,
+      );
     }
 
     switch (defaultTargetPlatform) {