在 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 作为向量搜索后端执行向量搜索，但对于较大的数据集，我们应该在进行本地搜索时使用 `vss` 扩展创建索引，并使用 Hub 作为存储后端。

了解更多

< > 在 GitHub 上更新

开源 AI 食谱