Additional options native embeddings (#23)

* Update NativeEmbeddings.swift * update update * Update Sources/SimilaritySearchKit/Core/Embeddings/Models/NativeEmbeddings.swift Co-authored-by: Zach Nagengast <[email protected]> * Update Sources/SimilaritySearchKit/Core/Embeddings/Models/NativeEmbeddings.swift Co-authored-by: Zach Nagengast <[email protected]> --------- Co-authored-by: Zach Nagengast <[email protected]>
ZachNagengast · Nov 7, 2023 · 61a2e09 · 61a2e09
1 parent 18b52c6
commit 61a2e09
Showing 1 changed file with 20 additions and 6 deletions.
diff --git a/Sources/SimilaritySearchKit/Core/Embeddings/Models/NativeEmbeddings.swift b/Sources/SimilaritySearchKit/Core/Embeddings/Models/NativeEmbeddings.swift
@@ -8,26 +8,40 @@
 import Foundation
 import NaturalLanguage
 
+public enum NativeEmbeddingType {
+    case wordEmbedding
+    case sentenceEmbedding
+}
+
 @available(macOS 11.0, iOS 15.0, *)
 public class NativeEmbeddings: EmbeddingsProtocol {
     public let model: ModelActor
     public let tokenizer: any TokenizerProtocol
 
-    public init(language: NLLanguage = .english) {
+    public init(language: NLLanguage = .english, type: NativeEmbeddingType = .sentenceEmbedding) {
         self.tokenizer = NativeTokenizer()
-        self.model = ModelActor(language: language)
+        self.model = ModelActor(language: language, type: type)
     }
 
     // MARK: - Dense Embeddings
 
     public actor ModelActor {
         private let model: NLEmbedding
 
-        init(language: NLLanguage) {
-            guard let nativeModel = NLEmbedding.sentenceEmbedding(for: language) else {
-                fatalError("Failed to load the Core ML model.")
+        init(language: NLLanguage, type:NativeEmbeddingType = .sentenceEmbedding) {
+            switch type {
+                case .sentenceEmbedding:
+                    guard let nativeModel = NLEmbedding.sentenceEmbedding(for: language) else {
+                        fatalError("Failed to load the Core ML model.")
+                    }
+                    model = nativeModel
+                case .wordEmbedding:
+                    guard let nativeModel = NLEmbedding.wordEmbedding(for: language) else {
+                        fatalError("Failed to load the Core ML model.")
+                    }
+                    model = nativeModel
             }
-            model = nativeModel
+
         }
 
         func vector(for sentence: String) -> [Float]? {