Allow Faster-whisper evals to run on multiple GPUs in parallel (#15)

Co-authored-by: yoad <[email protected]>

Author: yoadsn

engines/faster_whisper_engine.py

@@ -18,14 +18,35 @@ def transcribe(model, entry: Dict[str, Any]) -> str:
     return " ".join(texts)
 
 
+def get_device_and_index(device: str) -> tuple[str, int | None]:
+    if len(device.split(":")) == 2:
+        device, device_index = device.split(":")
+        device_index = int(device_index)
+        return device, device_index
+
+    return device, None
+
+
 def create_app(**kwargs) -> Callable:
     model_path = kwargs.get("model_path")
     device: str = kwargs.get("device", "auto")
     device_index = None
-    if len(device.split(":")) == 2:
-        device, device_index = device.split(":")
-        device_index = int(device_index)
 
+    if len(device.split(",")) > 1:
+        device_indexes = []
+        base_device = None
+        for device_instance in device.split(","):
+            device, device_index = get_device_and_index(device_instance)
+            base_device = base_device or device
+            if base_device != device:
+                raise ValueError("Multiple devices must be instances of the same base device (e.g cuda:0, cuda:1 etc.)")
+            device_indexes.append(device_index)
+        device = base_device
+        device_index = device_indexes
+    else:
+        device, device_index = get_device_and_index(device)
+
+    print(f'Loading faster-whisper model: {model_path} on {device} with index: {device_index or "auto"}')
     model = faster_whisper.WhisperModel(model_path, device=device, device_index=device_index)
 
     def transcribe_fn(entry):