一个 Swift 包,为 iOS 提供本地优先的文本嵌入和语义搜索功能

许可证

相似性搜索套件

SSK徽标

SimilaritySearchKit 是一个 Swift 软件包,为 iOS 和 macOS 应用程序提供本地优先的文本嵌入和语义搜索功能。它强调速度、灵活性和隐私性,除了与自带选项无缝集成外,还支持各种内置距离指标和金属加速机器学习模型。

与文件聊天示例视频

iOS 基本示例视频

使用案例

SimilaritySearchKit的一些潜在用例包括:

  • 注重隐私的文档搜索引擎:创建一个搜索引擎,在本地处理敏感文档,而不会将用户数据暴露给外部服务。(请参阅目录中的示例项目“ChatWithFiles”。Examples

  • 文档聚类和组织:根据文档的文本内容自动分组和组织文档,所有操作均在设备上完成。

  • 问答系统:实现一个问答系统,在本地数据集中查找与用户查询最相关的答案。

通过利用 SimilaritySearchKit,开发人员可以轻松创建功能强大的应用程序,优先将数据保存在离家较近的位置,而无需在功能或性能方面进行重大权衡。

安装

要安装 SimilaritySearchKit,只需使用 Swift Package Manager 将其作为依赖项添加到您的 Swift 项目中即可。将以下行添加到 Package.swift 文件中的依赖项数组中:

.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(id: "id1", score: 0.86216, metadata: ["source": "example.pdf"])]

可用型号

用例 大小
NaturalLanguage 文本相似度,推理速度更快 内置 苹果
MiniLMAll 文本相似性,最快推理 46兆字节 拥抱脸
Distilbert 问答搜索,精度高 86 MB(量化) 拥抱脸
MiniLMMultiQA 问答搜索,最快推理 46兆字节 拥抱脸

模型符合 并且可以与类互换使用。EmbeddingProtocolSimilarityIndex

可用指标

度量 描述
DotProduct 测量两个向量之间的相似性作为其幅度的乘积
CostineSimilarity 通过测量两个向量之间角度的余弦来计算相似性
EuclideanDistance 计算欧几里得空间中两点之间的直线距离
NLDistance 内置余弦相似性

指标符合 ,可以与类互换使用。DistanceMetricProtocolSimilarityIndex

自带

可以使用符合以下协议的自定义实现覆盖 的所有主要部分:SimilarityIndex

嵌入协议

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

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

距离度量协议

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

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

矢量存储协议

保存和加载索引项。默认实现使用 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 出现后自然语言服务和应用程序的巨大进步。虽然这些服务已经打开了一个强大的基于文本的应用程序的全新世界,但它们通常依赖于云服务。具体来说,许多“与数据聊天”服务要求用户将数据上传到远程服务器进行处理和存储。虽然这适用于某些人,但它可能不是最适合那些处于低连接环境或处理机密或敏感信息的人。考虑到这一点,SimilaritySearchKit旨在提供一个强大的,本地优先的解决方案,使开发人员能够在Apple生态系统中创建最先进的NLP应用程序。

未来工作

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

  • 内存中索引
  • 磁盘支持的索引
    • 对于不适合内存的大型数据集
  • 全方位性能改进
  • Swift-DocC 网站
  • HSNW / 烦人的索引选项
  • 查询筛选器
    • 仅返回具有特定元数据的结果
  • 稀疏/密集混合搜索
    • 使用稀疏搜索查找候选结果,然后使用密集搜索重新排名
    • 更多信息在这里
  • 更多嵌入模型
  • 汇总模型
    • 可用于将多个查询结果合并为一个,并清理不相关的文本
  • 用于距离计算的金属加速度

贡献

如果您有任何想法、建议或错误要报告,请从您的分叉中打开问题或提交拉取请求。随时欢迎贡献!

关于文件结构的说明:

  • Sources/SimilaritySearchKit/Core包含主要的相似性搜索逻辑和 100% 本机运行的帮助程序方法(即无依赖项)。
  • Sources/SimilaritySearchKit/AddOns包含可选的嵌入模型,以及任何其他需要外部依赖项的逻辑,这些逻辑应作为单独的目标和导入添加。这旨在为不需要二进制文件的用户减小二进制文件的大小。

我很想知道人们如何使用这个库以及哪些其他功能会很有用,所以请不要犹豫,通过Twitter@ZachNagengast联系或发送电子邮件至znagengast (at) gmail (dot) com。

GitHub

查看 Github