311 lines
11 KiB
Swift
311 lines
11 KiB
Swift
import Foundation
|
|
import WhisperKit
|
|
|
|
struct LocalTranscriptionModel: Identifiable, Hashable {
|
|
let id: String
|
|
let url: URL
|
|
let isInstalled: Bool
|
|
|
|
init(id: String, url: URL, isInstalled: Bool = true) {
|
|
self.id = id
|
|
self.url = url
|
|
self.isInstalled = isInstalled
|
|
}
|
|
|
|
var displayName: String {
|
|
Self.displayName(for: id)
|
|
}
|
|
|
|
var installStateLabel: String {
|
|
isInstalled ? "Installiert" : "Nicht installiert"
|
|
}
|
|
|
|
var shortDisplayName: String {
|
|
if id.contains("small") {
|
|
return "Whisper Small"
|
|
}
|
|
if id.contains("base") {
|
|
return "Whisper Base"
|
|
}
|
|
if id.contains("tiny") {
|
|
return "Whisper Tiny"
|
|
}
|
|
if id.contains("turbo") {
|
|
return "Whisper Turbo"
|
|
}
|
|
if id.contains("large-v3") {
|
|
return "Whisper Large"
|
|
}
|
|
return displayName
|
|
}
|
|
|
|
static func displayName(for modelName: String) -> String {
|
|
if modelName.contains("small") {
|
|
return "Whisper Small"
|
|
}
|
|
if modelName.contains("base") {
|
|
return "Whisper Base"
|
|
}
|
|
if modelName.contains("tiny") {
|
|
return "Whisper Tiny"
|
|
}
|
|
if modelName.contains("turbo") {
|
|
return "Whisper Large v3 Turbo"
|
|
}
|
|
if modelName.contains("large-v3") {
|
|
return "Whisper Large v3"
|
|
}
|
|
return modelName
|
|
.replacingOccurrences(of: "openai_", with: "")
|
|
.replacingOccurrences(of: "_", with: " ")
|
|
.replacingOccurrences(of: "-", with: " ")
|
|
}
|
|
}
|
|
|
|
enum LocalTranscriptionError: LocalizedError {
|
|
case modelMissing(URL)
|
|
case downloadedModelInvalid(String)
|
|
case noText
|
|
|
|
var errorDescription: String? {
|
|
switch self {
|
|
case .modelMissing(let url):
|
|
return "Lokales Modell fehlt: \(url.path)"
|
|
case .downloadedModelInvalid(let modelName):
|
|
return "Das geladene Modell ist unvollständig: \(modelName)"
|
|
case .noText:
|
|
return "Das lokale Modell hat keinen Text erkannt."
|
|
}
|
|
}
|
|
}
|
|
|
|
actor LocalTranscriptionService {
|
|
static let shared = LocalTranscriptionService()
|
|
|
|
static let defaultModelName = "openai_whisper-large-v3-v20240930_626MB"
|
|
static let fastModelName = "openai_whisper-large-v3-v20240930_turbo_632MB"
|
|
static let recommendedFastModelName = "openai_whisper-small_216MB"
|
|
static let modelRepo = "argmaxinc/whisperkit-coreml"
|
|
static let supportedModelNames = [
|
|
recommendedFastModelName,
|
|
fastModelName,
|
|
defaultModelName
|
|
]
|
|
static let modelPageURL = URL(
|
|
string: "https://huggingface.co/argmaxinc/whisperkit-coreml/tree/main/openai_whisper-large-v3-v20240930_626MB"
|
|
)!
|
|
static let fastModelPageURL = URL(
|
|
string: "https://huggingface.co/argmaxinc/whisperkit-coreml/tree/main/openai_whisper-large-v3-v20240930_turbo_632MB"
|
|
)!
|
|
static let recommendedFastModelPageURL = URL(
|
|
string: "https://huggingface.co/argmaxinc/whisperkit-coreml/tree/main/openai_whisper-small_216MB"
|
|
)!
|
|
|
|
static func modelPageURL(for modelName: String) -> URL {
|
|
switch normalizedModelName(modelName) {
|
|
case recommendedFastModelName:
|
|
return recommendedFastModelPageURL
|
|
case fastModelName:
|
|
return fastModelPageURL
|
|
case defaultModelName:
|
|
return modelPageURL
|
|
default:
|
|
return URL(string: "https://huggingface.co/\(modelRepo)/tree/main/\(normalizedModelName(modelName))")!
|
|
}
|
|
}
|
|
|
|
private var whisperKit: WhisperKit?
|
|
private var loadedModelName: String?
|
|
|
|
static var isModelInstalled: Bool {
|
|
isModelInstalled(defaultModelName)
|
|
}
|
|
|
|
static func modelURL(named modelName: String) -> URL {
|
|
AppSupportPaths.whisperKitModelsDirectoryURL.appendingPathComponent(normalizedModelName(modelName), isDirectory: true)
|
|
}
|
|
|
|
static func isModelInstalled(_ modelName: String) -> Bool {
|
|
isUsableModel(at: modelURL(named: modelName))
|
|
}
|
|
|
|
static func normalizedModelName(_ modelName: String) -> String {
|
|
let trimmed = modelName.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
return trimmed.isEmpty ? recommendedFastModelName : trimmed
|
|
}
|
|
|
|
static func installedModels() -> [LocalTranscriptionModel] {
|
|
let directory = AppSupportPaths.whisperKitModelsDirectoryURL
|
|
let urls = (try? FileManager.default.contentsOfDirectory(
|
|
at: directory,
|
|
includingPropertiesForKeys: [.isDirectoryKey],
|
|
options: [.skipsHiddenFiles]
|
|
)) ?? []
|
|
|
|
return urls
|
|
.filter { isUsableModel(at: $0) }
|
|
.map { LocalTranscriptionModel(id: $0.lastPathComponent, url: $0) }
|
|
.sorted { lhs, rhs in
|
|
if lhs.id == recommendedFastModelName { return true }
|
|
if rhs.id == recommendedFastModelName { return false }
|
|
if lhs.id == fastModelName { return true }
|
|
if rhs.id == fastModelName { return false }
|
|
if lhs.id == defaultModelName { return true }
|
|
if rhs.id == defaultModelName { return false }
|
|
return lhs.displayName.localizedCaseInsensitiveCompare(rhs.displayName) == .orderedAscending
|
|
}
|
|
}
|
|
|
|
static func modelOptions() -> [LocalTranscriptionModel] {
|
|
var seen = Set<String>()
|
|
let installed = installedModels()
|
|
let installedByID = Dictionary(uniqueKeysWithValues: installed.map { ($0.id, $0) })
|
|
let orderedIDs = supportedModelNames + installed.map(\.id)
|
|
|
|
return orderedIDs.compactMap { modelName in
|
|
let normalizedName = normalizedModelName(modelName)
|
|
guard seen.insert(normalizedName).inserted else { return nil }
|
|
|
|
if let installedModel = installedByID[normalizedName] {
|
|
return installedModel
|
|
}
|
|
|
|
return LocalTranscriptionModel(
|
|
id: normalizedName,
|
|
url: modelURL(named: normalizedName),
|
|
isInstalled: false
|
|
)
|
|
}
|
|
}
|
|
|
|
static func resolvedModelName(_ preferredModelName: String) -> String {
|
|
let normalizedName = normalizedModelName(preferredModelName)
|
|
if isModelInstalled(normalizedName) {
|
|
return normalizedName
|
|
}
|
|
|
|
return installedModels().first?.id ?? normalizedName
|
|
}
|
|
|
|
static func shouldAutoSelectRecommendedFastModel(currentModelName: String) -> Bool {
|
|
guard isModelInstalled(recommendedFastModelName) else {
|
|
return false
|
|
}
|
|
|
|
return currentModelName == defaultModelName || currentModelName == fastModelName
|
|
}
|
|
|
|
func prepare(modelName: String) async throws {
|
|
_ = try await pipeline(modelName: modelName)
|
|
}
|
|
|
|
func downloadAndInstall(
|
|
modelName: String,
|
|
progressHandler: @escaping @Sendable (Double) -> Void
|
|
) async throws -> URL {
|
|
let normalizedName = Self.normalizedModelName(modelName)
|
|
let destinationURL = Self.modelURL(named: normalizedName)
|
|
|
|
if Self.isUsableModel(at: destinationURL) {
|
|
progressHandler(1)
|
|
return destinationURL
|
|
}
|
|
|
|
let fileManager = FileManager.default
|
|
try fileManager.createDirectory(
|
|
at: AppSupportPaths.whisperKitModelsDirectoryURL,
|
|
withIntermediateDirectories: true
|
|
)
|
|
|
|
let downloadRoot = AppSupportPaths.localModelsDirectoryURL
|
|
.appendingPathComponent("downloads", isDirectory: true)
|
|
.appendingPathComponent(UUID().uuidString, isDirectory: true)
|
|
try fileManager.createDirectory(at: downloadRoot, withIntermediateDirectories: true)
|
|
|
|
do {
|
|
let downloadedURL = try await WhisperKit.download(
|
|
variant: normalizedName,
|
|
downloadBase: downloadRoot,
|
|
from: Self.modelRepo
|
|
) { progress in
|
|
let fraction = progress.fractionCompleted
|
|
progressHandler(fraction.isFinite ? fraction : 0)
|
|
}
|
|
|
|
guard Self.isUsableModel(at: downloadedURL) else {
|
|
throw LocalTranscriptionError.downloadedModelInvalid(normalizedName)
|
|
}
|
|
|
|
if fileManager.fileExists(atPath: destinationURL.path) {
|
|
try fileManager.removeItem(at: destinationURL)
|
|
}
|
|
try fileManager.moveItem(at: downloadedURL, to: destinationURL)
|
|
try? fileManager.removeItem(at: downloadRoot)
|
|
|
|
if loadedModelName == normalizedName {
|
|
whisperKit = nil
|
|
loadedModelName = nil
|
|
}
|
|
|
|
progressHandler(1)
|
|
return destinationURL
|
|
} catch {
|
|
try? fileManager.removeItem(at: downloadRoot)
|
|
throw error
|
|
}
|
|
}
|
|
|
|
func transcribe(audioURL: URL, language: String, modelName: String) async throws -> String {
|
|
let resolvedLanguage = language.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
let decodeOptions = DecodingOptions(
|
|
task: .transcribe,
|
|
language: resolvedLanguage.isEmpty ? nil : resolvedLanguage
|
|
)
|
|
|
|
let pipeline = try await pipeline(modelName: modelName)
|
|
let results = try await pipeline.transcribe(
|
|
audioPath: audioURL.path,
|
|
decodeOptions: decodeOptions
|
|
)
|
|
let text = results
|
|
.map(\.text)
|
|
.joined(separator: " ")
|
|
.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
|
|
guard !text.isEmpty else {
|
|
throw LocalTranscriptionError.noText
|
|
}
|
|
|
|
return text
|
|
}
|
|
|
|
private func pipeline(modelName: String) async throws -> WhisperKit {
|
|
let resolvedModelName = Self.resolvedModelName(modelName)
|
|
if let whisperKit, loadedModelName == resolvedModelName {
|
|
return whisperKit
|
|
}
|
|
|
|
let url = Self.modelURL(named: resolvedModelName)
|
|
guard Self.isUsableModel(at: url) else {
|
|
throw LocalTranscriptionError.modelMissing(url)
|
|
}
|
|
|
|
let loaded = try await WhisperKit(
|
|
modelFolder: url.path,
|
|
verbose: false,
|
|
prewarm: true,
|
|
load: true,
|
|
download: false
|
|
)
|
|
whisperKit = loaded
|
|
loadedModelName = resolvedModelName
|
|
return loaded
|
|
}
|
|
|
|
private static func isUsableModel(at url: URL) -> Bool {
|
|
FileManager.default.fileExists(atPath: url.appendingPathComponent("AudioEncoder.mlmodelc").path) &&
|
|
FileManager.default.fileExists(atPath: url.appendingPathComponent("MelSpectrogram.mlmodelc").path) &&
|
|
FileManager.default.fileExists(atPath: url.appendingPathComponent("TextDecoder.mlmodelc").path)
|
|
}
|
|
}
|