Speaker Diarization

ArgmaxSpeakerDiarizationBasicExample
├── Package.swift
└── Sources
    └── ArgmaxTestCLI
        └── ArgmaxTestCommand.swift

// swift-tools-version: 6.0
// The swift-tools-version declares the minimum version of Swift required to build this package.
 
import PackageDescription
 
let package = Package(
    name: "Argmax Test CLI",
    platforms: [
        .macOS(.v13)
    ],
    products: [
        .executable(
            name: "argmax-test-cli",
            targets: ["ArgmaxTestCLI"]
        )
    ],
    dependencies: [
        .package(id: "argmaxinc.argmax-sdk-swift", .upToNextMinor(from: "2.0.9")),
        .package(url: "https://github.com/apple/swift-argument-parser.git", exact: "1.3.0")
    ],
    targets: [
        .executableTarget(
            name: "ArgmaxTestCLI",
            dependencies: [
                .product(name: "Argmax", package: "argmaxinc.argmax-sdk-swift"),
                .product(name: "ArgumentParser", package: "swift-argument-parser")
            ]
        ),
    ]
)
 

import Foundation
@preconcurrency import ArgumentParser
import Argmax
 
@main
struct ArgmaxTestCommand: AsyncParsableCommand {
    static let configuration = CommandConfiguration(
        abstract: "An example CLI tool for Argmax Pro SDK",
        subcommands: [Pyannote.self, Sortformer.self]
    )
 
    struct Pyannote: AsyncParsableCommand {
        static let configuration = CommandConfiguration(
            abstract: "Speaker diarization using Pyannote (macOS 13+)"
        )
 
        @Option(name: .long, help: "Argmax Pro SDK API key")
        var apiKey: String
 
        @Option(help: "Path to the audio file to process")
        var audioPath: String
 
        @Option(help: "Path to save the diarization output in RTTM format")
        var rttmPath: String?
 
        @Option(help: "Number of speakers to detect if known (default: automatic)")
        var numSpeakers: Int?
 
        mutating func run() async throws {
            print("Initializing Argmax Pro SDK...")
            let sdkConfig = ArgmaxConfig(apiKey: apiKey)
            await ArgmaxSDK.with(sdkConfig)
 
            print("Loading audio file from \(audioPath)...")
            let audioFrames = try AudioProcessor.loadAudioAsFloatArray(fromPath: audioPath)
 
            print("Initializing SpeakerKitPro (Pyannote)...")
            let diarizer = SpeakerKitDiarizer.pyannote()
            try await diarizer.downloadModels()
            try await diarizer.loadModels()
            let config = PyannoteConfig(diarizer: diarizer)
            let speakerKit = try await SpeakerKitPro(config)
            defer { Task { await speakerKit.unloadModels() } }
 
            print("Running diarization...")
            let options: (any DiarizationOptions)? = numSpeakers.map {
                PyannoteDiarizationOptions(numberOfSpeakers: $0)
            }
            let result = try await speakerKit.diarize(audioArray: audioFrames, options: options)
 
            outputRTTM(result: result, audioPath: audioPath, rttmPath: rttmPath)
 
            if let timings = result.timings {
                print(timings.debugDescription)
            }
        }
    }
 
    struct Sortformer: AsyncParsableCommand {
        static let configuration = CommandConfiguration(
            abstract: "Speaker diarization using Sortformer (macOS 15+)"
        )
 
        @Option(name: .long, help: "Argmax Pro SDK API key")
        var apiKey: String
 
        @Option(help: "Path to the audio file to process")
        var audioPath: String
 
        @Option(help: "Path to save the diarization output in RTTM format")
        var rttmPath: String?
 
        mutating func run() async throws {
            if #unavailable(macOS 15, iOS 18, watchOS 11, visionOS 2) {
                throw ValidationError("Sortformer requires macOS 15 / iOS 18 or later.")
            }
 
            print("Initializing Argmax Pro SDK...")
            let sdkConfig = ArgmaxConfig(apiKey: apiKey)
            await ArgmaxSDK.with(sdkConfig)
 
            print("Loading audio file from \(audioPath)...")
            let audioFrames = try AudioProcessor.loadAudioAsFloatArray(fromPath: audioPath)
 
            print("Initializing SpeakerKitPro (Sortformer)...")
            let speakerKit = try await SpeakerKitPro()
            defer { Task { await speakerKit.unloadModels() } }
 
            print("Running diarization...")
            let options = SortformerDiarizationOptions(sortformerMode: .prerecorded)
            let result = try await speakerKit.diarize(audioArray: audioFrames, options: options)
 
            outputRTTM(result: result, audioPath: audioPath, rttmPath: rttmPath)
 
            if let timings = result.timings {
                print(timings.debugDescription)
            }
        }
    }
}
 
private func outputRTTM(result: DiarizationResult, audioPath: String, rttmPath: String?) {
    let fileName = URL(filePath: audioPath).deletingPathExtension().lastPathComponent
    let rttmLines = SpeakerKit.generateRTTM(from: result, fileName: fileName)
    let rttmContent = rttmLines.map { $0.description }.joined(separator: "\n")
 
    if let rttmPath {
        do {
            try rttmContent.write(to: URL(filePath: rttmPath), atomically: true, encoding: .utf8)
            print("RTTM saved to \(rttmPath)")
        } catch {
            print("Failed to write RTTM: \(error)")
        }
    } else {
        print(rttmContent)
    }
}
 

# Pyannote (macOS 13+)
swift run argmax-test-cli pyannote --api-key <API_KEY> --audio-path <AUDIO_PATH>
 
# Sortformer (macOS 15+)
swift run argmax-test-cli sortformer --api-key <API_KEY> --audio-path <AUDIO_PATH>