Add --audio-folder cli argument

argmaxinc · Apr 18, 2024 · 076b670 · 076b670
1 parent 1d376e3
commit 076b670
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/Sources/WhisperKitCLI/CLIArguments.swift b/Sources/WhisperKitCLI/CLIArguments.swift
@@ -6,6 +6,9 @@ import ArgumentParser
 struct CLIArguments: ParsableArguments {
     @Option(help: "Paths to audio files")
     var audioPath = [String]()
+
+    @Option(help: "Path to a folder containing audio files")
+    var audioFolder: String?
 
     @Option(help: "Path of model files")
     var modelPath: String?

diff --git a/Sources/WhisperKitCLI/Transcribe.swift b/Sources/WhisperKitCLI/Transcribe.swift
@@ -21,6 +21,21 @@ struct Transcribe: AsyncParsableCommand {
                 throw ValidationError("Invalid language code \"\(language)\". Supported languages: \(Constants.languages.values)")
             }
         }
+
+        if cliArguments.audioPath.isEmpty {
+            guard let audioFolder = cliArguments.audioFolder else {
+                throw ValidationError("Either audioPath or audioFolder must be provided.")
+            }
+            let fileManager = FileManager.default
+            let audioExtensions = ["mp3", "wav", "m4a", "flac", "aiff", "aac", "ogg"]
+            let audioFiles = try fileManager.contentsOfDirectory(atPath: audioFolder)
+                .filter { fileName in
+                    let fileExtension = fileName.lowercased().components(separatedBy: ".").last
+                    return audioExtensions.contains(fileExtension ?? "")
+                }
+
+            cliArguments.audioPath = audioFiles.map { audioFolder + "/" + $0 }
+        }
     }
 
     mutating func run() async throws {