SimilaritySearchKit

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.

Use Cases

Some potential use cases for SimilaritySearchKit include:

Privacy-focused document search engines: Create a search engine that processes sensitive documents locally, without exposing user data to external services. (See example project "ChatWithFilesExample" in the Examples directory.)
Offline question-answering systems: Implement a question-answering system that finds the most relevant answers to a user's query within a local dataset.
Document clustering and recommendation engines: Automatically group and organize documents based on their textual content on the edge.

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:

File → Add Packages... → Search or Enter Package Url → https://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):

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.

Example	Description	Requirements
`BasicExample`	A basic multiplatform application that indexes and compares similarity of a small set of hardcoded strings.	iOS 16.0+, macOS 13.0+
`PDFExample`	A mac-catalyst application that enables semantic search on the contents of individual PDF files.	iOS 16.0+
`ChatWithFilesExample`	An advanced macOS application that indexes any/all text files on your computer.	macOS 13.0+

Available Models

Model	Use Case	Size	Source
`NaturalLanguage`	Text similarity, faster inference	Built-in	Apple
`MiniLMAll`	Text similarity, fastest inference	46 MB	HuggingFace
`Distilbert`	Q&A search, highest accuracy	86 MB (quantized)	HuggingFace
`MiniLMMultiQA`	Q&A search, fastest inference	46 MB	HuggingFace

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

Available Metrics

Metric	Description
`DotProduct`	Measures the similarity between two vectors as the product of their magnitudes
`CostineSimilarity`	Calculates similarity by measuring the cosine of the angle between two vectors
`EuclideanDistance`	Computes 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:

EmbeddingProtocol

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:

HuggingFace Transformers
- https://github.com/huggingface/transformers
- https://github.com/huggingface/swift-coreml-transformers
Sentence Transformers https://github.com/UKPLab/sentence-transformers
LangChain https://github.com/hwchase17/langchain
Chroma https://github.com/chroma-core/chroma
Pinecone https://github.com/pinecone-io/examples
OpenAI Plugins https://github.com/openai/chatgpt-retrieval-plugin

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.

kebbbnnn/similarity-search-kit