Docs
Managing Models

Managing Models

Downloading and compiling models in your application

Argmax Pro SDK provides utilities for managing model downloads and on-device compilation in your application. There are two main approaches:

  • ModelStore Approach (Recommended)
  • Direct Initializer Approach (Legacy)

ModelStore Approach

The ModelStore class provides a feature-rich way to download and manage models with several user experience benefits.

Key Benefits

  • Progress Tracking: Real-time download progress updates
  • Cancellation Support: Ability to cancel downloads
  • Repository Management: Utility to manage models across multiple model repositories
  • Decoupled model download and compilation: On-device model compilation (Specialization) happens separately from download, reducing blocking time

Basic Usage

@available(iOS 17.0, macOS 14.0, watchOS 10.0, visionOS 1.0, *)
public func setupWhisperKitPro() async throws -> WhisperKitPro {
 
    // First, obfuscate your Argmax API Key and Hugging Face model token using `ObfuscatedKeyProvider.generateCodeExample`
    // `ObfuscatedKeyProvider` protects against casual inspection and static analysis tools.
    // In production, consider fetching apiKey and modelToken securely from your backend.
    var keyProvider = ObfuscatedKeyProvider(mask: 37)
    keyProvider.huggingFaceObfuscated = [1, 2, 3]
    keyProvider.apiKeyObfuscated = [4, 5, 6]
    
    guard let apiKey = keyProvider.apiKey, let modelToken = keyProvider.huggingFaceToken else {
        fatalError("Missing API key or HuggingFace token")
    }
 
    // Initialize Argmax SDK to ensure Pro SDK licensed
    // This requires a connection during first use and once every 30 days
    await ArgmaxSDK.with(ArgmaxConfig(apiKey: apiKey))
 
    let modelRepo = RepoType.parakeetRepo  // e.g. RepoType.openSourceRepo, RepoType.proRepo
    
    // Fetch the recommended model for the current device
    let model = await WhisperKit.recommendedRemoteModels(
        from: modelRepo.rawValue,
        token: modelToken // modelToken is only needed for models from RepoType.parakeetRepo
    ).default
    
    let config = WhisperKitProConfig(
        model: model,
        modelRepo: modelRepo.rawValue,
    )
    let modelStore = ModelStore(whisperKitConfig: config)
    
    // Download model with progress tracking, print the progress
    // Or track the @Published `progress` in SwiftUI
    let cancellable = modelStore.$progress.sink { progress in
        if let progress = progress {
            let percentage = Int(progress.fractionCompleted * 100)
            print("\rDownload progress: \(percentage)%", terminator: "")
        }
    }
    
    // Download model (blocks until complete, track progress via modelStore.$progress)
    let modelURL = try await modelStore.downloadModel() 
 
    // Load model (blocks until complete, also compiles the model during first use or after OS update)
    let whisperKitPro = try await WhisperKitPro(config) 
    return whisperKitPro
}

Migration Guide

If you're currently using the direct initializer approach, here's how to migrate:

Before (Direct Initializer)

// Old approach - blocks for entire process
let config = WhisperKitProConfig(model: "large-v3-v20240930_626MB")
let whisperKitPro = try await WhisperKitPro(config)

After (ModelStore)

// New approach - better user experience
let config = WhisperKitProConfig(
    model: "large-v3-v20240930_626MB",
    modelRepo: RepoType.openSourceRepo.rawValue,
)
let modelStore = ModelStore(whisperKitConfig: config)
 
// Track progress and allow cancellation
let cancellable = modelStore.$progress.sink { progress in
    // Update UI with progress
}
 
let modelURL = try await modelStore.downloadModel()
 
// Initialize with downloaded model
let config = WhisperKitProConfig(
    modelFolder: modelURL.path(percentEncoded: false)
)
let whisperKitPro = try await WhisperKitPro(config)