Home

Awesome

SimilaritySearchKit

License

ssk-logo

SimilaritySearchKit is a Swift package enabling on-device text embeddings and semantic search functionality for iOS and macOS applications in just a few lines. Emphasizing speed, extensibility, and privacy, it supports a variety of built-in state-of-the-art NLP models and similarity metrics, in addition to seamless integration for bring-your-own options.

Chat With Files Example

Use Cases

Some potential use cases for SimilaritySearchKit include:

By leveraging SimilaritySearchKit, developers can easily create powerful applications that keep data close to home without major tradeoffs in functionality or performance.

Installation

To install SimilaritySearchKit, simply add it as a dependency to your Swift project using the Swift Package Manager. I recommend using the Xcode method personally via:

FileAdd Packages...Search or Enter Package Urlhttps://github.com/ZachNagengast/similarity-search-kit.git

Xcode should give you the following options to choose which model you'd like to add (see available models below for help choosing):

Xcode Swift Package Manager Import

If you want to add it via Package.swift, add the following line to your dependencies array:

.package(url: "https://github.com/ZachNagengast/similarity-search-kit.git", from: "0.0.1")

Then, add the appropriate target dependency to the desired target:

.target(name: "YourTarget", dependencies: [
    "SimilaritySearchKit", 
    "SimilaritySearchKitDistilbert", 
    "SimilaritySearchKitMiniLMMultiQA", 
    "SimilaritySearchKitMiniLMAll"
])

If you only want to use a subset of the available models, you can omit the corresponding dependency. This will reduce the size of your final binary.

Usage

To use SimilaritySearchKit in your project, first import the framework:

import SimilaritySearchKit

Next, create an instance of SimilarityIndex with your desired distance metric and embedding model (see below for options):

let similarityIndex = await SimilarityIndex(
    model: NativeEmbeddings(),
    metric: CosineSimilarity()
)

Then, add your text that you want to make searchable to the index:

await similarityIndex.addItem(
    id: "id1", 
    text: "Metal was released in June 2014.", 
    metadata: ["source": "example.pdf"]
)

Finally, query the index for the most similar items to a given query:

let results = await similarityIndex.search("When was metal released?")
print(results)

Which outputs a SearchResult array: [SearchResult(id: "id1", score: 0.86216, metadata: ["source": "example.pdf"])]

Examples

The Examples directory contains multple sample iOS and macOS applications that demonstrates how to use SimilaritySearchKit to it's fullest extent.

ExampleDescriptionRequirements
BasicExampleA basic multiplatform application that indexes and compares similarity of a small set of hardcoded strings.iOS 16.0+, macOS 13.0+
PDFExampleA mac-catalyst application that enables semantic search on the contents of individual PDF files.iOS 16.0+
ChatWithFilesExampleAn advanced macOS application that indexes any/all text files on your computer.macOS 13.0+

Available Models

ModelUse CaseSizeSource
NaturalLanguageText similarity, faster inferenceBuilt-inApple
MiniLMAllText similarity, fastest inference46 MBHuggingFace
DistilbertQ&A search, highest accuracy86 MB (quantized)HuggingFace
MiniLMMultiQAQ&A search, fastest inference46 MBHuggingFace

Models conform the the EmbeddingProtocol and can be used interchangeably with the SimilarityIndex class.

A small but growing list of pre-converted models can be found in this repo on HuggingFace. If you have a model that you would like to see added to the list, please open an issue or submit a pull request.

Available Metrics

MetricDescription
DotProductMeasures the similarity between two vectors as the product of their magnitudes
CosineSimilarityCalculates similarity by measuring the cosine of the angle between two vectors
EuclideanDistanceComputes the straight-line distance between two points in Euclidean space

Metrics conform to the DistanceMetricProtocol and can be used interchangeably with the SimilarityIndex class.

Bring Your Own

All the main parts of the SimilarityIndex can be overriden with custom implementations that conform to the following protocols:

EmbeddingsProtocol

Accepts a string and returns an array of floats representing the embedding of the input text.

func encode(sentence: String) async -> [Float]?

DistanceMetricProtocol

Accepts a query embedding vector and a list of embeddings vectors and returns a tuple of the distance metric score and index of the nearest neighbor.

func findNearest(for queryEmbedding: [Float], in neighborEmbeddings: [[Float]], resultsCount: Int) -> [(Float, Int)]

TextSplitterProtocol

Splits a string into chunks of a given size, with a given overlap. This is useful for splitting long documents into smaller chunks for embedding. It returns the list of chunks and an optional list of tokensIds for each chunk.

func split(text: String, chunkSize: Int, overlapSize: Int) -> ([String], [[String]]?)

TokenizerProtocol

Tokenizes and detokenizes text. Use this for custom models that use different tokenizers than are available in the current list.

func tokenize(text: String) -> [String]
func detokenize(tokens: [String]) -> String

VectorStoreProtocol

Save and load index items. The default implementation uses JSON files, but this can be overriden to use any storage mechanism.

func saveIndex(items: [IndexItem], to url: URL, as name: String) throws -> URL
func loadIndex(from url: URL) throws -> [IndexItem]
func listIndexes(at url: URL) -> [URL]

Acknowledgements

Many parts of this project were derived from the existing code, either already in swift, or translated into swift thanks to ChatGPT. These are some of the main projects that were referenced:

Motivation

This project has been inspired by the incredible advancements in natural language services and applications that have come about with the emergence of ChatGPT. While these services have unlocked a whole new world of powerful text-based applications, they often rely on cloud services. Specifically, many "Chat with Data" services necessitate users to upload their data to remote servers for processing and storage. Although this works for some, it might not be the best fit for those in low connectivity environments, or handling confidential or sensitive information. While Apple does have bundled library NaturalLanguage for similar tasks, the CoreML model conversion process opens up a much wider array of models and use cases. With this in mind, SimilaritySearchKit aims to provide a robust, on-device solution that enables developers to create state-of-the-art NLP applications within the Apple ecosystem.

Future Work

Here's a short list of some features that are planned for future releases:

I'm curious to see how people use this library and what other features would be useful, so please don't hesitate to reach out over twitter @ZachNagengast or email znagengast (at) gmail (dot) com.