在 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 食譜