在 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 上更新

开源 AI 食谱