File size: 1,122 Bytes
902bb3d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import CoreML
import SQLiteVec
import Testing
@testable import EmbeddingsBenchmarkLib
func createDatabase(_ data: [[Float]]) async throws -> Database {
try SQLiteVec.initialize()
let db = try Database(.inMemory)
try await db.execute("CREATE VIRTUAL TABLE embeddings USING vec0(embedding float[3])")
for (index, row) in data.enumerated() {
try await db.execute(
"""
INSERT INTO embeddings(rowid, embedding)
VALUES (?, ?)
""",
params: [index, row]
)
}
return db
}
@Test func testEmbeddingMethods() async throws {
let data: [[Float]] = [
[1.0, 2.0, 3.0],
[4.0, 5.0, 6.0],
[7.0, 8.0, 9.0]
]
let embeddings = MLTensor(shape: [3, 3], scalars: data.flatMap { $0 })
let coreMLResult = await queryEmbeddings(embeddings: embeddings, tokenIds: [0, 2])
let db = try await createDatabase(data)
let sqliteResult = try await queryEmbeddings(
db: db,
query: "(?, ?)",
tokenIds: [0, 2],
vectorSize: 3)
#expect(coreMLResult == sqliteResult)
}
|