以 Hub 为后端，在 Hugging Face 上进行向量搜索

Hugging Face Hub上的数据集使用Parquet文件。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秒缩短到不足1秒；无需部署重量级向量搜索引擎，存储则由Hub承担。