106 lines
3.7 KiB
Swift
106 lines
3.7 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 let storage: ModelStorage
|
|
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 = ModelCatalog.defaultVariant, storage: ModelStorage = ModelStorage()) {
|
|
self.modelName = model
|
|
self.storage = storage
|
|
}
|
|
|
|
public var variant: String { modelName }
|
|
|
|
private func loadPipeline(from folder: URL) async throws -> WhisperKit {
|
|
try await WhisperKit(WhisperKitConfig(
|
|
downloadBase: storage.base,
|
|
modelFolder: folder.path,
|
|
tokenizerFolder: storage.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 }
|
|
try FileManager.default.createDirectory(at: storage.base, withIntermediateDirectories: true)
|
|
|
|
if let folder = storage.folder(for: modelName),
|
|
let loaded = try? await loadPipeline(from: folder) {
|
|
pipeline = loaded
|
|
return
|
|
}
|
|
let folder = try await WhisperKit.download(
|
|
variant: modelName,
|
|
downloadBase: storage.base,
|
|
from: ModelStorage.repo
|
|
) { [onProgress] progress in
|
|
onProgress?(progress.fractionCompleted)
|
|
}
|
|
pipeline = try await loadPipeline(from: folder)
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|