开源 AI 食谱文档

在 Hugging Face 上使用 Hub 作为后端的向量搜索

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

Open In Colab

在 Hugging Face 上使用 Hub 作为后端的向量搜索

Hugging Face Hub 上的数据集依赖于 parquet 文件。我们可以使用 DuckDB 作为快速的内存数据库系统来与这些文件进行交互。DuckDB 的一个功能是 向量相似性搜索,它可以与索引一起使用,也可以不使用索引。

安装依赖

!pip install datasets duckdb sentence-transformers model2vec -q

为数据集创建嵌入

首先,我们需要为要搜索的数据集创建嵌入。我们将使用 sentence-transformers 库为数据集创建嵌入。

from sentence_transformers import SentenceTransformer
from sentence_transformers.models import StaticEmbedding

static_embedding = StaticEmbedding.from_model2vec("minishlab/potion-base-8M")
model = SentenceTransformer(modules=[static_embedding])

现在,让我们从 Hub 加载 ai-blueprint/fineweb-bbc-news 数据集。

from datasets import load_dataset

ds = load_dataset("ai-blueprint/fineweb-bbc-news")

我们现在可以为数据集创建嵌入。通常,我们可能希望将数据分块成更小的批次以避免损失精度,但对于此示例,我们将仅为数据集的全文创建嵌入。

def create_embeddings(batch):
    embeddings = model.encode(batch["text"], convert_to_numpy=True)
    batch["embeddings"] = embeddings.tolist()
    return batch


ds = ds.map(create_embeddings, batched=True)

我们现在可以将带有嵌入的数据集上传回 Hub。

ds.push_to_hub("ai-blueprint/fineweb-bbc-news-embeddings")

向量搜索 Hugging Face Hub

我们现在可以使用 duckdb 对数据集执行向量搜索。这样做时,我们可以使用索引,也可以不使用索引。不使用 索引搜索速度较慢,但更精确,而 使用 索引搜索速度更快,但精度较低。

不使用索引

要不使用索引进行搜索,我们可以使用 duckdb 库连接到数据集并执行向量搜索。这是一个缓慢的操作,但通常对于最多 10 万行的小数据集来说足够快。这意味着查询我们的数据集会稍微慢一些。

import duckdb
from typing import List


def similarity_search_without_duckdb_index(
    query: str,
    k: int = 5,
    dataset_name: str = "ai-blueprint/fineweb-bbc-news-embeddings",
    embedding_column: str = "embeddings",
):
    # Use same model as used for indexing
    query_vector = model.encode(query)
    embedding_dim = model.get_sentence_embedding_dimension()

    sql = f"""
        SELECT 
            *,
            array_cosine_distance(
                {embedding_column}::float[{embedding_dim}], 
                {query_vector.tolist()}::float[{embedding_dim}]
            ) as distance
        FROM 'hf://datasets/{dataset_name}/**/*.parquet'
        ORDER BY distance
        LIMIT {k}
    """
    return duckdb.sql(sql).to_df()


similarity_search_without_duckdb_index("What is the future of AI?")

使用索引

这种方法创建数据集的本地副本,并使用它来创建索引。这有一些小的开销,但一旦创建索引,它将显着加快搜索速度。

import duckdb


def _setup_vss():
    duckdb.sql(
        query="""
        INSTALL vss;
        LOAD vss;
        """
    )


def _drop_table(table_name):
    duckdb.sql(
        query=f"""
        DROP TABLE IF EXISTS {table_name};
        """
    )


def _create_table(dataset_name, table_name, embedding_column):
    duckdb.sql(
        query=f"""
        CREATE TABLE {table_name} AS 
        SELECT *, {embedding_column}::float[{model.get_sentence_embedding_dimension()}] as {embedding_column}_float 
        FROM 'hf://datasets/{dataset_name}/**/*.parquet';
        """
    )


def _create_index(table_name, embedding_column):
    duckdb.sql(
        query=f"""
        CREATE INDEX my_hnsw_index ON {table_name} USING HNSW ({embedding_column}_float) WITH (metric = 'cosine');
        """
    )


def create_index(dataset_name, table_name, embedding_column):
    _setup_vss()
    _drop_table(table_name)
    _create_table(dataset_name, table_name, embedding_column)
    _create_index(table_name, embedding_column)


create_index(
    dataset_name="ai-blueprint/fineweb-bbc-news-embeddings",
    table_name="fineweb_bbc_news_embeddings",
    embedding_column="embeddings",
)

现在我们可以使用索引执行向量搜索,这将立即返回结果。

def similarity_search_with_duckdb_index(
    query: str, k: int = 5, table_name: str = "fineweb_bbc_news_embeddings", embedding_column: str = "embeddings"
):
    embedding = model.encode(query).tolist()
    return duckdb.sql(
        query=f"""
        SELECT *, array_cosine_distance({embedding_column}_float, {embedding}::FLOAT[{model.get_sentence_embedding_dimension()}]) as distance 
        FROM {table_name}
        ORDER BY distance 
        LIMIT {k};
    """
    ).to_df()


similarity_search_with_duckdb_index("What is the future of AI?")

查询时间从 30 秒减少到亚秒级响应时间,并且不需要您部署重量级的向量搜索引擎,而存储由 Hub 处理。

结论

我们已经了解了如何使用 duckdb 在 Hub 上执行向量搜索。对于小于 10 万行的小数据集,我们可以使用 Hub 作为向量搜索后端执行无索引向量搜索,但对于更大的数据集,我们应该在使用本地搜索和使用 Hub 作为存储后端时使用 vss 扩展创建索引。

了解更多

< > 在 GitHub 上更新