#最近邻搜索 #向量搜索 #搜索引擎 #索引 #单文件 #近似 #度量

usearch

来自Unum的更小、更快的单文件向量搜索引擎

128个版本 (66个稳定版)

2.14.0 2024年8月19日
2.12.0 2024年4月29日
2.10.0 2024年3月31日
2.8.14 2023年11月26日
0.22.3 2023年7月20日

#86 in 调试

Download history 439/week @ 2024-05-03 511/week @ 2024-05-10 439/week @ 2024-05-17 520/week @ 2024-05-24 653/week @ 2024-05-31 694/week @ 2024-06-07 534/week @ 2024-06-14 568/week @ 2024-06-21 483/week @ 2024-06-28 536/week @ 2024-07-05 524/week @ 2024-07-12 462/week @ 2024-07-19 1376/week @ 2024-07-26 719/week @ 2024-08-02 618/week @ 2024-08-09 759/week @ 2024-08-16

每月3,567次下载
用于 3 个代码库

Apache-2.0

235KB
1.5K SLoC

Rust 1K SLoC // 0.1% comments C++ 155 SLoC // 0.0% comments Python 68 SLoC

USearch

更小 & 更快的 单文件
相似度搜索 & 聚类引擎,用于 向量 & 🔜 文本


Discord     LinkedIn     Twitter     Blog     GitHub

空间 • 二进制 • 概率性 • 用户自定义度量
C++ 11Python 3JavaScriptJavaRustC 99Objective-CSwiftC#GoLangWolfram
Linux • MacOS • Windows • iOS • WebAssembly • SQLite3


  • 比FAISS快10倍 的HNSW实现。
  • ✅ 简单易扩展的单个C++11头文件
  • ✅ 获得像Google这样的巨头和ClickHouse、DuckDB这样的数据库的信任。
  • ✅ SIMD优化和JIT编译的用户自定义度量。
  • ✅ 支持硬件无关的f16和i8 - 半精度和四分之一精度支持
  • ✅ 无需加载到RAM即可从磁盘查看大型索引。
  • ✅ 异构查找、重命名/重新标记和即时删除。
  • ✅ 二进制Tanimoto和Sorensen系数,适用于基因组学和化学应用
  • ✅ 使用uint40_t的内存高效点云,可容纳4B+大小。
  • ✅ 兼容OpenMP和自定义"执行器"以实现细粒度并行。
  • 语义搜索连接
  • 🔄 对数十或数百万个集群进行近实时聚类和子聚类

技术见解和相关文章

与FAISS的比较

FAISS是高性能向量搜索引擎的公认标准。USearch和FAISS都采用相同的HNSW算法,但在设计原则上存在显著差异。USearch紧凑且兼容性广泛,没有牺牲性能,主要关注用户自定义指标和更少的依赖项。

FAISS USearch 改进
索引时间 ⁰
100百万96d f32f16i8向量 2.6 · 2.6 · 2.6 h 0.3 · 0.2 · 0.2 h 9.6 · 10.4 · 10.7 x
100百万1536d f32f16i8向量 5.0 · 4.1 · 3.8 h 2.1 · 1.1 · 0.8 h 2.3 · 3.6 · 4.4 x
代码库长度 ¹ 84 K SLOC 3 K SLOC 可维护性
支持的指标 ² 9 个固定指标 任何指标 可扩展性
支持的语言 ³ C++, Python 10 种语言 可移植性
支持的 ID 类型 ⁴ 32 位,64 位 32 位,40 位,64 位 高效性
过滤 ⁵ 禁止列表 任何谓词 可组合性
所需依赖 ⁶ BLAS, OpenMP - 轻量级
绑定 ⁷ SWIG 本地 低延迟
Python 绑定大小 ⁸ ~ 10 MB < 1 MB 可部署

测试 在 Intel Sapphire Rapids 上,使用最简单的内积距离,等效召回率和内存消耗,同时提供远超搜索速度。 ¹ 相比于 faiss/ 中的 usearch/ 代码库更短,使得项目更易于维护和审核。 ² 用户定义的指标允许您针对各种应用自定义搜索,从 GIS 到创建来自多个 AI 模型或混合全文和语义搜索的复合嵌入的自定义指标。 ³ 使用 USearch,您可以在各种编程语言中重用相同的预构建索引。 ⁴ 40 位整数允许您存储 4B+ 向量,而无需为邻近图中的每个邻居引用分配 8 字节。 ⁵ 使用 USearch,索引可以与任意外部容器(如 Bloom 过滤器或第三方数据库)结合使用,以在索引遍历时过滤掉无关键。 ⁶ 缺乏强制依赖使得 USearch 更易于移植。 ⁷ 本地绑定引入的调用延迟低于更直接的方法。 ⁸ 更轻量级的绑定使下载和部署更快。

