Merge pull request #72 from jkrukowski/tokenizer-download-path

argmaxinc · Mar 15, 2024 · 8588a38 · 8588a38
2 parents 0b78c52 + 4b83f89
commit 8588a38
Show file tree

Hide file tree

Showing 11 changed files with 100 additions and 37 deletions.
diff --git a/Makefile b/Makefile
@@ -79,7 +79,7 @@ build:
 
 build-cli:
 	@echo "Building WhisperKit CLI..."
-	@swift build -c release --product transcribe
+	@swift build -c release --product whisperkit-cli
 
 
 test:

diff --git a/Package.resolved b/Package.resolved
@@ -14,8 +14,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/huggingface/swift-transformers.git",
       "state" : {
-        "revision" : "564442fba36b0b694d730a62d0593e5f54043b55",
-        "version" : "0.1.2"
+        "revision" : "24605a8c0cc974bec5b94a6752eb687bae77db31",
+        "version" : "0.1.3"
       }
     }
   ],

diff --git a/Package.swift b/Package.swift
@@ -15,12 +15,12 @@ let package = Package(
             targets: ["WhisperKit"]
         ),
         .executable(
-            name: "transcribe",
+            name: "whisperkit-cli",
             targets: ["WhisperKitCLI"]
         ),
     ],
     dependencies: [
-        .package(url: "https://github.com/huggingface/swift-transformers.git", exact: "0.1.2"),
+        .package(url: "https://github.com/huggingface/swift-transformers.git", exact: "0.1.3"),
         .package(url: "https://github.com/apple/swift-argument-parser.git", exact: "1.3.0"),
     ],
     targets: [

diff --git a/README.md b/README.md
@@ -121,13 +121,13 @@ make download-models
 You can then run them via the CLI with:
 
 ```bash
-swift run transcribe --model-path "Models/whisperkit-coreml/openai_whisper-large-v3" --audio-path "path/to/your/audio.{wav,mp3,m4a,flac}" 
+swift run whisperkit-cli transcribe --model-path "Models/whisperkit-coreml/openai_whisper-large-v3" --audio-path "path/to/your/audio.{wav,mp3,m4a,flac}" 
 ```
 
 Which should print a transcription of the audio file. If you would like to stream the audio directly from a microphone, use:
 
 ```bash
-swift run transcribe --model-path "Models/whisperkit-coreml/openai_whisper-large-v3" --stream
+swift run whisperkit-cli transcribe --model-path "Models/whisperkit-coreml/openai_whisper-large-v3" --stream
 ```
 
 ## Contributing & Roadmap

diff --git a/Sources/WhisperKit/Core/Utils.swift b/Sources/WhisperKit/Core/Utils.swift
@@ -5,6 +5,7 @@ import AVFoundation
 import CoreML
 import Foundation
 import Tokenizers
+import Hub
 #if canImport(UIKit)
 import UIKit
 #elseif canImport(AppKit)
@@ -269,10 +270,13 @@ public func resolveAbsolutePath(_ inputPath: String) -> String {
     return inputPath
 }
 
-func loadTokenizer(for pretrained: ModelVariant) async throws -> Tokenizer {
-    // TODO: Cache tokenizer config to avoid repeated network requests
+func loadTokenizer(
+    for pretrained: ModelVariant,
+    tokenizerFolder: URL? = nil
+) async throws -> Tokenizer {
     let tokenizerName = tokenizerNameForVariant(pretrained)
-    return try await AutoTokenizer.from(pretrained: tokenizerName)
+    let hubApi = HubApi(downloadBase: tokenizerFolder)
+    return try await AutoTokenizer.from(pretrained: tokenizerName, hubApi: hubApi)
 }
 
 func formatTimestamp(_ timestamp: Float) -> String {

diff --git a/Sources/WhisperKit/Core/WhisperKit.swift b/Sources/WhisperKit/Core/WhisperKit.swift
@@ -22,6 +22,7 @@ public class WhisperKit: Transcriber {
     public var modelState: ModelState = .unloaded
     public var modelCompute: ModelComputeOptions
     public var modelFolder: URL?
+    public var tokenizerFolder: URL?
     public var tokenizer: Tokenizer?
 
     /// Protocols
@@ -54,6 +55,7 @@ public class WhisperKit: Transcriber {
         downloadBase: URL? = nil,
         modelRepo: String? = nil,
         modelFolder: String? = nil,
+        tokenizerFolder: URL? = nil,
         computeOptions: ModelComputeOptions? = nil,
         audioProcessor: (any AudioProcessing)? = nil,
         featureExtractor: (any FeatureExtracting)? = nil,
@@ -74,10 +76,17 @@ public class WhisperKit: Transcriber {
         self.textDecoder = textDecoder ?? TextDecoder()
         self.logitsFilters = logitsFilters ?? []
         self.segmentSeeker = segmentSeeker ?? SegmentSeeker()
+        self.tokenizerFolder = tokenizerFolder
         Logging.shared.logLevel = verbose ? logLevel : .none
         currentTimings = TranscriptionTimings()
 
-        try await setupModels(model: model, downloadBase: downloadBase, modelRepo: modelRepo, modelFolder: modelFolder, download: download)
+        try await setupModels(
+            model: model,
+            downloadBase: downloadBase,
+            modelRepo: modelRepo,
+            modelFolder: modelFolder,
+            download: download
+        )
 
         if let prewarm = prewarm, prewarm {
             Logging.info("Prewarming models...")
@@ -283,7 +292,7 @@ public class WhisperKit: Transcriber {
         {
             modelVariant = detectVariant(logitsDim: logitsDim, encoderDim: encoderDim)
             Logging.debug("Loading tokenizer for \(modelVariant)")
-            tokenizer = try await loadTokenizer(for: modelVariant)
+            tokenizer = try await loadTokenizer(for: modelVariant, tokenizerFolder: tokenizerFolder)
             textDecoder.tokenizer = tokenizer
             Logging.debug("Loaded tokenizer")
         } else {

diff --git a/Sources/WhisperKitCLI/CLIArguments.swift b/Sources/WhisperKitCLI/CLIArguments.swift
@@ -8,7 +8,16 @@ struct CLIArguments: ParsableArguments {
     var audioPath: String = "Tests/WhisperKitTests/Resources/jfk.wav"
 
     @Option(help: "Path of model files")
-    var modelPath: String = "Models/whisperkit-coreml/openai_whisper-tiny"
+    var modelPath: String?
+
+    @Option(help: "Model to download if no modelPath is provided")
+    var model: String?
+
+    @Option(help: "Path to save the downloaded model")
+    var downloadModelPath: String?
+
+    @Option(help: "Path to save the downloaded tokenizer files")
+    var downloadTokenizerPath: String?
 
     @Option(help: "Compute units for audio encoder model with {all,cpuOnly,cpuAndGPU,cpuAndNeuralEngine,random}")
     var audioEncoderComputeUnits: ComputeUnits = .cpuAndNeuralEngine

diff --git a/Sources/WhisperKitCLI/CLIUtils.swift b/Sources/WhisperKitCLI/CLIUtils.swift
@@ -4,7 +4,6 @@
 import ArgumentParser
 import CoreML
 import Foundation
-import WhisperKit
 
 enum ComputeUnits: String, ExpressibleByArgument, CaseIterable {
     case all, cpuAndGPU, cpuOnly, cpuAndNeuralEngine, random

diff --git a/Sources/WhisperKitCLI/transcribe.swift → Sources/WhisperKitCLI/Transcribe.swift b/Sources/WhisperKitCLI/transcribe.swift → Sources/WhisperKitCLI/Transcribe.swift
@@ -7,48 +7,56 @@ import Foundation
 import WhisperKit
 
 @available(macOS 13, iOS 16, watchOS 10, visionOS 1, *)
-@main
-struct WhisperKitCLI: AsyncParsableCommand {
+struct Transcribe: AsyncParsableCommand {
     static let configuration = CommandConfiguration(
-        commandName: "transcribe",
-        abstract: "WhisperKit Transcribe CLI",
-        discussion: "Swift native speech recognition with Whisper for Apple Silicon"
+        abstract: "Transcribe audio to text using WhisperKit"
     )
 
     @OptionGroup 
     var cliArguments: CLIArguments
 
     mutating func run() async throws {
         if cliArguments.stream {
-            try await transcribeStream(modelPath: cliArguments.modelPath)
+            try await transcribeStream()
         } else {
-            let audioURL = URL(fileURLWithPath: cliArguments.audioPath)
-            if cliArguments.verbose {
-                print("Transcribing audio at \(audioURL)")
-            }
-            try await transcribe(audioPath: cliArguments.audioPath, modelPath: cliArguments.modelPath)
+            try await transcribe()
         }
     }
 
-    private func transcribe(audioPath: String, modelPath: String) async throws {
-        let resolvedModelPath = resolveAbsolutePath(modelPath)
-        guard FileManager.default.fileExists(atPath: resolvedModelPath) else {
-            fatalError("Model path does not exist \(resolvedModelPath)")
-        }
-
-        let resolvedAudioPath = resolveAbsolutePath(audioPath)
+    private func transcribe() async throws {
+        let resolvedAudioPath = resolveAbsolutePath(cliArguments.audioPath)
         guard FileManager.default.fileExists(atPath: resolvedAudioPath) else {
-            fatalError("Resource path does not exist \(resolvedAudioPath)")
+            throw CocoaError.error(.fileNoSuchFile)
+        }
+        if cliArguments.verbose {
+            print("Transcribing audio at \(cliArguments.audioPath)")
         }
 
         let computeOptions = ModelComputeOptions(
             audioEncoderCompute: cliArguments.audioEncoderComputeUnits.asMLComputeUnits,
             textDecoderCompute: cliArguments.textDecoderComputeUnits.asMLComputeUnits
         )
+
+        let downloadTokenizerFolder: URL? =
+            if let filePath = cliArguments.downloadTokenizerPath {
+                URL(filePath: filePath)
+            } else {
+                nil
+            }
+
+        let downloadModelFolder: URL? =
+            if let filePath = cliArguments.downloadModelPath {
+                URL(filePath: filePath)
+            } else {
+                nil
+            }
 
         print("Initializing models...")
         let whisperKit = try await WhisperKit(
-            modelFolder: modelPath,
+            model: cliArguments.model,
+            downloadBase: downloadModelFolder,
+            modelFolder: cliArguments.modelPath,
+            tokenizerFolder: downloadTokenizerFolder,
             computeOptions: computeOptions,
             verbose: cliArguments.verbose,
             logLevel: .debug
@@ -82,7 +90,7 @@ struct WhisperKitCLI: AsyncParsableCommand {
         let transcription = transcribeResult?.text ?? "Transcription failed"
 
         if cliArguments.report, let result = transcribeResult {
-            let audioFileName = URL(fileURLWithPath: audioPath).lastPathComponent.components(separatedBy: ".").first!
+            let audioFileName = URL(fileURLWithPath: cliArguments.audioPath).lastPathComponent.components(separatedBy: ".").first!
 
             // Write SRT (SubRip Subtitle Format) for the transcription
             let srtReportWriter = WriteSRT(outputDir: cliArguments.reportPath)
@@ -116,15 +124,32 @@ struct WhisperKitCLI: AsyncParsableCommand {
         }
     }
 
-    private func transcribeStream(modelPath: String) async throws {
+    private func transcribeStream() async throws {
         let computeOptions = ModelComputeOptions(
             audioEncoderCompute: cliArguments.audioEncoderComputeUnits.asMLComputeUnits,
             textDecoderCompute: cliArguments.textDecoderComputeUnits.asMLComputeUnits
         )
 
+        let downloadTokenizerFolder: URL? =
+            if let filePath = cliArguments.downloadTokenizerPath {
+                URL(filePath: filePath)
+            } else {
+                nil
+            }
+
+        let downloadModelFolder: URL? =
+            if let filePath = cliArguments.downloadModelPath {
+                URL(filePath: filePath)
+            } else {
+                nil
+            }
+
         print("Initializing models...")
         let whisperKit = try await WhisperKit(
-            modelFolder: modelPath,
+            model: cliArguments.model,
+            downloadBase: downloadModelFolder,
+            modelFolder: cliArguments.modelPath,
+            tokenizerFolder: downloadTokenizerFolder,
             computeOptions: computeOptions,
             verbose: cliArguments.verbose,
             logLevel: .debug

diff --git a/Sources/WhisperKitCLI/WhisperKitCLI.swift b/Sources/WhisperKitCLI/WhisperKitCLI.swift
@@ -0,0 +1,16 @@
+//  For licensing see accompanying LICENSE.md file.
+//  Copyright © 2024 Argmax, Inc. All rights reserved.
+
+import ArgumentParser
+import Foundation
+
+@available(macOS 13, iOS 16, watchOS 10, visionOS 1, *)
+@main
+struct WhisperKitCLI: AsyncParsableCommand {
+    static let configuration = CommandConfiguration(
+        commandName: "whisperkit-cli",
+        abstract: "WhisperKit CLI",
+        discussion: "Swift native speech recognition with Whisper for Apple Silicon",
+        subcommands: [Transcribe.self]
+    )
+}
diff --git a/Tests/WhisperKitTests/UnitTests.swift b/Tests/WhisperKitTests/UnitTests.swift
@@ -4,6 +4,7 @@
 import AVFoundation
 import CoreML
 import Tokenizers
+import Hub
 @testable import WhisperKit
 import XCTest