Upload New File

Tests/SwiftNLPTests/AllMiniLM_sampleTest.swift

+import XCTest
+@testable import SwiftNLP
+@testable import SwiftAnnoy
+
+final class BERT_test: XCTestCase {
+    
+    
+    func testBasicExample() async throws {
+        
+        let docs = [
+            "The quick brown fox jumps over the lazy dog",
+            "I enjoy taking long walks along the beach at sunset",
+            "Advances in neural networks have enabled new AI capabilities",
+            "The stock market experienced a significant downturn last week",
+            "Cooking a good meal can be both an art and a science",
+            "The exploration of space is both challenging and rewarding",
+            "Machine learning models are becoming increasingly sophisticated",
+            "I love reading about history and ancient civilizations"
+        ]
+        
+        let query = [
+            "I like to read about new technology and artificial intelligence"
+        ]
+        //        let docs = ["cat dog", "bee fly"]
+        
+        var database_embedding: [[Float]] = []
+        var query_embedding: [Float] = []
+        var embedding_dim: Int = 384
+        
+        var model = MiniLMEmbeddings()
+        
+        query_embedding = await model.encode(sentence: query[0])!
+        
+        var i = 1
+        //append sentence embedding to database_embedding
+        for string in docs {
+            if let vector = await model.encode(sentence: string) {
+                database_embedding.append(vector)
+                print(i)
+                i += 1
+            } else {
+                fatalError("Error occurred1")
+            }
+            
+        }
+        
+        let index = AnnoyIndex<Float>(itemLength: embedding_dim, metric: .euclidean)
+        
+        try? index.addItems(items: &database_embedding)
+        try? index.build(numTrees: 50)
+        
+        let results = index.getNNsForVector(vector: &query_embedding, neighbors: 8)
+        
+        if let finalresult = results {
+            let extractedIndeices = finalresult.indices
+            for index in extractedIndeices {
+                if index < docs.count {
+                    print(docs[index])
+                } else {
+                    print("Index \(index) out of range.")
+                }
+            }
+        }
+        print(results)
+    }
+}