基本功能与 FAISS 相同,如果您曾研究过近似最近邻搜索,则界面必须很熟悉

# pip install usearch

import numpy as np
from usearch.index import Index

index = Index(ndim=3)               # Default settings for 3D vectors
vector = np.array([0.2, 0.6, 0.4])  # Can be a matrix for batch operations
index.add(42, vector)               # Add one or many vectors in parallel
matches = index.search(vector, 10)  # Find 10 nearest neighbors

assert matches[0].key == 42
assert matches[0].distance <= 0.001
assert np.allclose(index[42], vector, atol=0.1) # Ensure high tolerance in mixed-precision comparisons

始终有更多设置可用,API 设计得尽可能灵活。 默认的存储/量化级别取决于硬件以实现效率,但推荐大多数现代 CPU 使用 bf16

index = Index(
    ndim=3, # Define the number of dimensions in input vectors
    metric='cos', # Choose 'l2sq', 'ip', 'haversine' or other metric, default = 'cos'
    dtype='bf16', # Store as 'f64', 'f32', 'f16', 'i8', 'b1'..., default = None
    connectivity=16, # Optional: Limit number of neighbors per graph node
    expansion_add=128, # Optional: Control the recall of indexing
    expansion_search=64, # Optional: Control the quality of the search
    multi=False, # Optional: Allow multiple vectors per key, default = False
)

从磁盘序列化 & 服务器 Index

USearch 支持多种序列化形式

  • 到由路径定义的 文件
  • 到由回调定义的 ,序列化或重建增量。
  • 到固定长度或支持随机访问的内存映射文件的 缓冲区

后者允许您从外部内存提供服务索引,从而使您能够优化服务器选择以实现索引速度和服务的成本。 这可以在 AWS 和其他公共云上实现 20x 成本降低

index.save("index.usearch")

loaded_copy = index.load("index.usearch")
view = Index.restore("index.usearch", view=True)

other_view = Index(ndim=..., metric=...)
other_view.view("index.usearch")

近似搜索方法(如 HNSW)在精确穷举搜索变得过于资源密集时主要使用。 这通常发生在您有一个包含数百万条条目的集合时。 对于较小的集合,我们提供了更直接的方法,使用 search 方法。

from usearch.index import search, MetricKind, Matches, BatchMatches
import numpy as np

# Generate 10'000 random vectors with 1024 dimensions
vectors = np.random.rand(10_000, 1024).astype(np.float32)
vector = np.random.rand(1024).astype(np.float32)

one_in_many: Matches = search(vectors, vector, 50, MetricKind.L2sq, exact=True)
many_in_many: BatchMatches = search(vectors, vectors, 50, MetricKind.L2sq, exact=True)

如果您传递 exact=True 参数,则系统完全绕过索引并使用来自 SimSIMD 的 SIMD 优化相似度度量对整个数据集进行穷举搜索。 当与 Google Colab 中的 FAISS 的 IndexFlatL2 相比时,USearch 可能提供高达 20 倍的性能改进

  • faiss.IndexFlatL2: 55.3 ms.
  • usearch.index.search: 2.54 毫秒

用户自定义指标

虽然大多数向量搜索包只关注两种指标,“内积距离”和“欧几里得距离”,但USearch允许任意用户自定义指标。这种灵活性使您可以根据各种应用定制搜索,从使用罕见的Haversine距离来计算地理空间坐标,到为来自多个AI模型(如联合图像-文本嵌入)创建自定义指标。您可以使用NumbaCppyyPeachPy在Python中定义您的自定义指标

from numba import cfunc, types, carray
from usearch.index import Index, MetricKind, MetricSignature, CompiledMetric

@cfunc(types.float32(types.CPointer(types.float32), types.CPointer(types.float32)))
def python_inner_product(a, b):
    a_array = carray(a, ndim)
    b_array = carray(b, ndim)
    c = 0.0
    for i in range(ndim):
        c += a_array[i] * b_array[i]
    return 1 - c

metric = CompiledMetric(pointer=python_inner_product.address, kind=MetricKind.IP, signature=MetricSignature.ArrayArray)
index = Index(ndim=ndim, metric=metric, dtype=np.float32)

在C、C++和Rust接口中实现类似效果甚至更容易。此外,与像KD-Trees和局部敏感哈希这样的旧方法不同,HNSW不需要向量具有相同的长度。它们只需可比即可。因此,您可以在不为人知的应用中使用它,例如使用GZip压缩比作为距离函数,进行相似集合的搜索或模糊文本匹配。

过滤和谓词函数

有时您可能希望将搜索结果与某些外部数据库进行交叉引用或根据某些标准进行过滤。在大多数引擎中,您必须手动执行分页请求,逐个过滤结果。在USearch中,您可以简单地将一个谓词函数传递给搜索方法,该函数将在图遍历期间直接应用。在Rust中,它看起来像这样

let is_odd = |key: Key| key % 2 == 1;
let query = vec![0.2, 0.1, 0.2, 0.1, 0.3];
let results = index.filtered_search(&query, 10, is_odd).unwrap();
assert!(
    results.keys.iter().all(|&key| key % 2 == 1),
    "All keys must be odd"
);

内存效率、向下转型和量化

训练量化模型和降维是加速向量搜索的常用方法。然而,这些方法并不总是可靠的,可能会显著影响数据的统计特性,并且在数据分布发生变化时需要定期调整。相反,我们专注于在低精度向下转型的向量上执行高精度算术。相同的索引以及addsearch操作将自动在f64_tf32_tf16_ti8_t和单比特b1x8_t表示之间进行下转型或上转型。您可以使用以下命令来检查是否启用了硬件加速

$ python -c 'from usearch.index import Index; print(Index(ndim=768, metric="cos", dtype="f16").hardware_acceleration)'
> sapphire
$ python -c 'from usearch.index import Index; print(Index(ndim=166, metric="tanimoto").hardware_acceleration)'
> ice

在大多数情况下,建议在现代硬件上使用半精度浮点数。当启用量化时,“get”之类的函数将无法恢复原始数据,因此您可能希望在别处复制原始向量。当量化到i8_t整数时,请注意,它仅适用于余弦类指标。作为量化过程的一部分,向量被归一化到单位长度,然后缩放到[-127, 127]范围以占用完整的8位范围。当量化到b1x8_t单比特表示时,请注意,它仅适用于Jaccard、Hamming等二进制指标。作为量化过程的一部分,大于零的标量分量被设置为true,其余设置为false

USearch uint40_t support

使用更小的数字类型可以节省存储向量的内存,但你也可以压缩形成我们的邻近图中的邻居列表。默认情况下,使用32位的uint32_t来枚举这些,但如果你需要处理超过40亿的条目,这就不够用了。对于这种情况,我们提供了一个自定义的uint40_t类型,它仍然比常用的8字节整数节省37.5%的空间,并且可以扩展到1000亿条条目。

Indexes用于多索引查找

对于针对数十亿甚至数千亿个向量的较大工作负载,并行多索引查找变得非常有价值。你不需要构建一个庞大的索引,而是可以构建多个较小的索引并将它们一起查看。

from usearch.index import Indexes

multi_index = Indexes(
    indexes: Iterable[usearch.index.Index] = [...],
    paths: Iterable[os.PathLike] = [...],
    view: bool = False,
    threads: int = 0,
)
multi_index.search(...)

聚类

一旦构建了索引,USearch可以比像SciPy、UMap和tSNE这样的独立聚类库更快地执行K-Nearest Neighbors聚类。对于PCA的降维也是如此。本质上,Index本身可以看作是一种聚类,允许迭代深化。

clustering = index.cluster(
    min_count=10, # Optional
    max_count=15, # Optional
    threads=..., # Optional
)

# Get the clusters and their sizes
centroid_keys, sizes = clustering.centroids_popularity

# Use Matplotlib to draw a histogram
clustering.plot_centroids_popularity()

# Export a NetworkX graph of the clusters
g = clustering.network

# Get members of a specific cluster
first_members = clustering.members_of(centroid_keys[0])

