fix: store model in Application Support, load at startup, skip network when cached
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
@@ -12,13 +12,68 @@ public final class WhisperKitTranscriber: Transcriber {
|
||||
private let modelName: String
|
||||
private var pipeline: WhisperKit?
|
||||
|
||||
/// Download progress (0...1) during the one-time model download. May be called from any thread.
|
||||
public var onProgress: ((Double) -> Void)?
|
||||
|
||||
public init(model: String = "large-v3-v20240930_turbo") {
|
||||
self.modelName = model
|
||||
}
|
||||
|
||||
/// Downloads the model on first run (≈1.5 GB) and loads it into memory.
|
||||
/// Hugging Face hub root; WhisperKit stores models in `<root>/models/<repo>/<variant>`.
|
||||
/// Kept in Application Support so it is neither in ~/Documents (TCC prompt, iCloud) nor re-fetched.
|
||||
private static var storageBase: URL {
|
||||
FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0]
|
||||
.appending(path: "Sst/models", directoryHint: .isDirectory)
|
||||
}
|
||||
|
||||
private static let repo = "argmaxinc/whisperkit-coreml"
|
||||
|
||||
/// A previously downloaded variant folder, if it contains the compiled Core ML models.
|
||||
private func localModelFolder(base: URL) -> URL? {
|
||||
let repoDir = base.appending(path: "models/\(Self.repo)", directoryHint: .isDirectory)
|
||||
let folders = (try? FileManager.default.contentsOfDirectory(atPath: repoDir.path)) ?? []
|
||||
return folders
|
||||
.filter { $0.hasSuffix(modelName) }
|
||||
.map { repoDir.appending(path: $0, directoryHint: .isDirectory) }
|
||||
.first { folder in
|
||||
["MelSpectrogram", "AudioEncoder", "TextDecoder"].allSatisfy {
|
||||
FileManager.default.fileExists(
|
||||
atPath: folder.appending(path: "\($0).mlmodelc/coremldata.bin").path)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func loadPipeline(from folder: URL, base: URL) async throws -> WhisperKit {
|
||||
try await WhisperKit(WhisperKitConfig(
|
||||
downloadBase: base,
|
||||
modelFolder: folder.path,
|
||||
tokenizerFolder: base,
|
||||
prewarm: true,
|
||||
load: true,
|
||||
download: false
|
||||
))
|
||||
}
|
||||
|
||||
/// Loads the model into memory. Downloads it (≈1.5 GB) only if it is not stored locally yet;
|
||||
/// once downloaded, launching needs no network.
|
||||
public func load() async throws {
|
||||
pipeline = try await WhisperKit(WhisperKitConfig(model: modelName))
|
||||
guard pipeline == nil else { return }
|
||||
let base = Self.storageBase
|
||||
try FileManager.default.createDirectory(at: base, withIntermediateDirectories: true)
|
||||
|
||||
if let folder = localModelFolder(base: base),
|
||||
let loaded = try? await loadPipeline(from: folder, base: base) {
|
||||
pipeline = loaded
|
||||
return
|
||||
}
|
||||
let folder = try await WhisperKit.download(
|
||||
variant: modelName,
|
||||
downloadBase: base,
|
||||
from: Self.repo
|
||||
) { [onProgress] progress in
|
||||
onProgress?(progress.fractionCompleted)
|
||||
}
|
||||
pipeline = try await loadPipeline(from: folder, base: base)
|
||||
}
|
||||
|
||||
public func transcribe(_ samples: [Float]) async throws -> String {
|
||||
|
||||
Reference in New Issue
Block a user