Initial commit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+embeddings.db filter=lfs diff=lfs merge=lfs -text

EmbeddingsBenchmark/.gitignore

@@ -0,0 +1,8 @@
+.DS_Store
+/.build
+/Packages
+xcuserdata/
+DerivedData/
+.swiftpm/configuration/registries.json
+.swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata
+.netrc

EmbeddingsBenchmark/Benchmarks/EmbeddingsBenchmark/EmbeddingsBenchmark.swift

@@ -0,0 +1,34 @@
+import CoreML
+import Benchmark
+import Embeddings
+import EmbeddingsBenchmarkLib
+import Foundation
+import SQLiteVec
+
+@MainActor let benchmarks = {
+    let parameterization = [10, 1000, 100000]
+    parameterization.forEach { count in
+        Benchmark("SQLite benchmark", configuration: .init(tags: ["count" : count.description])) { benchmark, db in
+            let tokenIds = (0..<count).map { _ in Int.random(in: 0..<1000) }
+            let query = "(\(tokenIds.map { _ in "?" }.joined(separator: ",")))"
+            benchmark.startMeasurement()
+            let result = try await queryEmbeddings(db: db, query: query, tokenIds: tokenIds)
+            blackHole(result)
+        } setup: {
+            try initializeDatabase("../embeddings.db")
+        }
+
+        Benchmark("CoreML benchmark", configuration: .init(tags: ["count" : count.description])) { (benchmark: _, embeddings: MLTensor) in
+            let tokenIds = (0..<count).map { _ in Int32.random(in: 0..<1000) }
+            benchmark.startMeasurement()
+            let result = await queryEmbeddings(embeddings: embeddings, tokenIds: tokenIds)
+            blackHole(result)
+        } setup: {
+            let modelBundle = try await StaticEmbeddings.loadModelBundle(
+                from: "sentence-transformers/static-retrieval-mrl-en-v1",
+                loadConfig: LoadConfig.staticEmbeddings
+            )
+            return modelBundle.model.embeddings
+        }
+    }
+}

EmbeddingsBenchmark/Package.resolved

@@ -0,0 +1,141 @@
+{
+  "originHash" : "c0c5da26b668f59a78fa7ea6e6755ca0f61bb28d38c1907156dbe4a6e49b39ec",
+  "pins" : [
+    {
+      "identity" : "hdrhistogram-swift",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/HdrHistogram/hdrhistogram-swift",
+      "state" : {
+        "revision" : "a69fa24d7b70421870cafa86340ece900489e17e",
+        "version" : "0.1.2"
+      }
+    },
+    {
+      "identity" : "jinja",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/johnmai-dev/Jinja",
+      "state" : {
+        "revision" : "bbddb92fc51ae420b87300298370fd1dfc308f73",
+        "version" : "1.1.1"
+      }
+    },
+    {
+      "identity" : "package-benchmark",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/ordo-one/package-benchmark",
+      "state" : {
+        "revision" : "51fffb6bc58bcbaca9c1954e6e52b56f3c9c2d54",
+        "version" : "1.27.4"
+      }
+    },
+    {
+      "identity" : "package-jemalloc",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/ordo-one/package-jemalloc",
+      "state" : {
+        "revision" : "e8a5db026963f5bfeac842d9d3f2cc8cde323b49",
+        "version" : "1.0.0"
+      }
+    },
+    {
+      "identity" : "sqlitevec",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/jkrukowski/SQLiteVec.git",
+      "state" : {
+        "revision" : "eaa734a8c7615e2ef791c218280292f27d33db39",
+        "version" : "0.0.13"
+      }
+    },
+    {
+      "identity" : "swift-argument-parser",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-argument-parser.git",
+      "state" : {
+        "revision" : "41982a3656a71c768319979febd796c6fd111d5c",
+        "version" : "1.5.0"
+      }
+    },
+    {
+      "identity" : "swift-atomics",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-atomics",
+      "state" : {
+        "revision" : "cd142fd2f64be2100422d658e7411e39489da985",
+        "version" : "1.2.0"
+      }
+    },
+    {
+      "identity" : "swift-collections",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-collections.git",
+      "state" : {
+        "revision" : "671108c96644956dddcd89dd59c203dcdb36cec7",
+        "version" : "1.1.4"
+      }
+    },
+    {
+      "identity" : "swift-embeddings",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/jkrukowski/swift-embeddings",
+      "state" : {
+        "revision" : "b9a8e6b536cb7a4541ab7ba411ea5c4fed6e9f18",
+        "version" : "0.0.11"
+      }
+    },
+    {
+      "identity" : "swift-numerics",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-numerics",
+      "state" : {
+        "revision" : "0a5bc04095a675662cf24757cc0640aa2204253b",
+        "version" : "1.0.2"
+      }
+    },
+    {
+      "identity" : "swift-safetensors",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/jkrukowski/swift-safetensors.git",
+      "state" : {
+        "revision" : "718b0f38f912e0bf9d92130fa1e1fe2ae5136dd6",
+        "version" : "0.0.7"
+      }
+    },
+    {
+      "identity" : "swift-sentencepiece",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/jkrukowski/swift-sentencepiece",
+      "state" : {
+        "revision" : "75d725019ff0b75fbbd7128314fe6710c5a86df0",
+        "version" : "0.0.5"
+      }
+    },
+    {
+      "identity" : "swift-system",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-system",
+      "state" : {
+        "revision" : "c8a44d836fe7913603e246acab7c528c2e780168",
+        "version" : "1.4.0"
+      }
+    },
+    {
+      "identity" : "swift-transformers",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/huggingface/swift-transformers.git",
+      "state" : {
+        "revision" : "8a83416cc00ab07a5de9991e6ad817a9b8588d20",
+        "version" : "0.1.15"
+      }
+    },
+    {
+      "identity" : "texttable",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/ordo-one/TextTable",
+      "state" : {
+        "revision" : "a27a07300cf4ae322e0079ca0a475c5583dd575f",
+        "version" : "0.0.2"
+      }
+    }
+  ],
+  "version" : 3
+}

EmbeddingsBenchmark/Package.swift

@@ -0,0 +1,54 @@
+// swift-tools-version: 6.0
+
+import PackageDescription
+
+let package = Package(
+    name: "EmbeddingsBenchmark",
+    platforms: [
+        .macOS(.v15)
+    ],
+    dependencies: [
+        .package(
+            url: "https://github.com/jkrukowski/SQLiteVec.git",
+            from: "0.0.13"
+        ),
+        .package(
+            url: "https://github.com/jkrukowski/swift-embeddings",
+            from: "0.0.11"
+        ),
+        .package(
+            url: "https://github.com/ordo-one/package-benchmark",
+            from: "1.27.4"
+        )
+    ],
+    targets: [
+        .target(
+            name: "EmbeddingsBenchmarkLib",
+            dependencies: [
+                .product(name: "SQLiteVec", package: "SQLiteVec"),
+                .product(name: "Embeddings", package: "swift-embeddings")
+            ]
+        ),
+        .testTarget(
+            name: "EmbeddingsBenchmarkTests",
+            dependencies: [
+                "EmbeddingsBenchmarkLib",
+                .product(name: "SQLiteVec", package: "SQLiteVec"),
+                .product(name: "Embeddings", package: "swift-embeddings")
+            ]
+        ),
+        .executableTarget(
+            name: "EmbeddingsBenchmark",
+            dependencies: [
+                "EmbeddingsBenchmarkLib",
+                .product(name: "SQLiteVec", package: "SQLiteVec"),
+                .product(name: "Embeddings", package: "swift-embeddings"),
+                .product(name: "Benchmark", package: "package-benchmark")
+            ],
+            path: "Benchmarks/EmbeddingsBenchmark",
+            plugins: [
+                .plugin(name: "BenchmarkPlugin", package: "package-benchmark")
+            ]
+        ),
+    ]
+)

EmbeddingsBenchmark/Sources/EmbeddingsBenchmarkLib/EmbeddingsBenchmarkLib.swift