# Deepen into that cluster, splitting it into more parts, all the same arguments supported
sub_clustering = clustering.subcluster(min_count=..., max_count=...)

生成的聚类与K-Means或其他传统方法并不相同,但具有相同的目的。另一方面,在包含100万个点的数据集上使用Scikit-Learn,你可能预计查询可能需要从几分钟到几小时,具体取决于你想要突出显示的聚类数量。对于50,000个聚类,USearch与传统聚类方法之间的性能差异可能轻易达到100倍。

连接,一对一,一对多,和多对多映射

如今,一个重要的问题是AI将如何改变数据库和数据管理的世界。大多数数据库仍在努力实现高质量的模糊搜索,并且它们所知道的唯一连接类型是确定性的。一个join与搜索每个条目不同,需要一个一对一的映射,以防止不同搜索结果之间的冲突。

精确搜索 模糊搜索 语义搜索?
精确连接 模糊连接? 语义连接?

使用USearch,可以实施亚二次复杂度的近似、模糊和语义连接。这在数据库管理系统软件中常见的模糊匹配任务中可能很有用。

men = Index(...)
women = Index(...)
pairs: dict = men.join(women, max_proposals=0, exact=False)

更多信息请参阅帖子:用于语义搜索的组合稳定婚姻💍

功能

到目前为止,核心功能已支持所有绑定。更广泛的功能可以根据需求进行移植。在某些情况下,如批处理操作,功能兼容性没有意义,因为宿主语言具有完整的多线程能力,而USearch索引结构是并发设计的,因此用户可以根据其应用程序以最优化方式实现批处理/调度/负载均衡。

C++ 11 Python 3 C 99 Java JavaScript Rust GoLang Swift
添加、搜索、删除
保存、加载、查看
用户定义的度量
批处理操作
过滤谓词
连接
可变长度向量
4B+容量

应用实例

AI有越来越多的应用,但其中最酷的古典想法之一是将其用于语义搜索。你可以使用一个编码模型,如多模态的UForm,以及一个Web编程框架,如UCall,只需20行Python代码即可构建一个文本到图像搜索平台。

from ucall import Server
from uform import get_model, Modality
from usearch.index import Index

import numpy as np
import PIL as pil

processors, models = get_model('unum-cloud/uform3-image-text-english-small')
model_text = models[Modality.TEXT_ENCODER]
model_image = models[Modality.IMAGE_ENCODER]
processor_text = processors[Modality.TEXT_ENCODER]
processor_image = processors[Modality.IMAGE_ENCODER]

server = Server()
index = Index(ndim=256)

@server
def add(key: int, photo: pil.Image.Image):
    image = processor_image(photo)
    vector = model_image(image)
    index.add(key, vector.flatten(), copy=True)

@server
def search(query: str) -> np.ndarray:
    tokens = processor_text(query)
    vector = model_text(tokens)
    matches = index.search(vector.flatten(), 3)
    return matches.keys

server.run()

相似的经验也可以在其他语言和客户端实现,从而消除网络延迟。对于 Swift 和 iOS,请查看 ashvardanian/SwiftSemanticSearch 仓库。

SwiftSemanticSearch demo Dog SwiftSemanticSearch demo with Flowers

GitHub 上有更完整的 Streamlit 示例。我们预处理了一些常用数据集,清理了图像,生成了向量,并预先构建了索引。

数据集 模态 图像 下载
Unsplash 图像 & 描述 25 K HuggingFace / Unum
概念性标题 图像 & 描述 3 M HuggingFace / Unum
Arxiv 标题 & 摘要 2 M HuggingFace / Unum

比较分子图并搜索相似结构既昂贵又缓慢。它可以被视为 NP-完全子图同构问题的特例。幸运的是,存在一些领域特定的近似方法。化学中常用的方法是从 SMILES 生成结构,并将其后来哈希到二进制指纹中。后者可以用二进制相似度指标(如 Tanimoto 系数)进行搜索。下面是使用 RDKit 包的示例。

from usearch.index import Index, MetricKind
from rdkit import Chem
from rdkit.Chem import AllChem

import numpy as np

molecules = [Chem.MolFromSmiles('CCOC'), Chem.MolFromSmiles('CCO')]
encoder = AllChem.GetRDKitFPGenerator()

fingerprints = np.vstack([encoder.GetFingerprint(x) for x in molecules])
fingerprints = np.packbits(fingerprints, axis=1)

