diff --git a/Package.swift b/Package.swift index 4acbe23..3b647c8 100644 --- a/Package.swift +++ b/Package.swift @@ -8,7 +8,7 @@ let package = Package( .library(name: "TranscriberCore", targets: ["TranscriberCore"]) ], dependencies: [ - .package(url: "https://github.com/Blaizzy/mlx-audio-swift.git", revision: "7734cd1fbbe86460083c1d24199737a24cadfcc8") + .package(url: "https://github.com/Newarr/mlx-audio-swift.git", revision: "b8ec43083e4c5535594dbf9274893f9e6fe4a506") ], targets: [ .target( diff --git a/Sources/TranscriberCore/Engines/CohereMLXBackend.swift b/Sources/TranscriberCore/Engines/CohereMLXBackend.swift index 10af77c..e6fc127 100644 --- a/Sources/TranscriberCore/Engines/CohereMLXBackend.swift +++ b/Sources/TranscriberCore/Engines/CohereMLXBackend.swift @@ -69,6 +69,57 @@ public protocol CohereMLXAudioDurationReading: Sendable { func durationSeconds(for audioURL: URL) async throws -> Double } +/// Terminal errors surfaced by `CohereMLXBackend.transcribe(_:)` that the +/// recovery layer must NOT treat as transient. The MLX path runs locally with +/// deterministic params, so a degenerate decode is a true failure — retrying +/// without changing inputs would loop again. `TranscriptionWorker.isTransient` +/// does not classify these, so they reach the user as a clear failed status. +public enum CohereMLXBackendError: Error, Sendable, Equatable { + case degenerateOutput(reason: String, sample: String) +} + +/// Heuristic guard against the Cohere/MLX decode-loop failure mode where the +/// model spends the entire 583-second recording emitting the same phrase +/// dozens of times. The two checks below are pure-Swift and cheap; either one +/// firing means the output is unusable as a transcript. +/// +/// Compression-ratio (gzip) is intentionally omitted — it would require a new +/// dependency to catch a failure mode already covered by the two checks here. +/// Add it later only if the field shows misses. +enum DegenerateOutputDetector { + /// Returns `nil` when the text looks like a healthy transcript; otherwise + /// returns a short human-readable reason that callers can put in an error + /// payload or log line. Texts shorter than 30 words always return `nil` — + /// the worker has a separate "No speech detected" terminal path for those. + static func evaluate(_ text: String) -> String? { + let words = text + .split(whereSeparator: { $0.isWhitespace }) + .map(String.init) + guard words.count >= 30 else { return nil } + + if words.count >= 3 { + var counts: [String: Int] = [:] + let total = words.count - 2 + for i in 0.. 0.08 { + return "tri-gram \"\(top)\" repeats \(n)/\(total) times" + } + } + + let unique = Set(words.map { $0.lowercased() }).count + let fraction = Double(unique) / Double(words.count) + if fraction < 0.10 { + return "unique-word fraction \(String(format: "%.3f", fraction)) below 0.10" + } + + return nil + } +} + public final class CohereMLXBackend: TranscriptionEngine, @unchecked Sendable { public static let modelID = "beshkenadze/cohere-transcribe-03-2026-mlx-fp16" public static let defaultRequestModelID = modelID @@ -87,6 +138,17 @@ public final class CohereMLXBackend: TranscriptionEngine, @unchecked Sendable { public static let inferenceSampleRate = 16_000 public static let inferenceChannelCount = 1 + // Inference parameter overrides. The upstream `mlx-audio-swift` defaults + // (`chunkDuration=1200`, `repetitionPenalty=1.0`) are unsafe for real + // recordings: the model card documents 35 s training distribution and + // greedy decoding without a penalty collapses into loop output on longer + // audio. These values are deliberately wrapper-side so the upstream fork + // can stay close to Blaizzy/mlx-audio-swift. + public static let inferenceChunkDurationSeconds: Float = 30.0 + public static let inferenceMinChunkDurationSeconds: Float = 1.0 + public static let inferenceRepetitionPenalty: Float = 1.2 + public static let inferenceRepetitionContextSize: Int = 32 + private let adapter: any CohereMLXTranscribing private let durationReader: any CohereMLXAudioDurationReading private let modelDirectoryURL: URL @@ -115,6 +177,11 @@ public final class CohereMLXBackend: TranscriptionEngine, @unchecked Sendable { audioDurationSeconds: duration ) let output = try await adapter.transcribe(localRequest) + if let reason = DegenerateOutputDetector.evaluate(output.text) { + let sample = String(output.text.prefix(120)) + Log.engine.error("Cohere MLX produced degenerate output: \(reason, privacy: .public)") + throw CohereMLXBackendError.degenerateOutput(reason: reason, sample: sample) + } let utterances: [EngineResponse.Utterance] if output.segments.isEmpty { utterances = [EngineResponse.Utterance( @@ -178,18 +245,9 @@ public struct NativeCohereMLXAdapter: CohereMLXTranscribing { public func transcribe(_ request: CohereMLXAdapterRequest) async throws -> CohereMLXAdapterResponse { let (_, audio) = try loadAudioArray(from: request.audioURL, sampleRate: request.inputSampleRate) let model = try CohereTranscribeModel.fromDirectory(request.modelDirectoryURL) - var parameters = model.defaultGenerationParameters - parameters = STTGenerateParameters( - maxTokens: parameters.maxTokens, - temperature: parameters.temperature, - topP: parameters.topP, - topK: parameters.topK, - verbose: parameters.verbose, - language: request.languageCode, - chunkDuration: parameters.chunkDuration, - minChunkDuration: parameters.minChunkDuration, - repetitionPenalty: parameters.repetitionPenalty, - repetitionContextSize: parameters.repetitionContextSize + let parameters = Self.makeGenerationParameters( + modelDefaults: model.defaultGenerationParameters, + languageCode: request.languageCode ) let output = model.generate(audio: audio, generationParameters: parameters) return CohereMLXAdapterResponse( @@ -199,6 +257,28 @@ public struct NativeCohereMLXAdapter: CohereMLXTranscribing { ) } + /// Builds the generation parameters used at inference time. Exposed as a + /// pure static helper so tests can verify the wrapper overrides upstream + /// defaults without instantiating MLX. Only `maxTokens` is taken from the + /// model defaults; the other fields are pinned to wrapper-side constants. + static func makeGenerationParameters( + modelDefaults: STTGenerateParameters, + languageCode: String + ) -> STTGenerateParameters { + STTGenerateParameters( + maxTokens: modelDefaults.maxTokens, + temperature: 0.0, + topP: 1.0, + topK: 0, + verbose: false, + language: languageCode, + chunkDuration: CohereMLXBackend.inferenceChunkDurationSeconds, + minChunkDuration: CohereMLXBackend.inferenceMinChunkDurationSeconds, + repetitionPenalty: CohereMLXBackend.inferenceRepetitionPenalty, + repetitionContextSize: CohereMLXBackend.inferenceRepetitionContextSize + ) + } + public static func segments(from rawSegments: [[String: Any]]?) -> [CohereMLXSegment] { guard let rawSegments else { return [] } return rawSegments.compactMap { raw in diff --git a/Tests/TranscriberCoreTests/Engines/CohereRustBackendTests.swift b/Tests/TranscriberCoreTests/Engines/CohereRustBackendTests.swift index 2958368..0512978 100644 --- a/Tests/TranscriberCoreTests/Engines/CohereRustBackendTests.swift +++ b/Tests/TranscriberCoreTests/Engines/CohereRustBackendTests.swift @@ -1,5 +1,6 @@ import XCTest import AVFoundation +import MLXAudioSTT @testable import TranscriberCore final class CohereMLXBackendTests: XCTestCase { @@ -268,6 +269,148 @@ final class CohereMLXBackendTests: XCTestCase { XCTAssertFalse(source.localizedCaseInsensitiveContains(forbidden), "Local Cohere/MLX backend must not depend on subprocess inference: \(forbidden)") } } + + // MARK: - Decode-loop fix regression guards + + func testGenerationParametersOverrideUpstreamDefaults() { + // Fabricate the upstream model defaults exactly as Blaizzy/mlx-audio-swift + // returns them today: chunkDuration=1200, repetitionPenalty=1.0. The + // wrapper MUST replace those before generation runs. + let upstreamDefaults = STTGenerateParameters( + maxTokens: 1024, + temperature: 0.0, + topP: 1.0, + topK: 0, + verbose: false, + language: "en", + chunkDuration: 1200.0, + minChunkDuration: 1.0, + repetitionPenalty: 1.0, + repetitionContextSize: 32 + ) + + let parameters = NativeCohereMLXAdapter.makeGenerationParameters( + modelDefaults: upstreamDefaults, + languageCode: "en" + ) + + XCTAssertEqual(parameters.chunkDuration, CohereMLXBackend.inferenceChunkDurationSeconds) + XCTAssertEqual(parameters.chunkDuration, 30.0, "30 s matches the Cohere model card's training distribution; 1200 s loops") + XCTAssertEqual(parameters.minChunkDuration, CohereMLXBackend.inferenceMinChunkDurationSeconds) + XCTAssertEqual(parameters.repetitionPenalty, CohereMLXBackend.inferenceRepetitionPenalty) + XCTAssertEqual(parameters.repetitionPenalty, 1.2, "Penalty must be active to break the degenerate decode") + XCTAssertEqual(parameters.repetitionContextSize, CohereMLXBackend.inferenceRepetitionContextSize) + XCTAssertEqual(parameters.temperature, 0.0, "Local transcription stays deterministic") + XCTAssertEqual(parameters.topP, 1.0) + XCTAssertEqual(parameters.topK, 0) + XCTAssertEqual(parameters.language, "en") + XCTAssertEqual(parameters.maxTokens, 1024, "maxTokens is the only value taken from model defaults") + } + + func testDegenerateOutputDetectorFlagsRepetitiveTranscripts() { + let looped = String(repeating: "I think that's what I'm hearing ", count: 50) + let loopedReason = DegenerateOutputDetector.evaluate(looped) + XCTAssertNotNil(loopedReason, "Detector must catch the observed 'I think that's what I'm hearing' loop") + XCTAssertTrue(loopedReason?.contains("tri-gram") ?? false, "Reason should identify the dominant tri-gram, got: \(loopedReason ?? "nil")") + + // A clean ≥30-word sentence with healthy unique-word distribution. + // Cribbed from the pangram zoo plus padding to exceed the 30-word floor. + let healthy = """ + The quick brown fox jumps over the lazy dog while sphinx of black quartz \ + judges my vow and pack my box with five dozen liquor jugs as vexingly \ + quick daft zebras jump over a chilled fence near the meadow at dawn. + """ + XCTAssertNil(DegenerateOutputDetector.evaluate(healthy), "Detector must not fire on natural diverse text") + + // Fixture used by the cloud backend — guards against the detector + // accidentally flagging legitimate short transcripts. + let fixtureURL = URL(fileURLWithPath: #filePath) + .deletingLastPathComponent() + .appendingPathComponent("Fixtures/elevenlabs-success.json") + let fixtureData = try? Data(contentsOf: fixtureURL) + let fixtureText = (try? JSONSerialization.jsonObject(with: fixtureData ?? Data())) as? [String: Any] + let cloudText = (fixtureText?["text"] as? String) ?? "" + XCTAssertFalse(cloudText.isEmpty, "Fixture must still be present") + XCTAssertNil(DegenerateOutputDetector.evaluate(cloudText), "Detector must not fire on cloud fixture text") + + XCTAssertNil(DegenerateOutputDetector.evaluate(""), "Empty input is handled by the worker's 'no speech' path") + XCTAssertNil(DegenerateOutputDetector.evaluate("only a few words here"), "Below 30-word floor") + } + + /// Real-MLX integration: load the actual model and run it against the + /// recording that surfaced the loop bug. Skipped unless + /// `SCRIBE_RUN_MLX_INTEGRATION=1` is set in the environment because it + /// needs the model weights downloaded and takes ~minutes to complete. + func testFailingRecordingTranscribesWithoutDegenerationIntegration() async throws { + try XCTSkipUnless( + ProcessInfo.processInfo.environment["SCRIBE_RUN_MLX_INTEGRATION"] == "1", + "Integration test; set SCRIBE_RUN_MLX_INTEGRATION=1 to run" + ) + + let audioURL = FileManager.default + .homeDirectoryForCurrentUser + .appendingPathComponent("Scribe/2026-05-11-1756/audio.m4a") + try XCTSkipUnless( + FileManager.default.fileExists(atPath: audioURL.path), + "Real recording fixture missing at \(audioURL.path)" + ) + + let modelDir = CohereMLXBackend.defaultModelDirectoryURL + try XCTSkipUnless( + FileManager.default.fileExists(atPath: modelDir.path), + "Cohere MLX model weights missing at \(modelDir.path)" + ) + + let backend = CohereMLXBackend() + let response = try await backend.transcribe(EngineRequest( + audioURL: audioURL, + mode: .singleChannelDiarized(numSpeakers: nil), + languageCode: "en", + keyterms: [], + modelID: CohereMLXBackend.modelID + )) + + let combined = response.utterances.map(\.text).joined(separator: " ") + let wordCount = combined.split(whereSeparator: { $0.isWhitespace }).count + print("MLX integration: \(wordCount) words; first 200 chars: \(combined.prefix(200))") + + XCTAssertGreaterThan( + wordCount, + 1_000, + "Local engine must produce >1000 words on the 583s recording (got \(wordCount))" + ) + XCTAssertNil( + DegenerateOutputDetector.evaluate(combined), + "Local engine output must not look degenerate" + ) + } + + func testBackendThrowsDegenerateOutputErrorOnLoopedAdapterOutput() async { + let looped = String(repeating: "I think that's what I'm hearing ", count: 50) + let adapter = RecordingLocalAdapter(output: .init(text: looped, detectedLanguage: "en")) + let backend = CohereMLXBackend(adapter: adapter, durationReader: FixedDurationReader(duration: 583)) + + do { + _ = try await backend.transcribe(EngineRequest( + audioURL: URL(fileURLWithPath: "/tmp/audio.m4a"), + mode: .singleChannelDiarized(numSpeakers: nil), + languageCode: "en", + keyterms: [], + modelID: CohereMLXBackend.modelID + )) + XCTFail("Backend must throw on degenerate output, not silently return") + } catch let error as CohereMLXBackendError { + switch error { + case .degenerateOutput(let reason, let sample): + XCTAssertTrue(reason.contains("tri-gram") || reason.contains("unique-word"), + "Reason should identify the failure mode, got: \(reason)") + XCTAssertFalse(sample.isEmpty, "Sample should include a snippet of the failing transcript") + XCTAssertLessThanOrEqual(sample.count, 120, "Sample is capped at 120 chars") + } + } catch { + XCTFail("Expected CohereMLXBackendError.degenerateOutput, got \(error)") + } + } } private final class RecordingLocalAdapter: CohereMLXTranscribing, @unchecked Sendable { diff --git a/TranscriberApp/Scribe.xcodeproj/project.pbxproj b/TranscriberApp/Scribe.xcodeproj/project.pbxproj index a75a621..6d664aa 100644 --- a/TranscriberApp/Scribe.xcodeproj/project.pbxproj +++ b/TranscriberApp/Scribe.xcodeproj/project.pbxproj @@ -198,7 +198,7 @@ mainGroup = CF1A6569B97AB964E08BAAEC; minimizedProjectReferenceProxies = 1; packageReferences = ( - 8486363FC92D4A5A3F677AE3 /* XCRemoteSwiftPackageReference "mlx-audio-swift" */, + A79E7498DC22BE4916B2781B /* XCRemoteSwiftPackageReference "mlx-audio-swift" */, 953AE7D0D2147AEBF865B983 /* XCLocalSwiftPackageReference ".." */, ); preferredProjectObjectVersion = 77; @@ -452,12 +452,12 @@ /* End XCLocalSwiftPackageReference section */ /* Begin XCRemoteSwiftPackageReference section */ - 8486363FC92D4A5A3F677AE3 /* XCRemoteSwiftPackageReference "mlx-audio-swift" */ = { + A79E7498DC22BE4916B2781B /* XCRemoteSwiftPackageReference "mlx-audio-swift" */ = { isa = XCRemoteSwiftPackageReference; - repositoryURL = "https://github.com/Blaizzy/mlx-audio-swift.git"; + repositoryURL = "https://github.com/Newarr/mlx-audio-swift.git"; requirement = { kind = revision; - revision = 7734cd1fbbe86460083c1d24199737a24cadfcc8; + revision = b8ec43083e4c5535594dbf9274893f9e6fe4a506; }; }; /* End XCRemoteSwiftPackageReference section */ @@ -469,12 +469,12 @@ }; 2F343A2DF12DC4D575CEA724 /* MLXAudioSTT */ = { isa = XCSwiftPackageProductDependency; - package = 8486363FC92D4A5A3F677AE3 /* XCRemoteSwiftPackageReference "mlx-audio-swift" */; + package = A79E7498DC22BE4916B2781B /* XCRemoteSwiftPackageReference "mlx-audio-swift" */; productName = MLXAudioSTT; }; BF1DD1FF53036826C2308046 /* MLXAudioCore */ = { isa = XCSwiftPackageProductDependency; - package = 8486363FC92D4A5A3F677AE3 /* XCRemoteSwiftPackageReference "mlx-audio-swift" */; + package = A79E7498DC22BE4916B2781B /* XCRemoteSwiftPackageReference "mlx-audio-swift" */; productName = MLXAudioCore; }; /* End XCSwiftPackageProductDependency section */ diff --git a/TranscriberApp/project.yml b/TranscriberApp/project.yml index 57b2c94..50bf4ce 100644 --- a/TranscriberApp/project.yml +++ b/TranscriberApp/project.yml @@ -15,8 +15,8 @@ packages: TranscriberCore: path: .. mlx-audio-swift: - url: https://github.com/Blaizzy/mlx-audio-swift.git - revision: 7734cd1fbbe86460083c1d24199737a24cadfcc8 + url: https://github.com/Newarr/mlx-audio-swift.git + revision: b8ec43083e4c5535594dbf9274893f9e6fe4a506 targets: Scribe: type: application