@@ -0,0 +1,33 @@
+import Accelerate
+import CoreML
+import Embeddings
+import Foundation
+import SQLiteVec
+
+public func initializeDatabase(_ filePath: String) throws -> Database {
+    try SQLiteVec.initialize()
+    return try Database(.uri(filePath))
+}
+
+@discardableResult
+public func queryEmbeddings(db: Database, query: String, tokenIds: [Int], vectorSize: Int = 1024) async throws -> [Float] {
+    let result = try await db.query("SELECT embedding FROM embeddings WHERE rowid IN \(query)", params: tokenIds)
+    var acc = [Float](repeating: 0, count: vectorSize)
+    for item in result {
+        guard let embedding = item["embedding"] as? Data else {
+            continue
+        }
+        let row: [Float] = embedding.toArray()
+        acc = vDSP.add(row, acc)
+    }
+    return vDSP.divide(acc, Float(result.count))
+}
+
+@discardableResult
+public func queryEmbeddings(embeddings: MLTensor, tokenIds: [Int32]) async -> [Float] {
+    let indices = MLTensor(shape: [tokenIds.count], scalars: tokenIds)
+    let data = embeddings
+        .gathering(atIndices: indices, alongAxis: 0)
+        .mean(alongAxes: 0)
+    return await data.shapedArray(of: Float.self).scalars
+}

EmbeddingsBenchmark/Tests/EmbeddingsBenchmarkTests/EmbeddingsBenchmarkTests.swift

@@ -0,0 +1,40 @@
+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)
+}

README.md

@@ -1,3 +1,123 @@
 ---
 license: mit
 ---
+
+# Static Embeddings using SQLiteVec
+
+- This repo contains [sentence-transformers/static-retrieval-mrl-en-v1](https://huggingface.co/sentence-transformers/static-retrieval-mrl-en-v1)
+model converted to [sqlite-vec](https://github.com/asg017/sqlite-vec) database.
+
+- This repo contains a Swift package that demonstrates how to use the [SQLiteVec](https://github.com/jkrukowski/SQLiteVec) to compute static embeddings.
+
+- This repo contains a benchmark that compares the performance of computing embeddings using [swift-embeddings](https://github.com/jkrukowski/swift-embeddings) and [SQLiteVec](https://github.com/jkrukowski/SQLiteVec).
+
+The sqlite-vec database was created using the following command:
+
+```sql
+CREATE VIRTUAL TABLE embeddings USING vec0(embedding float[1024])
+```
+
+and then the embeddings were inserted using the following command:
+
+```sql
+INSERT INTO embeddings(rowid, embedding) VALUES (?, ?)
+```
+
+## Benchmark
+
+To run the benchmark, execute the following command:
+
+```bash
+cd EmbeddingsBenchmark
+swift package --disable-sandbox benchmark
+```
+
+### Results
+
+```
+
+CoreML benchmark (count: 10)
+╒═══════════════════════════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╕
+│ Metric                        │      p0 │     p25 │     p50 │     p75 │     p90 │     p99 │    p100 │ Samples │
+╞═══════════════════════════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╡
+│ Instructions (K) *            │    1712 │    1779 │    1779 │    1791 │    1792 │    1805 │    1899 │    5136 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Malloc (total) *              │    1131 │    1132 │    1132 │    1132 │    1132 │    1132 │    1139 │    5136 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Memory (resident peak) (M)    │     156 │     158 │     158 │     159 │     159 │     159 │     159 │    5136 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Throughput (# / s) (#)        │   11121 │    9471 │    9287 │    8919 │    8607 │    7747 │    5810 │    5136 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Time (total CPU) (μs) *       │      92 │     107 │     109 │     114 │     118 │     131 │     174 │    5136 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Time (wall clock) (μs) *      │      90 │     106 │     108 │     112 │     116 │     129 │     172 │    5136 │
+╘═══════════════════════════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╛
+
+CoreML benchmark (count: 1000)
+╒═══════════════════════════════╤═════════╤═════════╤══��══════╤═════════╤═════════╤═════════╤═════════╤═════════╕
+│ Metric                        │      p0 │     p25 │     p50 │     p75 │     p90 │     p99 │    p100 │ Samples │
+╞═══════════════════════════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╡
+│ Instructions (K) *            │    3320 │    3422 │    3467 │    3516 │    3545 │    3566 │    3640 │    2698 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Malloc (total) *              │    1131 │    1132 │    1132 │    1132 │    1132 │    1132 │    1140 │    2698 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Memory (resident peak) (M)    │     159 │     162 │     162 │     162 │     162 │     162 │     162 │    2698 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Throughput (# / s) (#)        │    4418 │    4299 │    4231 │    4139 │    3999 │    3515 │    2093 │    2698 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Time (total CPU) (μs) *       │     228 │     234 │     238 │     243 │     252 │     287 │     480 │    2698 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Time (wall clock) (μs) *      │     226 │     233 │     237 │     242 │     250 │     285 │     478 │    2698 │
+╘═══════════════════════════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╛
+
+CoreML benchmark (count: 100000)
+╒═══════════════════════════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╕
+│ Metric                        │      p0 │     p25 │     p50 │     p75 │     p90 │     p99 │    p100 │ Samples │
+╞═══════════════════════════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╡
+│ Instructions (M) *            │     163 │     163 │     163 │     163 │     163 │     168 │     168 │      20 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Malloc (total) *              │    1131 │    1131 │    1131 │    1132 │    1133 │    1139 │    1139 │      20 │
+├───────────────────────────────┼────────��┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Memory (resident peak) (M)    │     566 │     567 │     567 │     568 │     568 │     568 │     568 │      20 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Throughput (# / s) (#)        │      21 │      21 │      21 │      21 │      21 │      20 │      20 │      20 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Time (total CPU) (ms) *       │      47 │      48 │      48 │      48 │      48 │      50 │      50 │      20 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Time (wall clock) (ms) *      │      47 │      48 │      48 │      48 │      48 │      50 │      50 │      20 │
+╘═══════════════════════════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╛
+
+SQLite benchmark (count: 10)
+╒═══════════════════════════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╕
+│ Metric                        │      p0 │     p25 │     p50 │     p75 │     p90 │     p99 │    p100 │ Samples │
+╞═══════════════════════════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╡
+│ Instructions (M) *            │     103 │     117 │     122 │     128 │     132 │     136 │     140 │     144 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Malloc (total) *              │      99 │     106 │     106 │     106 │     106 │     107 │     107 │     144 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Memory (resident peak) (M)    │      15 │      17 │      19 │      19 │      19 │      19 │      19 │     144 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Throughput (# / s) (#)        │     186 │     155 │     147 │     139 │     132 │     126 │     126 │     144 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Time (total CPU) (μs) *       │    5382 │    6451 │    6779 │    7180 │    7582 │    7922 │    7932 │     144 │
+├────────────────────────────���──┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Time (wall clock) (μs) *      │    5380 │    6443 │    6779 │    7180 │    7574 │    7918 │    7930 │     144 │
+╘═══════════════════════════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╛
+
+SQLite benchmark (count: 1000)
+╒═══════════════════════════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╤═════════╕
+│ Metric                        │      p0 │     p25 │     p50 │     p75 │     p90 │     p99 │    p100 │ Samples │
+╞═══════════════════════════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╪═════════╡
+│ Instructions (M) *            │    1517 │    1572 │    1599 │    1611 │    1618 │    1618 │    1618 │       9 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Malloc (total) *              │    6330 │    6339 │    6367 │    6379 │    6414 │    6414 │    6414 │       9 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Memory (resident peak) (M)    │      22 │      23 │      23 │      23 │      23 │      23 │      23 │       9 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Throughput (# / s) (#)        │       9 │       8 │       8 │       8 │       8 │       8 │       8 │       9 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Time (total CPU) (ms) *       │     114 │     118 │     120 │     121 │     121 │     121 │     121 │       9 │
+├───────────────────────────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤
+│ Time (wall clock) (ms) *      │     114 │     118 │     120 │     121 │     121 │     121 │     121 │       9 │
+╘═══════════════════════════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╧═════════╛
+```

embeddings.db

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:74e3a963cf8e8242835a38caf2e4aa78bb4b5ae2ab07b0e1e95e07c01d9147ca
+size 126672896