开源 AI 食谱文档
在 Hugging Face 上使用 Hub 作为后端的向量搜索
并获得增强的文档体验
开始使用
在 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
扩展创建索引。