Skip to content
Snippets Groups Projects
Commit 2faafc4c authored by Mingchung Xia's avatar Mingchung Xia
Browse files

Added basic hnsw query test (raw vector result)

parent b530eb8a
No related branches found
No related tags found
1 merge request!13HNSW Implementation with Testcases
Pipeline #111256 failed
......@@ -103,6 +103,39 @@ final class HNSWTests: XCTestCase {
XCTAssert(corpus.count == 17999)
}
// Refer to AllMiniLM_sampleTest.swift for reference
func testBasicQueryExample() 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 _documentEncoder = ContextFreeEncoder<Double>(source: .glove6B50d)
var corpus = HNSWCorpus(encoder: _documentEncoder)
corpus.addUntokenizedDocuments(docs)
let size = MemoryLayout.size(ofValue: corpus)
print("Approximate memory footprint: \(size) bytes")
do {
print("Attempting to query corpus.encodedDocuments.find()...")
let query_embedding: [Double] = _documentEncoder.encodeToken(query).map { Double($0) }
let results = try corpus.encodedDocuments.find(near: query_embedding, limit: 8)
print(results)
print("Query successful!")
} catch {
print("Error when trying corpus.encodedDocuments.find: \(error)")
}
}
}
#endif
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment