jkrukowski/swift-embeddings

Run embedding models locally in `Swift` using `MLTensor`.

Supported Models Architectures

BERT (Bidirectional Encoder Representations from Transformers)

Some of the supported models on Hugging Face:

NOTE: google-bert/bert-base-uncased is supported but weightKeyTransform must be provided in the LoadConfig:

let modelBundle = try await Bert.loadModelBundle(
    from: "google-bert/bert-base-uncased",
    loadConfig: .googleBert
)

ModernBERT

Some of the supported models on Hugging Face:

NOTE: answerdotai/ModernBERT-base is supported but weights must be prefixed with model.

let modelBundle = try await Bert.loadModelBundle(
    from: "answerdotai/ModernBERT-base",
    loadConfig: .addWeightKeyPrefix("model.")
)

NomicBERT (Nomic Embed)

Some of the supported models on Hugging Face:

nomic-ai/nomic-embed-text-v1.5

RoBERTa (Robustly Optimized BERT Approach)

Some of the supported models on Hugging Face:

FacebookAI/roberta-base

NOTE: Weights in FacebookAI/roberta-base must be prefixed with roberta., this has to be provided in the LoadConfig:

let modelBundle = try await Roberta.loadModelBundle(
    from: "FacebookAI/roberta-base",
    loadConfig: .addWeightKeyPrefix("roberta.")
)

XLM-RoBERTa (Cross-lingual Language Model - Robustly Optimized BERT Approach)

Some of the supported models on Hugging Face:

NOTE: Weights in FacebookAI/xlm-roberta-base must be prefixed with roberta., this has to be provided in the LoadConfig:

let modelBundle = try await XLMRoberta.loadModelBundle(
    from: "FacebookAI/xlm-roberta-base",
    loadConfig: .addWeightKeyPrefix("roberta.")
)

CLIP (Contrastive Language–Image Pre-training)

NOTE: only text encoding is supported for now. Some of the supported models on Hugging Face:

Word2Vec

NOTE: it's a word embedding model. It loads and keeps the whole model in memory. For the more memory efficient solution, you might want to use SQLiteVec. Some of the supported models on Hugging Face:

Model2Vec

More info here.

Some of the supported models on Hugging Face:

Static Embeddings

More info here.

Some of the supported models on Hugging Face:

Installation

Add the following to your Package.swift file. In the package dependencies add:

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

In the target dependencies add:

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

Usage

Encoding

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 = modelBundle.encode("The cat is black")
let result = await encoded.cast(to: Float.self).shapedArray(of: Float.self).scalars

// print result
print(result)

Batch Encoding

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 = modelBundle.batchEncode(texts)
let distance = cosineDistance(encoded, encoded)
let result = await distance.cast(to: Float.self).shapedArray(of: Float.self).scalars
print(result)

Command Line Demo

To run the command line demo, use the following command:

swift run embeddings-cli <subcommand> [--model-id <model-id>] [--model-file <model-file>] [--text <text>] [--max-length <max-length>]

Subcommands:

bert                    Encode text using BERT model
modern-bert             Encode text using ModernBERT model
nomic-bert              Encode text using Nomic embedding model
clip                    Encode text using CLIP model
model2vec               Encode text using Model2Vec model
roberta                 Encode text using RoBERTa model
static-embeddings       Encode text using Static Embeddings model
xlm-roberta             Encode text using XLMRoberta model
word2vec                Encode word using Word2Vec model

Command line options:

--model-id <model-id>                       Id of the model to use
--model-file <model-file>                   Path to the model file (only for `Word2Vec`)
--text <text>                               Text to encode
--max-length <max-length>                   Maximum length of the input (not for `Word2Vec`)
-h, --help                                  Show help information.

Code Formatting

This project uses swift-format. To format the code run:

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

Acknowledgements

This project is based on and uses some of the code from:

mlx-embeddings

Package Metadata

Repository: jkrukowski/swift-embeddings

Default branch: main

README: README.md