swift-embeddings

The benchmarking toolchain allows you to manage performance baselines to track regressions. Use the following commands to list, update, or compare benchmarks against a stored baseline (e.g., the main branch).

# List available benchmarks
swift package benchmark list

# Run and store a baseline
swift package --disable-sandbox benchmark baseline update main

# Compare against a stored baseline
swift package --disable-sandbox benchmark baseline compare main

Accuracy tests compare Swift outputs against a Python/transformers reference. These tests are skipped by default unless the UV_PATH environment variable points to a uv binary.

Important: You must run these tests with the --no-parallel flag. Running multiple model forward passes concurrently can cause crashes in the underlying MLTensor/BNNS matmul operations.

To run the accuracy tests, use the following command structure:

• Set PYTORCH_ENABLE_MPS_FALLBACK=1 to allow the PyTorch reference to fall back to CPU for operations unsupported on MPS.
• Set UV_PATH to your uv binary location.
• Set CI_DISABLE_NETWORK_MONITOR=1 to disable the network monitor (assuming models are already downloaded).

PYTORCH_ENABLE_MPS_FALLBACK=1 UV_PATH=$(which uv) CI_DISABLE_NETWORK_MONITOR=1 \
    swift test --no-parallel --filter AccuracyTests

To use swift-embeddings in your Swift project, add the package to your Package.swift file.

1. Add the package dependency:

dependencies: [
    .package(url: "https://github.com/jkrukowski/swift-embeddings", from: "0.0.16")
]

2. Add the Embeddings product to your target dependencies:

dependencies: [
    .product(name: "Embeddings", package: "swift-embeddings")
]

dependencies: [
    .package(url: "https://github.com/jkrukowski/swift-embeddings", from: "0.0.16")
]

targets: [
    .target(
        name: "YourTarget",
        dependencies: [
            .product(name: "Embeddings", package: "swift-embeddings")
        ]
    )
]

To execute the benchmarks for swift-embeddings, run the following command from the Benchmarks directory. Note that the benchmarking toolchain is maintained as a separate Swift package to keep the main package graph clean.

On the first run, the tool will download a small static embedding model (minishlab/potion-base-8M) from the Hugging Face Hub and cache it locally. Subsequent runs will use the cached model and will not require network access during measurement.

swift package --disable-sandbox benchmark

The Benchmarks package uses package-benchmark to measure the encoding performance of various model architectures.

To run all benchmarks, use the --disable-sandbox flag. This is required because the benchmark setup phase downloads the model (which is cached after the first run).

To see a list of all available benchmarks, use the benchmark list command.

# Run all benchmarks
swift package --package-path Benchmarks --disable-sandbox benchmark

# List available benchmarks
swift package --package-path Benchmarks benchmark list

This project uses swift-format for code styling. To apply formatting to the entire project using the provided configuration, run the following command:

swift format . -i -r --configuration .swift-format

You can load a model bundle from Hugging Face and encode a single string into an embedding. The resulting embedding can be converted to a Float array using MLTensor utilities.

import Embeddings

// load model and tokenizer from Hugging Face
let modelBundle = try await Bert.loadModelBundle(
    from: "sentence-transformers/all-MiniLM-L6-v2"
)

// encode text
let encoded = try modelBundle.encode("The cat is black")
let result = await encoded.cast(to: Float.self).shapedArray(of: Float.self).scalars

// print result
print(result)

For processing multiple strings at once, use batchEncode. This is useful for calculating distances (like cosine distance) between multiple embeddings.

import Embeddings
import MLTensorUtils

let texts = [
    "The cat is black",
    "The dog is black",
    "The cat sleeps well"
]
let modelBundle = try await Bert.loadModelBundle(
    from: "sentence-transformers/all-MiniLM-L6-v2"
)
let encoded = try modelBundle.batchEncode(texts)
let distance = cosineDistance(encoded, encoded)
let result = await distance.cast(to: Float.self).shapedArray(of: Float.self).scalars
print(result)

The word2vec command accepts the following options: