diff --git a/Documentation/Benchmarks.md b/Documentation/Benchmarks.md index 7a5976a41..704461a44 100644 --- a/Documentation/Benchmarks.md +++ b/Documentation/Benchmarks.md @@ -734,3 +734,52 @@ Both the English BART G2P and multilingual ByT5 G2P models run fastest on CPU-on | cpuOnly | **13.0** | | all (ANE+GPU+CPU) | 17.3 | | cpuAndGPU | 23.4 | + +## CTC zh-CN Mandarin ASR (Experimental) + +Parakeet CTC 0.6B zh-CN model converted to CoreML for on-device Mandarin Chinese transcription. + +> **⚠️ Experimental Feature**: This is an early preview of Mandarin Chinese ASR support. The API and performance characteristics may change in future releases. + +Model: [FluidInference/parakeet-ctc-0.6b-zh-cn-coreml](https://huggingface.co/FluidInference/parakeet-ctc-0.6b-zh-cn-coreml) + +Hardware: Apple M2, 2022, macOS 26 + +### THCHS-30 Test Set + +Full benchmark on the complete THCHS-30 test set — 2,495 utterances (250 unique sentences × 10 speakers) from the THCHS-30 corpus. + +Dataset: [FluidInference/THCHS-30-tests](https://huggingface.co/datasets/FluidInference/THCHS-30-tests) + +```bash +swift run -c release fluidaudiocli ctc-zh-cn-benchmark --auto-download +``` + +| Metric | int8 encoder (0.55 GB) | +|---|---| +| **Mean CER** | **8.23%** | +| **Median CER** | **6.45%** | +| CER = 0% (perfect) | 435 (17.4%) | +| CER < 5% | 947 (38.0%) | +| CER < 10% | 1,674 (67.1%) | +| CER < 20% | 2,325 (93.2%) | +| Mean Latency | 614 ms | +| Mean RTFx | 14.83x | + +### Error Analysis + +Error analysis from the 100 highest-CER samples (out of the full 2,495) identified 862 substitution errors. The dominant patterns: + +- **Homophones / near-homophones**: acoustically similar syllables (e.g. 呢/了, 了/的) account for the majority of substitutions — unavoidable without a language model +- **Digit representation**: the model may output Arabic digits (1, 5, 2011) when references use Chinese characters (一五, 二零一一); the benchmark normalizer converts digits before scoring to avoid penalizing this +- **Sentence-final particles**: 了/的/呢/吧 are frequently confused, contributing a disproportionate share of errors given their high occurrence + +### Beam Search + +Beam search does not improve CER for this model without a language model. Greedy decoding (beam width 1) is recommended. + +### Recommendations + +- **Greedy decoding** is sufficient for production use at this CER level. +- For applications requiring <8% CER, a character-level language model would be needed. +- Int8 encoder (0.55 GB) performs on par with FP32 (1.1 GB). diff --git a/Sources/FluidAudio/ASR/Parakeet/AsrManager.swift b/Sources/FluidAudio/ASR/Parakeet/AsrManager.swift index 634194458..a9dde46b2 100644 --- a/Sources/FluidAudio/ASR/Parakeet/AsrManager.swift +++ b/Sources/FluidAudio/ASR/Parakeet/AsrManager.swift @@ -293,6 +293,10 @@ public actor AsrManager { isLastChunk: isLastChunk, globalFrameOffset: globalFrameOffset ) + case .ctcZhCn: + throw ASRError.processingFailed( + "CTC-only model .ctcZhCn does not support TDT decoding. Use CtcZhCnManager instead." + ) } } diff --git a/Sources/FluidAudio/ASR/Parakeet/AsrModels.swift b/Sources/FluidAudio/ASR/Parakeet/AsrModels.swift index d28a372ad..c56caa239 100644 --- a/Sources/FluidAudio/ASR/Parakeet/AsrModels.swift +++ b/Sources/FluidAudio/ASR/Parakeet/AsrModels.swift @@ -7,12 +7,15 @@ public enum AsrModelVersion: Sendable { case v3 /// 110M parameter hybrid TDT-CTC model with fused preprocessor+encoder case tdtCtc110m + /// 600M parameter CTC-only model for Mandarin Chinese (zh-CN) + case ctcZhCn var repo: Repo { switch self { case .v2: return .parakeetV2 case .v3: return .parakeet case .tdtCtc110m: return .parakeetTdtCtc110m + case .ctcZhCn: return .parakeetCtcZhCn } } @@ -24,10 +27,19 @@ public enum AsrModelVersion: Sendable { } } + /// Whether this model is CTC-only (no TDT decoder+joint) + public var isCtcOnly: Bool { + switch self { + case .ctcZhCn: return true + default: return false + } + } + /// Encoder hidden dimension for this model version public var encoderHiddenSize: Int { switch self { case .tdtCtc110m: return 512 + case .ctcZhCn: return 1024 default: return 1024 } } @@ -37,6 +49,7 @@ public enum AsrModelVersion: Sendable { switch self { case .v2, .tdtCtc110m: return 1024 case .v3: return 8192 + case .ctcZhCn: return 7000 } } diff --git a/Sources/FluidAudio/ASR/Parakeet/CtcZhCnManager.swift b/Sources/FluidAudio/ASR/Parakeet/CtcZhCnManager.swift new file mode 100644 index 000000000..1ee8af6c7 --- /dev/null +++ b/Sources/FluidAudio/ASR/Parakeet/CtcZhCnManager.swift @@ -0,0 +1,207 @@ +@preconcurrency import CoreML +import Foundation + +/// Manager for Parakeet CTC zh-CN transcription +/// +/// This manager handles the full pipeline for Mandarin Chinese CTC transcription: +/// 1. Preprocessor: Audio → Mel spectrogram +/// 2. Encoder: Mel → Encoder features +/// 3. CTC Decoder: Encoder features → CTC logits +/// 4. Greedy CTC decoding: Logits → Text +public actor CtcZhCnManager { + + private let models: CtcZhCnModels + private let maxAudioSamples: Int + private let sampleRate: Int + + private static let logger = AppLogger(category: "CtcZhCnManager") + + /// Initialize with pre-loaded models + public init(models: CtcZhCnModels, maxAudioSamples: Int = 240_000, sampleRate: Int = 16_000) { + self.models = models + self.maxAudioSamples = maxAudioSamples + self.sampleRate = sampleRate + } + + /// Convenience initializer that loads models from default cache directory + public static func load( + useInt8Encoder: Bool = true, + configuration: MLModelConfiguration? = nil, + progressHandler: DownloadUtils.ProgressHandler? = nil + ) async throws -> CtcZhCnManager { + let models = try await CtcZhCnModels.downloadAndLoad( + useInt8Encoder: useInt8Encoder, + configuration: configuration, + progressHandler: progressHandler + ) + return CtcZhCnManager(models: models) + } + + /// Transcribe audio to text using CTC decoding + /// + /// - Parameters: + /// - audio: Audio samples (mono, 16kHz) + /// - audioLength: Optional audio length (if nil, uses audio.count) + /// - Returns: Transcribed text + public func transcribe( + audio: [Float], + audioLength: Int? = nil + ) throws -> String { + let actualLength = audioLength ?? audio.count + + // Pad or truncate audio to maxAudioSamples + let paddedAudio = padOrTruncateAudio(audio, targetLength: maxAudioSamples) + + // Step 1: Preprocessor (audio → mel spectrogram) + let melOutput = try runPreprocessor(audio: paddedAudio, audioLength: actualLength) + + // Step 2: Encoder (mel → encoder features) + let encoderOutput = try runEncoder(mel: melOutput.mel, melLength: melOutput.melLength) + + // Step 3: CTC Decoder (encoder features → CTC logits) + let ctcLogits = try runCtcDecoder(encoderOutput: encoderOutput) + + // Step 4: CTC decoding (logits → text) + let text = greedyCtcDecode(logits: ctcLogits) + + return text + } + + /// Transcribe audio file to text + /// + /// - Parameters: + /// - audioURL: URL to audio file (will be resampled to 16kHz mono) + /// - Returns: Transcribed text + public func transcribe(audioURL: URL) throws -> String { + // Load and convert audio + let converter = AudioConverter(sampleRate: Double(sampleRate)) + let samples = try converter.resampleAudioFile(audioURL) + + return try transcribe(audio: samples) + } + + // MARK: - Private Pipeline Methods + + private struct MelOutput { + let mel: MLMultiArray + let melLength: MLMultiArray + } + + private func runPreprocessor(audio: [Float], audioLength: Int) throws -> MelOutput { + // Create input arrays + let audioArray = try MLMultiArray(shape: [1, maxAudioSamples as NSNumber], dataType: .float32) + for (i, sample) in audio.enumerated() where i < maxAudioSamples { + audioArray[i] = NSNumber(value: sample) + } + + let audioLengthArray = try MLMultiArray(shape: [1], dataType: .int32) + audioLengthArray[0] = NSNumber(value: min(audioLength, maxAudioSamples)) + + // Run preprocessor + let input = try MLDictionaryFeatureProvider( + dictionary: [ + "audio_signal": MLFeatureValue(multiArray: audioArray), + "audio_length": MLFeatureValue(multiArray: audioLengthArray), + ] + ) + let output = try models.preprocessor.prediction(from: input) + + guard + let mel = output.featureValue(for: "mel")?.multiArrayValue, + let melLength = output.featureValue(for: "mel_length")?.multiArrayValue + else { + throw ASRError.processingFailed("Failed to extract mel or mel_length from preprocessor output") + } + + return MelOutput(mel: mel, melLength: melLength) + } + + private func runEncoder(mel: MLMultiArray, melLength: MLMultiArray) throws -> MLMultiArray { + // Run encoder + let input = try MLDictionaryFeatureProvider( + dictionary: [ + "audio_signal": MLFeatureValue(multiArray: mel), + "length": MLFeatureValue(multiArray: melLength), + ] + ) + let output = try models.encoder.prediction(from: input) + + guard let encoderOutput = output.featureValue(for: "encoder_output")?.multiArrayValue else { + throw ASRError.processingFailed("Failed to extract encoder_output from encoder") + } + + return encoderOutput + } + + private func runCtcDecoder(encoderOutput: MLMultiArray) throws -> MLMultiArray { + // Run CTC decoder head + let input = try MLDictionaryFeatureProvider( + dictionary: [ + "encoder_output": MLFeatureValue(multiArray: encoderOutput) + ] + ) + let output = try models.decoder.prediction(from: input) + + guard let ctcLogits = output.featureValue(for: "ctc_logits")?.multiArrayValue else { + throw ASRError.processingFailed("Failed to extract ctc_logits from decoder") + } + + return ctcLogits + } + + private func greedyCtcDecode(logits: MLMultiArray) -> String { + // logits shape: [1, T, vocab_size+1] where T is time steps (188) + // vocab_size = 7000, blank_id = 7000 + + let timeSteps = logits.shape[1].intValue + let vocabSize = logits.shape[2].intValue + + var decoded: [Int] = [] + var prevLabel: Int? = nil + + for t in 0.. maxLogit { + maxLogit = logit + maxLabel = v + } + } + + // CTC collapse: skip blanks and repeats + if maxLabel != models.blankId && maxLabel != prevLabel { + decoded.append(maxLabel) + } + prevLabel = maxLabel + } + + // Convert token IDs to text + var text = "" + for tokenId in decoded { + if let token = models.vocabulary[tokenId] { + text += token + } + } + + // Replace SentencePiece underscores with spaces + text = text.replacingOccurrences(of: "▁", with: " ") + + return text.trimmingCharacters(in: .whitespacesAndNewlines) + } + + private func padOrTruncateAudio(_ audio: [Float], targetLength: Int) -> [Float] { + var result = audio + if result.count < targetLength { + // Pad with zeros + result.append(contentsOf: Array(repeating: 0.0, count: targetLength - result.count)) + } else if result.count > targetLength { + // Truncate + result = Array(result.prefix(targetLength)) + } + return result + } +} diff --git a/Sources/FluidAudio/ASR/Parakeet/CtcZhCnModels.swift b/Sources/FluidAudio/ASR/Parakeet/CtcZhCnModels.swift new file mode 100644 index 000000000..8e901a79e --- /dev/null +++ b/Sources/FluidAudio/ASR/Parakeet/CtcZhCnModels.swift @@ -0,0 +1,265 @@ +@preconcurrency import CoreML +import Foundation + +/// Container for Parakeet CTC zh-CN CoreML models (full pipeline) +public struct CtcZhCnModels: Sendable { + + public let preprocessor: MLModel + public let encoder: MLModel + public let decoder: MLModel + public let configuration: MLModelConfiguration + public let vocabulary: [Int: String] + public let blankId: Int + + private static let logger = AppLogger(category: "CtcZhCnModels") + + public init( + preprocessor: MLModel, + encoder: MLModel, + decoder: MLModel, + configuration: MLModelConfiguration, + vocabulary: [Int: String], + blankId: Int = 7000 + ) { + self.preprocessor = preprocessor + self.encoder = encoder + self.decoder = decoder + self.configuration = configuration + self.vocabulary = vocabulary + self.blankId = blankId + } +} + +extension CtcZhCnModels { + + /// Load CTC zh-CN models from a directory. + /// + /// - Parameters: + /// - directory: Directory containing the downloaded CoreML bundles. + /// - useInt8Encoder: Whether to use int8 quantized encoder (default: true). + /// - configuration: Optional MLModel configuration. When nil, uses default configuration. + /// - progressHandler: Optional progress handler for model downloading. + /// - Returns: Loaded `CtcZhCnModels` instance. + public static func load( + from directory: URL, + useInt8Encoder: Bool = true, + configuration: MLModelConfiguration? = nil, + progressHandler: DownloadUtils.ProgressHandler? = nil + ) async throws -> CtcZhCnModels { + logger.info("Loading CTC zh-CN models from: \(directory.path)") + + let config = configuration ?? defaultConfiguration() + let parentDirectory = directory.deletingLastPathComponent() + + // Load preprocessor, encoder, and decoder + let encoderFileName = + useInt8Encoder + ? ModelNames.CTCZhCn.encoderFile + : ModelNames.CTCZhCn.encoderFp32File + + let modelNames = [ + ModelNames.CTCZhCn.preprocessorFile, + encoderFileName, + ModelNames.CTCZhCn.decoderFile, + ] + + let models = try await DownloadUtils.loadModels( + .parakeetCtcZhCn, + modelNames: modelNames, + directory: parentDirectory, + computeUnits: config.computeUnits, + progressHandler: progressHandler + ) + + guard + let preprocessorModel = models[ModelNames.CTCZhCn.preprocessorFile], + let encoderModel = models[encoderFileName], + let decoderModel = models[ModelNames.CTCZhCn.decoderFile] + else { + throw AsrModelsError.loadingFailed( + "Failed to load CTC zh-CN models (preprocessor, encoder, or decoder missing)" + ) + } + + logger.info("Loaded preprocessor, encoder (\(useInt8Encoder ? "int8" : "fp32")), and decoder") + + // Load vocabulary + let vocab = try loadVocabulary(from: directory) + + logger.info("Successfully loaded CTC zh-CN models with \(vocab.count) tokens") + + return CtcZhCnModels( + preprocessor: preprocessorModel, + encoder: encoderModel, + decoder: decoderModel, + configuration: config, + vocabulary: vocab, + blankId: 7000 + ) + } + + /// Download CTC zh-CN models to the default cache directory. + /// + /// - Parameters: + /// - directory: Custom cache directory (default: uses defaultCacheDirectory). + /// - useInt8Encoder: Whether to download int8 quantized encoder (default: true). + /// - downloadBothEncoders: If true, downloads both int8 and fp32 encoders (default: false). + /// - force: Whether to force re-download even if models exist. + /// - progressHandler: Optional progress handler for download progress. + /// - Returns: The directory where models were downloaded. + @discardableResult + public static func download( + to directory: URL? = nil, + useInt8Encoder: Bool = true, + downloadBothEncoders: Bool = false, + force: Bool = false, + progressHandler: DownloadUtils.ProgressHandler? = nil + ) async throws -> URL { + let targetDir = directory ?? defaultCacheDirectory() + logger.info("Preparing CTC zh-CN models at: \(targetDir.path)") + + let parentDir = targetDir.deletingLastPathComponent() + + if !force && modelsExist(at: targetDir) { + logger.info("CTC zh-CN models already present at: \(targetDir.path)") + return targetDir + } + + if force { + let fileManager = FileManager.default + if fileManager.fileExists(atPath: targetDir.path) { + try fileManager.removeItem(at: targetDir) + } + } + + // Download encoder variant(s) + let encoderFileName = + useInt8Encoder + ? ModelNames.CTCZhCn.encoderFile + : ModelNames.CTCZhCn.encoderFp32File + + var modelNames = [ + ModelNames.CTCZhCn.preprocessorFile, + encoderFileName, + ModelNames.CTCZhCn.decoderFile, + ] + + // Optionally download both encoder variants + if downloadBothEncoders { + let otherEncoder = + useInt8Encoder + ? ModelNames.CTCZhCn.encoderFp32File + : ModelNames.CTCZhCn.encoderFile + modelNames.append(otherEncoder) + } + + _ = try await DownloadUtils.loadModels( + .parakeetCtcZhCn, + modelNames: modelNames, + directory: parentDir, + progressHandler: progressHandler + ) + + logger.info("Successfully downloaded CTC zh-CN models") + return targetDir + } + + /// Convenience helper that downloads (if needed) and loads the CTC zh-CN models. + /// + /// - Parameters: + /// - directory: Custom cache directory (default: uses defaultCacheDirectory). + /// - useInt8Encoder: Whether to use int8 quantized encoder (default: true). + /// - configuration: Optional MLModel configuration. + /// - progressHandler: Optional progress handler. + /// - Returns: Loaded `CtcZhCnModels` instance. + public static func downloadAndLoad( + to directory: URL? = nil, + useInt8Encoder: Bool = true, + configuration: MLModelConfiguration? = nil, + progressHandler: DownloadUtils.ProgressHandler? = nil + ) async throws -> CtcZhCnModels { + let targetDir = try await download( + to: directory, + useInt8Encoder: useInt8Encoder, + progressHandler: progressHandler + ) + return try await load( + from: targetDir, + useInt8Encoder: useInt8Encoder, + configuration: configuration, + progressHandler: progressHandler + ) + } + + /// Default CoreML configuration for CTC zh-CN inference. + public static func defaultConfiguration() -> MLModelConfiguration { + MLModelConfigurationUtils.defaultConfiguration(computeUnits: .cpuAndNeuralEngine) + } + + /// Check whether required CTC zh-CN model bundles and vocabulary exist at a directory. + public static func modelsExist(at directory: URL) -> Bool { + let fileManager = FileManager.default + let repoPath = directory + + // Check if at least one encoder variant exists + let int8EncoderPath = repoPath.appendingPathComponent(ModelNames.CTCZhCn.encoderFile) + let fp32EncoderPath = repoPath.appendingPathComponent(ModelNames.CTCZhCn.encoderFp32File) + let encoderExists = + fileManager.fileExists(atPath: int8EncoderPath.path) + || fileManager.fileExists(atPath: fp32EncoderPath.path) + + let requiredFiles = [ + ModelNames.CTCZhCn.preprocessorFile, + ModelNames.CTCZhCn.decoderFile, + ] + + let modelsPresent = requiredFiles.allSatisfy { fileName in + let path = repoPath.appendingPathComponent(fileName) + return fileManager.fileExists(atPath: path.path) + } + + let vocabPath = repoPath.appendingPathComponent(ModelNames.CTCZhCn.vocabularyFile) + let vocabPresent = fileManager.fileExists(atPath: vocabPath.path) + + return encoderExists && modelsPresent && vocabPresent + } + + /// Default cache directory for CTC zh-CN models (within Application Support). + public static func defaultCacheDirectory() -> URL { + MLModelConfigurationUtils.defaultModelsDirectory(for: .parakeetCtcZhCn) + } + + /// Load vocabulary from vocab.json in the given directory. + private static func loadVocabulary(from directory: URL) throws -> [Int: String] { + let vocabPath = directory.appendingPathComponent(ModelNames.CTCZhCn.vocabularyFile) + guard FileManager.default.fileExists(atPath: vocabPath.path) else { + throw AsrModelsError.modelNotFound("vocab.json", vocabPath) + } + + let data = try Data(contentsOf: vocabPath) + + // Try parsing as array first (standard format: ["", "▁t", "he", ...]) + if let tokenArray = try? JSONSerialization.jsonObject(with: data) as? [String] { + var vocabulary: [Int: String] = [:] + for (index, token) in tokenArray.enumerated() { + vocabulary[index] = token + } + logger.info("Loaded CTC zh-CN vocabulary with \(vocabulary.count) tokens from \(vocabPath.path)") + return vocabulary + } + + // Fallback: try parsing as dictionary ({"0": "", "1": "▁t", ...}) + if let jsonDict = try? JSONSerialization.jsonObject(with: data) as? [String: String] { + var vocabulary: [Int: String] = [:] + for (key, value) in jsonDict { + if let tokenId = Int(key) { + vocabulary[tokenId] = value + } + } + logger.info("Loaded CTC zh-CN vocabulary with \(vocabulary.count) tokens from \(vocabPath.path)") + return vocabulary + } + + throw AsrModelsError.loadingFailed("Failed to parse vocab.json - expected array or dictionary format") + } +} diff --git a/Sources/FluidAudio/ModelNames.swift b/Sources/FluidAudio/ModelNames.swift index 5227c7712..a059323f8 100644 --- a/Sources/FluidAudio/ModelNames.swift +++ b/Sources/FluidAudio/ModelNames.swift @@ -7,6 +7,7 @@ public enum Repo: String, CaseIterable { case parakeetV2 = "FluidInference/parakeet-tdt-0.6b-v2-coreml" case parakeetCtc110m = "FluidInference/parakeet-ctc-110m-coreml" case parakeetCtc06b = "FluidInference/parakeet-ctc-0.6b-coreml" + case parakeetCtcZhCn = "FluidInference/parakeet-ctc-0.6b-zh-cn-coreml" case parakeetEou160 = "FluidInference/parakeet-realtime-eou-120m-coreml/160ms" case parakeetEou320 = "FluidInference/parakeet-realtime-eou-120m-coreml/320ms" case parakeetEou1280 = "FluidInference/parakeet-realtime-eou-120m-coreml/1280ms" @@ -35,6 +36,8 @@ public enum Repo: String, CaseIterable { return "parakeet-ctc-110m-coreml" case .parakeetCtc06b: return "parakeet-ctc-0.6b-coreml" + case .parakeetCtcZhCn: + return "parakeet-ctc-0.6b-zh-cn-coreml" case .parakeetEou160: return "parakeet-realtime-eou-120m-coreml/160ms" case .parakeetEou320: @@ -133,6 +136,8 @@ public enum Repo: String, CaseIterable { return "parakeet-ctc-110m-coreml" case .parakeetCtc06b: return "parakeet-ctc-0.6b-coreml" + case .parakeetCtcZhCn: + return "parakeet-ctc-zh-cn" case .parakeetTdtCtc110m: return "parakeet-tdt-ctc-110m" default: @@ -240,6 +245,30 @@ public enum ModelNames { ] } + /// CTC zh-CN model names (full pipeline: Preprocessor + Encoder + CTC Decoder) + public enum CTCZhCn { + public static let preprocessor = "Preprocessor" + public static let encoder = "Encoder-v2-int8" // Default to int8 quantized version + public static let encoderFp32 = "Encoder-v1-fp32" + public static let decoder = "Decoder" + + public static let preprocessorFile = preprocessor + ".mlmodelc" + public static let encoderFile = encoder + ".mlmodelc" + public static let encoderFp32File = encoderFp32 + ".mlmodelc" + public static let decoderFile = decoder + ".mlmodelc" + + // Vocabulary JSON path + public static let vocabularyFile = "vocab.json" + + // Download both encoder variants (int8 and fp32) so users can choose at runtime + public static let requiredModels: Set = [ + preprocessorFile, + encoderFile, // int8 encoder + encoderFp32File, // fp32 encoder + decoderFile, + ] + } + /// VAD model names public enum VAD { public static let sileroVad = "silero-vad-unified-256ms-v6.0.0" @@ -579,6 +608,8 @@ public enum ModelNames { return ModelNames.ASR.requiredModelsFused case .parakeetCtc110m, .parakeetCtc06b: return ModelNames.CTC.requiredModels + case .parakeetCtcZhCn: + return ModelNames.CTCZhCn.requiredModels case .parakeetEou160, .parakeetEou320, .parakeetEou1280: return ModelNames.ParakeetEOU.requiredModels case .nemotronStreaming1120, .nemotronStreaming560: diff --git a/Sources/FluidAudioCLI/Commands/ASR/CtcZhCnBenchmark.swift b/Sources/FluidAudioCLI/Commands/ASR/CtcZhCnBenchmark.swift new file mode 100644 index 000000000..62c28bd8d --- /dev/null +++ b/Sources/FluidAudioCLI/Commands/ASR/CtcZhCnBenchmark.swift @@ -0,0 +1,534 @@ +#if os(macOS) +import AVFoundation +import FluidAudio +import Foundation + +enum CtcZhCnBenchmark { + private static let logger = AppLogger(category: "CtcZhCnBenchmark") + + static func run(arguments: [String]) async { + var numSamples = 100 + var useInt8 = true + var outputFile: String? + var verbose = false + var datasetPath: String? + var autoDownload = false + + var i = 0 + while i < arguments.count { + let arg = arguments[i] + switch arg { + case "--samples", "-n": + if i + 1 < arguments.count { + numSamples = Int(arguments[i + 1]) ?? 100 + i += 1 + } + case "--fp32": + useInt8 = false + case "--int8": + useInt8 = true + case "--output", "-o": + if i + 1 < arguments.count { + outputFile = arguments[i + 1] + i += 1 + } + case "--dataset-path": + if i + 1 < arguments.count { + datasetPath = arguments[i + 1] + i += 1 + } + case "--auto-download": + autoDownload = true + case "--verbose", "-v": + verbose = true + case "--help", "-h": + printUsage() + return + default: + break + } + i += 1 + } + + logger.info("=== Parakeet CTC zh-CN Benchmark ===") + logger.info("Encoder: \(useInt8 ? "int8 (0.55GB)" : "fp32 (1.1GB)")") + logger.info("Samples: \(numSamples)") + logger.info("") + + do { + // Load models + logger.info("Loading CTC zh-CN models...") + let manager = try await CtcZhCnManager.load( + useInt8Encoder: useInt8, + progressHandler: verbose ? createProgressHandler() : nil + ) + logger.info("Models loaded successfully") + + // Load THCHS-30 dataset + logger.info("") + logger.info("Loading THCHS-30 test set...") + let samples = try await loadTHCHS30Samples( + maxSamples: numSamples, + datasetPath: datasetPath, + autoDownload: autoDownload + ) + logger.info("Loaded \(samples.count) samples") + + // Run benchmark + logger.info("") + logger.info("Running transcription benchmark...") + let results = try await runBenchmark(manager: manager, samples: samples) + + // Print results + printResults(results: results, encoderType: useInt8 ? "int8" : "fp32") + + // Save to JSON if requested + if let outputFile = outputFile { + try saveResults(results: results, outputFile: outputFile) + logger.info("") + logger.info("Results saved to: \(outputFile)") + } + + } catch { + logger.error("Benchmark failed: \(error.localizedDescription)") + if verbose { + logger.error("Error details: \(String(describing: error))") + } + } + } + + private struct BenchmarkSample { + let audioPath: String + let reference: String + let sampleId: Int + } + + private struct BenchmarkResult: Codable { + let sampleId: Int + let reference: String + let hypothesis: String + let normalizedRef: String + let normalizedHyp: String + let cer: Double + let latencyMs: Double + let audioDurationSec: Double + let rtfx: Double + } + + private struct MetadataEntry: Codable { + let file_name: String + let text: String + } + + private static func loadTHCHS30Samples( + maxSamples: Int, datasetPath: String?, autoDownload: Bool + ) async throws -> [BenchmarkSample] { + let baseDir: URL + + if let path = datasetPath { + // Use provided path + baseDir = URL(fileURLWithPath: path) + } else if autoDownload { + // Download from HuggingFace to cache directory + #if os(macOS) + let homeDir = FileManager.default.homeDirectoryForCurrentUser + let cacheDir = + homeDir + .appendingPathComponent("Library/Application Support/FluidAudio/Datasets/THCHS-30") + #else + let cacheDir = FileManager.default.temporaryDirectory + .appendingPathComponent("FluidAudio/Datasets/THCHS-30") + #endif + + try FileManager.default.createDirectory( + at: cacheDir, withIntermediateDirectories: true) + + logger.info("Downloading THCHS-30 from HuggingFace...") + try await downloadTHCHS30Dataset(to: cacheDir) + baseDir = cacheDir + } else { + throw NSError( + domain: "CtcZhCnBenchmark", + code: 1, + userInfo: [ + NSLocalizedDescriptionKey: + """ + THCHS-30 dataset not found. + + Options: + 1. Use --auto-download to download from HuggingFace + 2. Use --dataset-path to specify local dataset directory + + Expected directory structure: + / + ├── audio/ # WAV files + └── metadata.jsonl # Transcripts + """ + ] + ) + } + + // Load metadata.jsonl + let metadataPath = baseDir.appendingPathComponent("metadata.jsonl") + guard FileManager.default.fileExists(atPath: metadataPath.path) else { + throw NSError( + domain: "CtcZhCnBenchmark", + code: 2, + userInfo: [ + NSLocalizedDescriptionKey: + "metadata.jsonl not found at: \(metadataPath.path)" + ] + ) + } + + let metadataContent = try String(contentsOf: metadataPath, encoding: .utf8) + var samples: [BenchmarkSample] = [] + + for (index, line) in metadataContent.components(separatedBy: .newlines).enumerated() { + guard !line.isEmpty else { continue } + guard samples.count < maxSamples else { break } + + let decoder = JSONDecoder() + guard let data = line.data(using: .utf8), + let entry = try? decoder.decode(MetadataEntry.self, from: data) + else { + logger.warning("Failed to decode line \(index): \(line)") + continue + } + + let audioPath = baseDir.appendingPathComponent(entry.file_name).path + guard FileManager.default.fileExists(atPath: audioPath) else { + logger.warning("Audio file not found: \(audioPath)") + continue + } + + samples.append( + BenchmarkSample( + audioPath: audioPath, + reference: entry.text, + sampleId: index + )) + } + + return samples + } + + private static func downloadTHCHS30Dataset(to directory: URL) async throws { + // Download using git-lfs or HuggingFace Hub API + // For now, use a simple approach: shell out to huggingface-cli + let process = Process() + process.executableURL = URL(fileURLWithPath: "/usr/bin/env") + process.arguments = [ + "huggingface-cli", + "download", + "FluidInference/THCHS-30-tests", + "--repo-type", "dataset", + "--local-dir", directory.path, + ] + + try process.run() + process.waitUntilExit() + + guard process.terminationStatus == 0 else { + throw NSError( + domain: "CtcZhCnBenchmark", + code: 3, + userInfo: [ + NSLocalizedDescriptionKey: + """ + Failed to download THCHS-30 dataset from HuggingFace. + Make sure huggingface-cli is installed: pip install huggingface_hub + """ + ] + ) + } + } + + private static func runBenchmark( + manager: CtcZhCnManager, samples: [BenchmarkSample] + ) async throws -> [BenchmarkResult] { + var results: [BenchmarkResult] = [] + + for (index, sample) in samples.enumerated() { + let audioURL = URL(fileURLWithPath: sample.audioPath) + + let startTime = Date() + let hypothesis = try await manager.transcribe(audioURL: audioURL) + let elapsed = Date().timeIntervalSince(startTime) + + let normalizedRef = normalizeChineseText(sample.reference) + let normalizedHyp = normalizeChineseText(hypothesis) + + let cer = calculateCER(reference: normalizedRef, hypothesis: normalizedHyp) + + // Get audio duration + let audioFile = try AVAudioFile(forReading: audioURL) + let duration = Double(audioFile.length) / audioFile.processingFormat.sampleRate + + let rtfx = duration / elapsed + + let result = BenchmarkResult( + sampleId: sample.sampleId, + reference: sample.reference, + hypothesis: hypothesis, + normalizedRef: normalizedRef, + normalizedHyp: normalizedHyp, + cer: cer, + latencyMs: elapsed * 1000.0, + audioDurationSec: duration, + rtfx: rtfx + ) + + results.append(result) + + if (index + 1) % 10 == 0 { + logger.info("Processed \(index + 1)/\(samples.count) samples...") + } + } + + return results + } + + private static func normalizeChineseText(_ text: String) -> String { + var normalized = text + + // Remove Chinese punctuation (including curly quotes U+201C, U+201D, U+2018, U+2019) + let chinesePunct = "，。！？、；：\u{201C}\u{201D}\u{2018}\u{2019}" + for char in chinesePunct { + normalized = normalized.replacingOccurrences(of: String(char), with: "") + } + + // Remove Chinese brackets and quotes + let brackets = "「」『』（）《》【】" + for char in brackets { + normalized = normalized.replacingOccurrences(of: String(char), with: "") + } + + // Remove common symbols + let symbols = "…—·" + for char in symbols { + normalized = normalized.replacingOccurrences(of: String(char), with: "") + } + + // Remove English punctuation + let englishPunct = ",.!?;:()[]{}\\<>\"'-" + for char in englishPunct { + normalized = normalized.replacingOccurrences(of: String(char), with: "") + } + + // Convert Arabic digits to Chinese characters + let digitMap: [Character: String] = [ + "0": "零", + "1": "一", + "2": "二", + "3": "三", + "4": "四", + "5": "五", + "6": "六", + "7": "七", + "8": "八", + "9": "九", + ] + for (digit, chinese) in digitMap { + normalized = normalized.replacingOccurrences(of: String(digit), with: chinese) + } + + // Normalize whitespace and remove spaces + normalized = normalized.components(separatedBy: .whitespacesAndNewlines) + .filter { !$0.isEmpty } + .joined() + + return normalized.lowercased() + } + + private static func calculateCER(reference: String, hypothesis: String) -> Double { + let refChars = Array(reference) + let hypChars = Array(hypothesis) + + // Levenshtein distance + let distance = levenshteinDistance(refChars, hypChars) + + guard !refChars.isEmpty else { return hypChars.isEmpty ? 0.0 : 1.0 } + + return Double(distance) / Double(refChars.count) + } + + private static func levenshteinDistance(_ a: [T], _ b: [T]) -> Int { + let m = a.count + let n = b.count + + var dp = Array(repeating: Array(repeating: 0, count: n + 1), count: m + 1) + + for i in 0...m { + dp[i][0] = i + } + for j in 0...n { + dp[0][j] = j + } + + // Skip main loop if either array is empty (ranges 1...0 would be invalid) + guard m > 0 && n > 0 else { return dp[m][n] } + + for i in 1...m { + for j in 1...n { + if a[i - 1] == b[j - 1] { + dp[i][j] = dp[i - 1][j - 1] + } else { + dp[i][j] = min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1]) + 1 + } + } + } + + return dp[m][n] + } + + private static func printResults(results: [BenchmarkResult], encoderType: String) { + guard !results.isEmpty else { + logger.info("No results to display") + return + } + + let cers = results.map { $0.cer } + let latencies = results.map { $0.latencyMs } + let rtfxs = results.map { $0.rtfx } + + let meanCER = cers.reduce(0, +) / Double(cers.count) * 100.0 + let medianCER = median(cers) * 100.0 + let meanLatency = latencies.reduce(0, +) / Double(latencies.count) + let meanRTFx = rtfxs.reduce(0, +) / Double(rtfxs.count) + + logger.info("") + logger.info("=== Benchmark Results ===") + logger.info("Encoder: \(encoderType)") + logger.info("Samples: \(results.count)") + logger.info("") + logger.info("Mean CER: \(String(format: "%.2f", meanCER))%") + logger.info("Median CER: \(String(format: "%.2f", medianCER))%") + logger.info("Mean Latency: \(String(format: "%.1f", meanLatency))ms") + logger.info("Mean RTFx: \(String(format: "%.1f", meanRTFx))x") + + // CER distribution + let below5 = cers.filter { $0 < 0.05 }.count + let below10 = cers.filter { $0 < 0.10 }.count + let below20 = cers.filter { $0 < 0.20 }.count + + logger.info("") + logger.info("CER Distribution:") + logger.info( + " <5%: \(below5) samples (\(String(format: "%.1f", Double(below5) / Double(results.count) * 100.0))%)") + logger.info( + " <10%: \(below10) samples (\(String(format: "%.1f", Double(below10) / Double(results.count) * 100.0))%)") + logger.info( + " <20%: \(below20) samples (\(String(format: "%.1f", Double(below20) / Double(results.count) * 100.0))%)") + } + + private static func median(_ values: [Double]) -> Double { + let sorted = values.sorted() + let count = sorted.count + if count == 0 { return 0.0 } + if count % 2 == 0 { + return (sorted[count / 2 - 1] + sorted[count / 2]) / 2.0 + } else { + return sorted[count / 2] + } + } + + private struct BenchmarkOutput: Codable { + let summary: Summary + let results: [BenchmarkResult] + + struct Summary: Codable { + let mean_cer: Double + let median_cer: Double + let mean_latency_ms: Double + let mean_rtfx: Double + let total_samples: Int + let below_5_pct: Int + let below_10_pct: Int + let below_20_pct: Int + } + } + + private static func saveResults(results: [BenchmarkResult], outputFile: String) throws { + guard !results.isEmpty else { + logger.warning("No results to save") + return + } + + let cers = results.map { $0.cer } + let latencies = results.map { $0.latencyMs } + let rtfxs = results.map { $0.rtfx } + + let summary = BenchmarkOutput.Summary( + mean_cer: cers.reduce(0, +) / Double(cers.count), + median_cer: median(cers), + mean_latency_ms: latencies.reduce(0, +) / Double(latencies.count), + mean_rtfx: rtfxs.reduce(0, +) / Double(rtfxs.count), + total_samples: results.count, + below_5_pct: cers.filter { $0 < 0.05 }.count, + below_10_pct: cers.filter { $0 < 0.10 }.count, + below_20_pct: cers.filter { $0 < 0.20 }.count + ) + + let output = BenchmarkOutput(summary: summary, results: results) + let encoder = JSONEncoder() + encoder.outputFormatting = .prettyPrinted + let jsonData = try encoder.encode(output) + try jsonData.write(to: URL(fileURLWithPath: outputFile)) + } + + private static func createProgressHandler() -> DownloadUtils.ProgressHandler { + return { progress in + let percentage = progress.fractionCompleted * 100.0 + switch progress.phase { + case .listing: + logger.info("Listing files from repository...") + case .downloading(let completed, let total): + logger.info( + "Downloading models: \(completed)/\(total) files (\(String(format: "%.1f", percentage))%)" + ) + case .compiling(let modelName): + logger.info("Compiling \(modelName)...") + } + } + } + + private static func printUsage() { + logger.info( + """ + CTC zh-CN Benchmark - Measure Character Error Rate on THCHS-30 dataset + + Usage: fluidaudiocli ctc-zh-cn-benchmark [options] + + Options: + --samples, -n Number of samples to test (default: 100) + --int8 Use int8 quantized encoder (default) + --fp32 Use fp32 encoder + --output, -o Save results to JSON file + --dataset-path Path to THCHS-30 dataset directory + --auto-download Download THCHS-30 from HuggingFace (requires huggingface-cli) + --verbose, -v Show download progress + --help, -h Show this help message + + Examples: + # Auto-download from HuggingFace + fluidaudiocli ctc-zh-cn-benchmark --auto-download --samples 100 + + # Use local dataset + fluidaudiocli ctc-zh-cn-benchmark --dataset-path ./thchs30_test_hf + + # Save results to JSON + fluidaudiocli ctc-zh-cn-benchmark --auto-download --output results.json + + Expected Results (THCHS-30, 100 samples): + Int8 encoder: 8.37% mean CER, 6.67% median CER + FP32 encoder: Similar performance + + Dataset: FluidInference/THCHS-30-tests on HuggingFace + 2,495 Mandarin Chinese test utterances from THCHS-30 corpus + """ + ) + } +} + +#endif diff --git a/Sources/FluidAudioCLI/Commands/ASR/CtcZhCnTranscribeCommand.swift b/Sources/FluidAudioCLI/Commands/ASR/CtcZhCnTranscribeCommand.swift new file mode 100644 index 000000000..3e5cf6b5f --- /dev/null +++ b/Sources/FluidAudioCLI/Commands/ASR/CtcZhCnTranscribeCommand.swift @@ -0,0 +1,130 @@ +#if os(macOS) +import AVFoundation +import FluidAudio +import Foundation + +enum CtcZhCnTranscribeCommand { + private static let logger = AppLogger(category: "CtcZhCnTranscribe") + + static func run(arguments: [String]) async { + // Parse arguments + var audioPath: String? + var useInt8 = true + var verbose = false + + var i = 0 + while i < arguments.count { + let arg = arguments[i] + switch arg { + case "--fp32": + useInt8 = false + case "--int8": + useInt8 = true + case "--verbose", "-v": + verbose = true + case "--help", "-h": + printUsage() + return + default: + if audioPath == nil { + audioPath = arg + } + } + i += 1 + } + + guard let audioPath = audioPath else { + logger.error("Error: No audio file specified") + printUsage() + return + } + + let audioURL = URL(fileURLWithPath: audioPath) + guard FileManager.default.fileExists(atPath: audioURL.path) else { + logger.error("Error: Audio file not found: \(audioPath)") + return + } + + do { + logger.info("Loading CTC zh-CN models (encoder: \(useInt8 ? "int8" : "fp32"))...") + + let manager = try await CtcZhCnManager.load( + useInt8Encoder: useInt8, + progressHandler: verbose ? createProgressHandler() : nil + ) + + logger.info("Transcribing: \(audioPath)") + + let startTime = Date() + let text = try await manager.transcribe(audioURL: audioURL) + let elapsed = Date().timeIntervalSince(startTime) + + logger.info("Transcription completed in \(String(format: "%.2f", elapsed))s") + logger.info("") + logger.info("Result:") + print(text) + + } catch { + logger.error("Transcription failed: \(error.localizedDescription)") + if verbose { + logger.error("Error details: \(String(describing: error))") + } + } + } + + private static func createProgressHandler() -> DownloadUtils.ProgressHandler { + return { progress in + let percentage = progress.fractionCompleted * 100.0 + switch progress.phase { + case .listing: + logger.info("Listing files from repository...") + case .downloading(let completed, let total): + logger.info( + "Downloading models: \(completed)/\(total) files (\(String(format: "%.1f", percentage))%)" + ) + case .compiling(let modelName): + logger.info("Compiling \(modelName)...") + } + } + } + + private static func printUsage() { + logger.info( + """ + CTC zh-CN Transcribe - Mandarin Chinese speech recognition + + Usage: fluidaudiocli ctc-zh-cn-transcribe [options] + + Arguments: + Path to audio file (WAV, MP3, etc.) + + Options: + --int8 Use int8 quantized encoder (default, faster) + --fp32 Use fp32 encoder (higher precision) + --verbose, -v Show download progress and detailed logs + --help, -h Show this help message + + Examples: + # Basic transcription + fluidaudiocli ctc-zh-cn-transcribe audio.wav + + # Use fp32 encoder for higher precision + fluidaudiocli ctc-zh-cn-transcribe audio.wav --fp32 + + Model Info: + - Language: Mandarin Chinese (Simplified, zh-CN) + - Vocabulary: 7000 SentencePiece tokens + - Max audio: 15 seconds (longer audio is truncated) + - Int8 encoder: 0.55GB (recommended) + - FP32 encoder: 1.1GB + + Performance (FLEURS 100 samples): + - Int8 encoder: 10.54% CER + - FP32 encoder: 10.45% CER + + Note: Models auto-download from HuggingFace on first use. + """ + ) + } +} +#endif diff --git a/Sources/FluidAudioCLI/Commands/ASR/Parakeet/SlidingWindow/AsrBenchmark.swift b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/SlidingWindow/AsrBenchmark.swift index a9aab9c6e..ce551f7ad 100644 --- a/Sources/FluidAudioCLI/Commands/ASR/Parakeet/SlidingWindow/AsrBenchmark.swift +++ b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/SlidingWindow/AsrBenchmark.swift @@ -842,6 +842,7 @@ extension ASRBenchmark { case .v2: versionLabel = "v2" case .v3: versionLabel = "v3" case .tdtCtc110m: versionLabel = "tdt-ctc-110m" + case .ctcZhCn: versionLabel = "ctc-zh-cn" } logger.info(" Model version: \(versionLabel)") logger.info(" Debug mode: \(debugMode ? "enabled" : "disabled")") diff --git a/Sources/FluidAudioCLI/Commands/ASR/Parakeet/SlidingWindow/TranscribeCommand.swift b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/SlidingWindow/TranscribeCommand.swift index c07f21d2e..18f87326b 100644 --- a/Sources/FluidAudioCLI/Commands/ASR/Parakeet/SlidingWindow/TranscribeCommand.swift +++ b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/SlidingWindow/TranscribeCommand.swift @@ -430,6 +430,7 @@ enum TranscribeCommand { case .v2: modelVersionLabel = "v2" case .v3: modelVersionLabel = "v3" case .tdtCtc110m: modelVersionLabel = "tdt-ctc-110m" + case .ctcZhCn: modelVersionLabel = "ctc-zh-cn" } let output = TranscriptionJSONOutput( audioFile: audioFile, @@ -684,6 +685,7 @@ enum TranscribeCommand { case .v2: modelVersionLabel = "v2" case .v3: modelVersionLabel = "v3" case .tdtCtc110m: modelVersionLabel = "tdt-ctc-110m" + case .ctcZhCn: modelVersionLabel = "ctc-zh-cn" } let output = TranscriptionJSONOutput( audioFile: audioFile, diff --git a/Sources/FluidAudioCLI/FluidAudioCLI.swift b/Sources/FluidAudioCLI/FluidAudioCLI.swift index 0b226ac51..0714a6f64 100644 --- a/Sources/FluidAudioCLI/FluidAudioCLI.swift +++ b/Sources/FluidAudioCLI/FluidAudioCLI.swift @@ -70,6 +70,10 @@ struct FluidAudioCLI { await NemotronBenchmark.run(arguments: Array(arguments.dropFirst(2))) case "nemotron-transcribe": await NemotronTranscribe.run(arguments: Array(arguments.dropFirst(2))) + case "ctc-zh-cn-transcribe": + await CtcZhCnTranscribeCommand.run(arguments: Array(arguments.dropFirst(2))) + case "ctc-zh-cn-benchmark": + await CtcZhCnBenchmark.run(arguments: Array(arguments.dropFirst(2))) case "help", "--help", "-h": printUsage() default: @@ -107,6 +111,8 @@ struct FluidAudioCLI { g2p-benchmark Run multilingual G2P benchmark nemotron-benchmark Run Nemotron 0.6B streaming ASR benchmark nemotron-transcribe Transcribe custom audio files with Nemotron + ctc-zh-cn-transcribe Transcribe Mandarin Chinese audio with Parakeet CTC + ctc-zh-cn-benchmark Run CTC zh-CN benchmark on THCHS-30 dataset download Download evaluation datasets help Show this help message diff --git a/Tests/FluidAudioTests/ASR/Parakeet/CtcZhCnTests.swift b/Tests/FluidAudioTests/ASR/Parakeet/CtcZhCnTests.swift new file mode 100644 index 000000000..06d7c2c3e --- /dev/null +++ b/Tests/FluidAudioTests/ASR/Parakeet/CtcZhCnTests.swift @@ -0,0 +1,331 @@ +import Foundation +import XCTest + +@testable import FluidAudio + +/// Unit tests for CTC zh-CN text normalization and CER calculation +/// +/// These tests verify the pure functions used in CTC zh-CN benchmarking: +/// - Text normalization (punctuation removal, digit conversion, whitespace handling) +/// - Character Error Rate (CER) calculation +/// - Levenshtein distance algorithm +final class CtcZhCnTests: XCTestCase { + + // MARK: - Text Normalization Tests + + func testNormalizeChineseText_RemovesChinesePunctuation() { + let input = "你好，世界！这是、一个：测试。" + let expected = "你好世界这是一个测试" + + // Access via reflection since normalizeChineseText is private in CtcZhCnBenchmark + // For testing purposes, we'll test the logic inline + var normalized = input + + // Remove Chinese punctuation (including curly quotes U+201C, U+201D, U+2018, U+2019) + let chinesePunct = "，。！？、；：\u{201C}\u{201D}\u{2018}\u{2019}" + for char in chinesePunct { + normalized = normalized.replacingOccurrences(of: String(char), with: "") + } + + XCTAssertEqual(normalized, expected, "Should remove all Chinese punctuation") + } + + func testNormalizeChineseText_RemovesEnglishPunctuation() { + let input = "Hello, world! This is a test." + var normalized = input + + let englishPunct = ",.!?;:()[]{}\\<>\"'-" + for char in englishPunct { + normalized = normalized.replacingOccurrences(of: String(char), with: "") + } + + let expected = "Hello world This is a test" + XCTAssertEqual(normalized, expected, "Should remove all English punctuation") + } + + func testNormalizeChineseText_ConvertsDigitsToChineseCharacters() { + let input = "2021年8月15日" + var normalized = input + + let digitMap: [Character: String] = [ + "0": "零", + "1": "一", + "2": "二", + "3": "三", + "4": "四", + "5": "五", + "6": "六", + "7": "七", + "8": "八", + "9": "九", + ] + for (digit, chinese) in digitMap { + normalized = normalized.replacingOccurrences(of: String(digit), with: chinese) + } + + let expected = "二零二一年八月一五日" + XCTAssertEqual( + normalized, expected, + "Should convert Arabic digits to Chinese characters") + } + + func testNormalizeChineseText_RemovesBracketsAndQuotes() { + let input = "「你好」『世界』（测试）《书名》【注释】" + var normalized = input + + let brackets = "「」『』（）《》【】" + for char in brackets { + normalized = normalized.replacingOccurrences(of: String(char), with: "") + } + + let expected = "你好世界测试书名注释" + XCTAssertEqual(normalized, expected, "Should remove all brackets and quotation marks") + } + + func testNormalizeChineseText_NormalizesWhitespace() { + let input = "你好 世界\n这是\t测试" + let normalized = + input.components(separatedBy: .whitespacesAndNewlines) + .filter { !$0.isEmpty } + .joined() + + let expected = "你好世界这是测试" + XCTAssertEqual(normalized, expected, "Should normalize and remove all whitespace") + } + + func testNormalizeChineseText_CompleteExample() { + // This mimics the exact normalization logic from CtcZhCnBenchmark + let input = "桥下垂直净空15米，该项目于2011年8月完工。" + var normalized = input + + // Remove Chinese punctuation (including curly quotes U+201C, U+201D, U+2018, U+2019) + let chinesePunct = "，。！？、；：\u{201C}\u{201D}\u{2018}\u{2019}" + for char in chinesePunct { + normalized = normalized.replacingOccurrences(of: String(char), with: "") + } + + // Remove Chinese brackets and quotes + let brackets = "「」『』（）《》【】" + for char in brackets { + normalized = normalized.replacingOccurrences(of: String(char), with: "") + } + + // Remove common symbols + let symbols = "…—·" + for char in symbols { + normalized = normalized.replacingOccurrences(of: String(char), with: "") + } + + // Remove English punctuation + let englishPunct = ",.!?;:()[]{}\\<>\"'-" + for char in englishPunct { + normalized = normalized.replacingOccurrences(of: String(char), with: "") + } + + // Convert Arabic digits to Chinese characters + let digitMap: [Character: String] = [ + "0": "零", + "1": "一", + "2": "二", + "3": "三", + "4": "四", + "5": "五", + "6": "六", + "7": "七", + "8": "八", + "9": "九", + ] + for (digit, chinese) in digitMap { + normalized = normalized.replacingOccurrences(of: String(digit), with: chinese) + } + + // Normalize whitespace + normalized = + normalized.components(separatedBy: .whitespacesAndNewlines) + .filter { !$0.isEmpty } + .joined() + + let expected = "桥下垂直净空一五米该项目于二零一一年八月完工" + XCTAssertEqual( + normalized, expected, + "Full normalization should match expected output") + } + + // MARK: - Levenshtein Distance Tests + + func testLevenshteinDistance_IdenticalStrings() { + let a = Array("hello") + let b = Array("hello") + let distance = levenshteinDistance(a, b) + + XCTAssertEqual(distance, 0, "Identical strings should have distance 0") + } + + func testLevenshteinDistance_EmptyStrings() { + let a: [Character] = [] + let b: [Character] = [] + let distance = levenshteinDistance(a, b) + + XCTAssertEqual(distance, 0, "Empty strings should have distance 0") + } + + func testLevenshteinDistance_OneEmpty() { + let a = Array("hello") + let b: [Character] = [] + let distance = levenshteinDistance(a, b) + + XCTAssertEqual(distance, 5, "Distance should equal length of non-empty string") + } + + func testLevenshteinDistance_SingleSubstitution() { + let a = Array("kitten") + let b = Array("sitten") + let distance = levenshteinDistance(a, b) + + XCTAssertEqual(distance, 1, "Single substitution should have distance 1") + } + + func testLevenshteinDistance_SingleInsertion() { + let a = Array("cat") + let b = Array("cats") + let distance = levenshteinDistance(a, b) + + XCTAssertEqual(distance, 1, "Single insertion should have distance 1") + } + + func testLevenshteinDistance_SingleDeletion() { + let a = Array("cats") + let b = Array("cat") + let distance = levenshteinDistance(a, b) + + XCTAssertEqual(distance, 1, "Single deletion should have distance 1") + } + + func testLevenshteinDistance_ClassicExample() { + let a = Array("kitten") + let b = Array("sitting") + let distance = levenshteinDistance(a, b) + + XCTAssertEqual(distance, 3, "Classic kitten->sitting should have distance 3") + } + + func testLevenshteinDistance_ChineseCharacters() { + let a = Array("你好世界") + let b = Array("你好地球") + let distance = levenshteinDistance(a, b) + + XCTAssertEqual(distance, 2, "Two character substitutions should have distance 2") + } + + // MARK: - CER Calculation Tests + + func testCalculateCER_IdenticalStrings() { + let reference = "你好世界" + let hypothesis = "你好世界" + let cer = calculateCER(reference: reference, hypothesis: hypothesis) + + XCTAssertEqual(cer, 0.0, accuracy: 0.001, "Identical strings should have CER 0") + } + + func testCalculateCER_EmptyReference() { + let reference = "" + let hypothesis = "你好" + let cer = calculateCER(reference: reference, hypothesis: hypothesis) + + XCTAssertEqual(cer, 1.0, accuracy: 0.001, "Empty reference with non-empty hypothesis should have CER 1.0") + } + + func testCalculateCER_EmptyHypothesis() { + let reference = "你好" + let hypothesis = "" + let cer = calculateCER(reference: reference, hypothesis: hypothesis) + + XCTAssertEqual(cer, 1.0, accuracy: 0.001, "Non-empty reference with empty hypothesis should have CER 1.0") + } + + func testCalculateCER_BothEmpty() { + let reference = "" + let hypothesis = "" + let cer = calculateCER(reference: reference, hypothesis: hypothesis) + + XCTAssertEqual(cer, 0.0, accuracy: 0.001, "Both empty should have CER 0") + } + + func testCalculateCER_SingleCharacterError() { + let reference = "你好世界" // 4 characters + let hypothesis = "你好地界" // 1 substitution + let cer = calculateCER(reference: reference, hypothesis: hypothesis) + + // Distance = 1, Length = 4, CER = 1/4 = 0.25 + XCTAssertEqual(cer, 0.25, accuracy: 0.001, "Single character error in 4 chars should be 0.25") + } + + func testCalculateCER_MultipleErrors() { + let reference = "你好世界今天" // 6 characters + let hypothesis = "你好地球昨天" // 3 substitutions + let cer = calculateCER(reference: reference, hypothesis: hypothesis) + + // Distance = 3, Length = 6, CER = 3/6 = 0.5 + XCTAssertEqual(cer, 0.5, accuracy: 0.001, "3 errors in 6 chars should be 0.5") + } + + func testCalculateCER_InsertionErrors() { + let reference = "你好" // 2 characters + let hypothesis = "你好世界" // 2 insertions + let cer = calculateCER(reference: reference, hypothesis: hypothesis) + + // Distance = 2, Length = 2, CER = 2/2 = 1.0 + XCTAssertEqual(cer, 1.0, accuracy: 0.001, "2 insertions in 2 chars should be 1.0") + } + + func testCalculateCER_DeletionErrors() { + let reference = "你好世界" // 4 characters + let hypothesis = "你好" // 2 deletions + let cer = calculateCER(reference: reference, hypothesis: hypothesis) + + // Distance = 2, Length = 4, CER = 2/4 = 0.5 + XCTAssertEqual(cer, 0.5, accuracy: 0.001, "2 deletions in 4 chars should be 0.5") + } + + // MARK: - Helper Functions (matching CtcZhCnBenchmark implementation) + + private func levenshteinDistance(_ a: [T], _ b: [T]) -> Int { + let m = a.count + let n = b.count + + var dp = Array(repeating: Array(repeating: 0, count: n + 1), count: m + 1) + + for i in 0...m { + dp[i][0] = i + } + for j in 0...n { + dp[0][j] = j + } + + // Skip main loop if either array is empty (ranges 1...0 would be invalid) + guard m > 0 && n > 0 else { return dp[m][n] } + + for i in 1...m { + for j in 1...n { + if a[i - 1] == b[j - 1] { + dp[i][j] = dp[i - 1][j - 1] + } else { + dp[i][j] = min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1]) + 1 + } + } + } + + return dp[m][n] + } + + private func calculateCER(reference: String, hypothesis: String) -> Double { + let refChars = Array(reference) + let hypChars = Array(hypothesis) + + let distance = levenshteinDistance(refChars, hypChars) + + guard !refChars.isEmpty else { return hypChars.isEmpty ? 0.0 : 1.0 } + + return Double(distance) / Double(refChars.count) + } +}