Fixed Conformance of 'Float16' warning (#58)

* fixed warnings * removed fixme
argmaxinc · Mar 7, 2024 · 20c9001 · 20c9001
1 parent 880f95e
commit 20c9001
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/Sources/WhisperKit/Core/Models.swift b/Sources/WhisperKit/Core/Models.swift
@@ -1,18 +1,22 @@
 //  For licensing see accompanying LICENSE.md file.
 //  Copyright © 2024 Argmax, Inc. All rights reserved.
 
+import Accelerate
 import CoreML
 import Hub
 import NaturalLanguage
 import Tokenizers
 
-#if os(watchOS) || arch(arm64)
-@available(macOS 13, iOS 16, watchOS 10, visionOS 1, *)
+#if !((os(macOS) || targetEnvironment(macCatalyst)) && arch(x86_64))
 public typealias FloatType = Float16
 #else
 public typealias FloatType = Float
 #endif
 
+#if (os(macOS) || targetEnvironment(macCatalyst)) && arch(arm64)
+extension Float16: BNNSScalar {}
+#endif
+
 // MARK: - CoreML
 
 public protocol WhisperMLModel {

diff --git a/Sources/WhisperKit/Core/TokenSampler.swift b/Sources/WhisperKit/Core/TokenSampler.swift
@@ -44,7 +44,7 @@ public class GreedyTokenSampler: TokenSampling {
 
             let logitsDescriptor = BNNSNDArrayDescriptor(
                 data: logitsRawPointer,
-                scalarType: FloatType.self, // FIXME: Float16 here breaks in swift 6
+                scalarType: FloatType.self,
                 shape: .vector(logits.count, stride: 1)
             )!