127 lines
4.8 KiB
Swift
127 lines
4.8 KiB
Swift
import Foundation
|
|
import SstCore
|
|
import WhisperKit
|
|
|
|
public final class WhisperKitTranscriber: Transcriber {
|
|
public enum TranscriberError: Error, LocalizedError {
|
|
case notLoaded
|
|
|
|
public var errorDescription: String? { "Spraakmodel is nog niet geladen." }
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
/// 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 {
|
|
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 {
|
|
try await transcribe(samples, vocabulary: [])
|
|
}
|
|
|
|
/// Terms are added one by one while the encoded prompt stays within `maxTokens`.
|
|
/// Returns nil when nothing fits or the list is empty.
|
|
static func promptTokens(
|
|
vocabulary: [String],
|
|
maxTokens: Int = 64,
|
|
specialTokenBegin: Int,
|
|
encode: (String) -> [Int]
|
|
) -> [Int]? {
|
|
var kept: [String] = []
|
|
var tokens: [Int] = []
|
|
for term in vocabulary {
|
|
let candidate = kept + [term]
|
|
let encoded = encode(" " + candidate.joined(separator: ", ")).filter { $0 < specialTokenBegin }
|
|
if encoded.count > maxTokens { break }
|
|
kept = candidate
|
|
tokens = encoded
|
|
}
|
|
return tokens.isEmpty ? nil : tokens
|
|
}
|
|
|
|
public func transcribe(_ samples: [Float], vocabulary: [String]) async throws -> String {
|
|
guard let pipeline else { throw TranscriberError.notLoaded }
|
|
// language nil + detectLanguage: Whisper picks NL or EN per utterance.
|
|
var options = DecodingOptions(
|
|
task: .transcribe,
|
|
language: nil,
|
|
detectLanguage: true,
|
|
skipSpecialTokens: true,
|
|
withoutTimestamps: true
|
|
)
|
|
if !vocabulary.isEmpty, let tokenizer = pipeline.tokenizer {
|
|
options.promptTokens = Self.promptTokens(
|
|
vocabulary: vocabulary,
|
|
specialTokenBegin: tokenizer.specialTokens.specialTokenBegin,
|
|
encode: { tokenizer.encode(text: $0) }
|
|
)
|
|
}
|
|
let results = try await pipeline.transcribe(audioArray: samples, decodeOptions: options)
|
|
return results
|
|
.map(\.text)
|
|
.joined(separator: " ")
|
|
.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
}
|
|
}
|