Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 0 additions & 24 deletions Sources/Fluid/Services/ExternalCoreMLModelRegistry.swift
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,6 @@ enum ExternalCoreMLArtifactsValidationError: LocalizedError {
case manifestUnreadable(URL, Error)
case unexpectedModelID(expected: String, actual: String)
case unexpectedSampleRate(expected: Int, actual: Int)
case unexpectedMaxAudioSamples(expected: Int, actual: Int)
case unexpectedMaxAudioSeconds(expected: Double, actual: Double)

var errorDescription: String? {
switch self {
Expand All @@ -43,10 +41,6 @@ enum ExternalCoreMLArtifactsValidationError: LocalizedError {
return "Unexpected model_id '\(actual)'. Expected '\(expected)'."
case let .unexpectedSampleRate(expected, actual):
return "Unexpected sample rate \(actual). Expected \(expected)."
case let .unexpectedMaxAudioSamples(expected, actual):
return "Unexpected max audio samples \(actual). Expected \(expected)."
case let .unexpectedMaxAudioSeconds(expected, actual):
return "Unexpected max audio seconds \(actual). Expected \(expected)."
}
}
}
Expand All @@ -62,8 +56,6 @@ struct ExternalCoreMLASRModelSpec {
let cachedDecoderFileName: String
let expectedModelID: String
let expectedSampleRate: Int
let expectedMaxAudioSamples: Int
let expectedMaxAudioSeconds: Double
let computeConfiguration: CohereTranscribeComputeConfiguration
let sourceURL: URL?
let repositoryOwner: String?
Expand Down Expand Up @@ -137,20 +129,6 @@ struct ExternalCoreMLASRModelSpec {
actual: manifest.sampleRate
)
Comment on lines 129 to 130
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Enforce sane manifest audio-window limits

validateArtifactsOrThrow now accepts any max_audio_samples/max_audio_seconds values as long as model_id and sample_rate match, so a bad or tampered manifest can pass validation and then drive unbounded allocations in paddedSamplesToModelLimit (it pads to loadedManifest.maxAudioSamples). In environments using downloaded or user-provided artifact directories, this can turn a manifest change into OOM/crash behavior during transcription; adding range checks (or expected bounds) for these manifest fields would prevent that regression.

Useful? React with 👍 / 👎.

}

guard manifest.maxAudioSamples == self.expectedMaxAudioSamples else {
throw ExternalCoreMLArtifactsValidationError.unexpectedMaxAudioSamples(
expected: self.expectedMaxAudioSamples,
actual: manifest.maxAudioSamples
)
}

guard manifest.maxAudioSeconds == self.expectedMaxAudioSeconds else {
throw ExternalCoreMLArtifactsValidationError.unexpectedMaxAudioSeconds(
expected: self.expectedMaxAudioSeconds,
actual: manifest.maxAudioSeconds
)
}
}
}

Expand All @@ -169,8 +147,6 @@ enum ExternalCoreMLModelRegistry {
cachedDecoderFileName: "cohere_decoder_cached.mlpackage",
expectedModelID: "CohereLabs/cohere-transcribe-03-2026",
expectedSampleRate: 16_000,
expectedMaxAudioSamples: 560_000,
expectedMaxAudioSeconds: 35.0,
computeConfiguration: .aneSmall,
sourceURL: URL(string: "https://huggingface.co/BarathwajAnandan/cohere-transcribe-03-2026-CoreML-6bit"),
repositoryOwner: "BarathwajAnandan",
Expand Down
46 changes: 45 additions & 1 deletion Sources/Fluid/Services/ExternalCoreMLTranscriptionProvider.swift
Original file line number Diff line number Diff line change
Expand Up @@ -58,18 +58,20 @@ final class ExternalCoreMLTranscriptionProvider: TranscriptionProvider {
case .cohereTranscribe:
let manager = CohereTranscribeAsrManager()
progressHandler?(0.9)
try self.invalidateCompiledCohereCacheIfNeeded(at: directory)
let computeSummary = [
String(describing: spec.computeConfiguration.frontend),
String(describing: spec.computeConfiguration.encoder),
String(describing: spec.computeConfiguration.crossKV),
String(describing: spec.computeConfiguration.decoder),
].joined(separator: "/")
DebugLogger.shared.info(
"ExternalCoreML: loading Cohere models [splitCompute=\(computeSummary), maxAudioSamples=\(self.loadedManifest?.maxAudioSamples ?? spec.expectedMaxAudioSamples)]",
"ExternalCoreML: loading Cohere models [splitCompute=\(computeSummary), maxAudioSamples=\(self.loadedManifest?.maxAudioSamples ?? 0)]",
source: "ExternalCoreML"
)
try await manager.loadModels(from: directory, computeConfiguration: spec.computeConfiguration)
self.cohereManager = manager
self.persistCompiledCohereCacheStamp(at: directory)
}

self.isReady = true
Expand Down Expand Up @@ -274,6 +276,48 @@ final class ExternalCoreMLTranscriptionProvider: TranscriptionProvider {
)
}

private func invalidateCompiledCohereCacheIfNeeded(at directory: URL) throws {
guard let manifest = self.loadedManifest else { return }

let compiledDirectory = CohereTranscribeAsrModels.compiledArtifactsDirectory(for: directory)
guard FileManager.default.fileExists(atPath: compiledDirectory.path) else { return }

let currentStamp = Self.compiledCohereCacheStamp(for: manifest)
let stampURL = Self.compiledCohereCacheStampURL(for: directory)
let previousStamp = try? String(contentsOf: stampURL, encoding: .utf8)
.trimmingCharacters(in: .whitespacesAndNewlines)

guard previousStamp != currentStamp else { return }

let reason = previousStamp == nil ? "missing cache stamp" : "manifest changed"
DebugLogger.shared.warning(
"ExternalCoreML: clearing stale compiled Cohere cache [reason=\(reason)]",
source: "ExternalCoreML"
)
try FileManager.default.removeItem(at: compiledDirectory)
}

private func persistCompiledCohereCacheStamp(at directory: URL) {
guard let manifest = self.loadedManifest else { return }
let stampURL = Self.compiledCohereCacheStampURL(for: directory)
let stamp = Self.compiledCohereCacheStamp(for: manifest)
try? stamp.write(to: stampURL, atomically: true, encoding: .utf8)
}

private static func compiledCohereCacheStampURL(for directory: URL) -> URL {
directory.appendingPathComponent(".cohere_compiled_cache_stamp", isDirectory: false)
}

private static func compiledCohereCacheStamp(for manifest: ExternalCoreMLManifestIdentity) -> String {
[
manifest.modelID,
String(manifest.sampleRate),
String(manifest.maxAudioSamples),
String(manifest.maxAudioSeconds),
String(manifest.overlapSamples ?? 0),
].joined(separator: "|")
}

private func previewSamples(for samples: [Float]) -> [Float] {
let sampleRate = self.loadedManifest?.sampleRate
?? (self.modelOverride ?? SettingsStore.shared.selectedSpeechModel).externalCoreMLSpec?.expectedSampleRate
Expand Down
Loading