SimilaritySearchKit

License

ssk-logo

SimilaritySearchKit 是一个 Swift 包，只需几行代码即可为 iOS 和 macOS 应用程序启用设备端文本嵌入和语义搜索功能。它强调速度、可扩展性和隐私，除了无缝集成您自己的选项外，还支持各种内置的最新 NLP 模型和相似度指标。

Chat With Files Example

使用案例

SimilaritySearchKit 的一些潜在使用案例包括

通过利用 SimilaritySearchKit，开发人员可以轻松创建功能强大的应用程序，这些应用程序可以将数据保存在本地，而不会在功能或性能方面做出重大妥协。

安装

要安装 SimilaritySearchKit，只需使用 Swift Package Manager 将其作为依赖项添加到您的 Swift 项目中。我个人建议使用 Xcode 方法，通过

File → Add Packages... → Search or Enter Package Url → https://github.com/ZachNagengast/similarity-search-kit.git

Xcode 应该会为您提供以下选项来选择要添加的模型（有关选择的帮助，请参见下面的 可用模型）

Xcode Swift Package Manager Import

如果您想通过 Package.swift 添加它，请将以下行添加到您的 dependencies 数组中

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

然后，将适当的目标依赖项添加到所需的目标

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

如果您只想使用可用模型的一个子集，您可以省略相应的依赖项。这将减小最终二进制文件的大小。

用法

要在您的项目中使用 SimilaritySearchKit，首先导入框架

import SimilaritySearchKit

接下来，使用您所需的距离度量和 嵌入模型 创建 SimilarityIndex 的实例（请参阅下面的选项）

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

然后，将您想要使其可搜索的文本添加到索引中

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

最后，查询索引以查找与给定查询最相似的项目

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

这会输出一个 SearchResult 数组：[SearchResult(id: "id1", score: 0.86216, metadata: ["source": "example.pdf"])]

示例

Examples 目录包含多个示例 iOS 和 macOS 应用程序，演示了如何最大限度地使用 SimilaritySearchKit。

示例 描述 要求
BasicExample 一个基本的多平台应用程序，用于索引和比较一小组硬编码字符串的相似性。 iOS 16.0+, macOS 13.0+
PDFExample 一个 mac-catalyst 应用程序，可在单个 PDF 文件的内容上启用语义搜索。 iOS 16.0+
ChatWithFilesExample 一个高级 macOS 应用程序，用于索引您计算机上的任何/所有文本文件。 macOS 13.0+

可用模型

模型 使用案例 大小 来源
NaturalLanguage 文本相似性，更快的推理 内置 Apple
MiniLMAll 文本相似性，最快的推理 46 MB HuggingFace
Distilbert 问答搜索，最高的准确性 86 MB（量化） HuggingFace
MiniLMMultiQA 问答搜索，最快的推理 46 MB HuggingFace

模型符合 EmbeddingProtocol，可以与 SimilarityIndex 类互换使用。

可以在 HuggingFace 上的此 repo 中找到一小部分但不断增长的预转换模型列表。如果您有希望添加到列表中的模型，请打开一个 issue 或提交一个 pull request。

可用指标

指标 描述
DotProduct 将两个向量之间的相似度度量为它们幅度的乘积
CosineSimilarity 通过测量两个向量之间角度的余弦值来计算相似度
EuclideanDistance 计算欧几里得空间中两点之间的直线距离

指标符合 DistanceMetricProtocol，可以与 SimilarityIndex 类互换使用。

自带选项

SimilarityIndex 的所有主要部分都可以通过符合以下协议的自定义实现来覆盖

EmbeddingsProtocol

接受一个字符串并返回一个浮点数组，表示输入文本的嵌入。

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

DistanceMetricProtocol

接受查询嵌入向量和嵌入向量列表，并返回距离度量分数和最近邻居的索引的元组。

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

TextSplitterProtocol

将字符串拆分为给定大小的块，并具有给定的重叠。这对于将长文档拆分为更小的块以进行嵌入非常有用。它返回块列表以及每个块的可选 tokensIds 列表。

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

TokenizerProtocol

标记化和反标记化文本。将其用于使用与当前列表中可用的标记器不同的自定义模型。

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

VectorStoreProtocol

保存和加载索引项。默认实现使用 JSON 文件，但可以覆盖它以使用任何存储机制。

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]

鸣谢

该项目的许多部分都源自现有代码，要么已经在 swift 中，要么由于 ChatGPT 而被翻译成 swift。以下是一些被引用的主要项目

动机

该项目的灵感来自于随着 ChatGPT 的出现而产生的自然语言服务和应用程序的令人难以置信的进步。虽然这些服务开启了一个强大的基于文本的应用程序的全新世界，但它们通常依赖于云服务。具体来说，许多“与数据聊天”服务需要用户将其数据上传到远程服务器进行处理和存储。虽然这适用于某些人，但对于那些处于低连接环境或处理机密或敏感信息的人来说，这可能不是最佳选择。虽然 Apple 确实为类似的任务捆绑了库 NaturalLanguage，但 CoreML 模型转换过程打开了更广泛的模型和用例阵列。考虑到这一点，SimilaritySearchKit 旨在提供一个强大的设备端解决方案，使开发人员能够在 Apple 生态系统中创建最先进的 NLP 应用程序。

未来工作

以下是计划在未来版本中推出的一些功能的简短列表

我很想看看人们如何使用这个库以及哪些其他功能有用，所以请随时通过 Twitter @ZachNagengast 或电子邮件 znagengast (at) gmail (dot) com 联系我。