index = Index(ndim=2048, metric=MetricKind.Tanimoto)
keys = np.arange(len(molecules))

index.add(keys, fingerprints)
matches = index.search(fingerprints, 10)

该方法用于构建 "USearch Molecules",这是最大的化学信息数据集之一,包含 70 亿个小分子和 280 亿个指纹。

USearch + POI 坐标 = 地理信息系统应用

与向量和分子搜索类似,USearch 可以用于地理信息系统。Haversine 距离是现成的,但您也可以定义更复杂的关系,如考虑地球扁率的 Vincenty 公式。

from numba import cfunc, types, carray
import math

# Define the dimension as 2 for latitude and longitude
ndim = 2

# Signature for the custom metric
signature = types.float32(
    types.CPointer(types.float32),
    types.CPointer(types.float32))

# WGS-84 ellipsoid parameters
a = 6378137.0  # major axis in meters
f = 1 / 298.257223563  # flattening
b = (1 - f) * a  # minor axis

def vincenty_distance(a_ptr, b_ptr):
    a_array = carray(a_ptr, ndim)
    b_array = carray(b_ptr, ndim)
    lat1, lon1, lat2, lon2 = a_array[0], a_array[1], b_array[0], b_array[1]
    L, U1, U2 = lon2 - lon1, math.atan((1 - f) * math.tan(lat1)), math.atan((1 - f) * math.tan(lat2))
    sinU1, cosU1, sinU2, cosU2 = math.sin(U1), math.cos(U1), math.sin(U2), math.cos(U2)
    lambda_, iterLimit = L, 100
    while iterLimit > 0:
        iterLimit -= 1
        sinLambda, cosLambda = math.sin(lambda_), math.cos(lambda_)
        sinSigma = math.sqrt((cosU2 * sinLambda) ** 2 + (cosU1 * sinU2 - sinU1 * cosU2 * cosLambda) ** 2)
        if sinSigma == 0: return 0.0  # Co-incident points
        cosSigma, sigma = sinU1 * sinU2 + cosU1 * cosU2 * cosLambda, math.atan2(sinSigma, cosSigma)
        sinAlpha, cos2Alpha = cosU1 * cosU2 * sinLambda / sinSigma, 1 - (cosU1 * cosU2 * sinLambda / sinSigma) ** 2
        cos2SigmaM = cosSigma - 2 * sinU1 * sinU2 / cos2Alpha if not math.isnan(cosSigma - 2 * sinU1 * sinU2 / cos2Alpha) else 0  # Equatorial line
        C = f / 16 * cos2Alpha * (4 + f * (4 - 3 * cos2Alpha))
        lambda_, lambdaP = L + (1 - C) * f * (sinAlpha * (sigma + C * sinSigma * (cos2SigmaM + C * cosSigma * (-1 + 2 * cos2SigmaM ** 2)))), lambda_
        if abs(lambda_ - lambdaP) <= 1e-12: break
    if iterLimit == 0: return float('nan')  # formula failed to converge
    u2 = cos2Alpha * (a ** 2 - b ** 2) / (b ** 2)
    A = 1 + u2 / 16384 * (4096 + u2 * (-768 + u2 * (320 - 175 * u2)))
    B = u2 / 1024 * (256 + u2 * (-128 + u2 * (74 - 47 * u2)))
    deltaSigma = B * sinSigma * (cos2SigmaM + B / 4 * (cosSigma * (-1 + 2 * cos2SigmaM ** 2) - B / 6 * cos2SigmaM * (-3 + 4 * sinSigma ** 2) * (-3 + 4 * cos2SigmaM ** 2)))
    s = b * A * (sigma - deltaSigma)
    return s / 1000.0  # Distance in kilometers

# Example usage:
index = Index(ndim=ndim, metric=CompiledMetric(
    pointer=vincenty_distance.address,
    kind=MetricKind.Haversine,
    signature=MetricSignature.ArrayArray,
))

集成 & 用户

参考文献

@software{Vardanian_USearch_2023,
doi = {10.5281/zenodo.7949416},
author = {Vardanian, Ash},
title = {{USearch by Unum Cloud}},
url = {https://github.com/unum-cloud/usearch},
version = {2.14.0},
year = {2023},
month = oct,
}

依赖

~0.5–2MB
~29K SLoC