Use canImport for MLTensor checks (#288)

argmaxinc · Dec 31, 2024 · a7e3858 · a7e3858
1 parent 112a023
commit a7e3858
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/Sources/WhisperKit/Core/Text/TokenSampler.swift b/Sources/WhisperKit/Core/Text/TokenSampler.swift
@@ -28,7 +28,7 @@ open class GreedyTokenSampler: TokenSampling {
         self.decodingOptions = decodingOptions
     }
 
-    #if swift(>=5.10)
+    #if canImport(CoreML.MLState)
     @available(macOS 15, iOS 18, watchOS 11, visionOS 2, *)
     private func sampleWithMLTensor(logits: MLMultiArray) -> (token: Int, logprob: Float) {
         // Use MLTensor operations if available for sampling
@@ -209,7 +209,7 @@ open class GreedyTokenSampler: TokenSampling {
         var completed = false
 
         var result: (token: Int, logprob: Float)
-        #if swift(>=5.10)
+        #if canImport(CoreML.MLState)
         if #available(macOS 15.0, iOS 18.0, watchOS 11.0, visionOS 2.0, *) {
             result = sampleWithMLTensor(logits: logits)
         } else {

diff --git a/Sources/WhisperKit/Core/Utils/Utils.swift b/Sources/WhisperKit/Core/Utils/Utils.swift
@@ -109,7 +109,7 @@ extension MLMultiArray {
     }
 }
 
-#if swift(>=5.10)
+#if canImport(CoreML.MLState)
 @available(macOS 15.0, iOS 18.0, watchOS 11.0, visionOS 2.0, *)
 public extension MLTensor {
     func asIntArray() -> [Int] {