Fleet Context 嵌入 - 为 Llamaindex 库构建混合搜索引擎

本指南将使用 Fleet Context 下载 LlamaIndex 文档的嵌入向量，并基于此构建一个稠密/稀疏向量混合检索引擎。

前置准备

!pip install llama-index
!pip install --upgrade fleet-context

import os
import openai

os.environ["OPENAI_API_KEY"] = "sk-..." # 在此处添加你的API密钥!
openai.api_key = os.environ["OPENAI_API_KEY"]

从 Fleet Context 下载嵌入向量

我们将使用 Fleet Context 下载 LlamaIndex 全部文档的嵌入向量（约12k个文本块，100MB内容）。通过指定库名称作为参数，你可以下载1220个主流库中任意一个的嵌入数据。完整支持库列表可在页面底部此处查看。

选择 Fleet 是因为其嵌入处理流程保留了多项提升检索与生成质量的关键信息，包括：页面位置（用于重排序）、块类型（类/函数/属性等）、父章节等。更多细节可参阅其Github页面。

from context import download_embeddings

df = download_embeddings("llamaindex")

输出:

    100%|██████████| 83.7M/83.7M [00:03<00:00, 27.4MiB/s]
                                         id  \
    0  e268e2a1-9193-4e7b-bb9b-7a4cb88fc735
    1  e495514b-1378-4696-aaf9-44af948de1a1
    2  e804f616-7db0-4455-9a06-49dd275f3139
    3  eb85c854-78f1-4116-ae08-53b2a2a9fa41
    4  edfc116e-cf58-4118-bad4-c4bc0ca1495e

# 展示部分元数据示例
df["metadata"][0]
display(Markdown(f"{df['metadata'][8000]['text']}"))

输出:

classmethod from_dict(data: Dict[str, Any], kwargs: Any) → Self classmethod from_json(data_str: str, kwargs: Any) → Self classmethod from_orm(obj: Any) → Model json(, include: Optional[Union[AbstractSetIntStr, MappingIntStrAny]] = None, exclude: Optional[Union[AbstractSetIntStr, MappingIntStrAny]] = None, by_alias: bool = False, skip_defaults: Optional[bool] = None, exclude_unset: bool = False, exclude_defaults: bool = False, exclude_none: bool = False, encoder: Optional[Callable[[Any], Any]] = None, models_as_dict: bool = True*, dumps_kwargs: Any) → unicode Generate a JSON representation of the model, include and exclude arguments as per dict().

为 LlamaIndex 创建 Pinecone 混合搜索索引

我们将创建 Pinecone 索引并上传向量，以实现稀疏向量与稠密向量的混合检索。开始前请确保已注册Pinecone账户。

import logging
import sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().handlers = []
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

import pinecone

api_key = "..."  # 在此处添加你的Pinecone API密钥
pinecone.init(
    api_key=api_key, environment="us-east-1-aws"
)  # 在此处添加你的数据库区域

# Fleet Context 使用 OpenAI 的 text-embedding-ada-002 模型（1536维）

# 注意：Pinecone要求混合搜索使用点积相似度
pinecone.create_index(
    "quickstart-fleet-context",
    dimension=1536,
    metric="dotproduct",
    pod_type="p1",
)

pinecone.describe_index(
    "quickstart-fleet-context"
)  # 确保已在pinecone创建索引

from llama_index.vector_stores.pinecone import PineconeVectorStore

pinecone_index = pinecone.Index("quickstart-fleet-context")
vector_store = PineconeVectorStore(pinecone_index, add_sparse_vector=True)

批量上传向量至 Pinecone

Pinecone 建议每次上传100个向量。我们将在调整数据格式后执行此操作。

import random
import itertools


def chunks(iterable, batch_size=100):
    """将可迭代对象分割为指定批次大小的辅助函数"""
    it = iter(iterable)
    chunk = tuple(itertools.islice(it, batch_size))
    while chunk:
        yield chunk
        chunk = tuple(itertools.islice(it, batch_size))


# 生成包含(id, vector, metadata, sparse_values)的生成器
data_generator = map(
    lambda row: {
        "id": row[1]["id"],
        "values": row[1]["values"],
        "metadata": row[1]["metadata"],
        "sparse_values": row[1]["sparse_values"],
    },
    df.iterrows(),
)

# 每次上传100个向量
for ids_vectors_chunk in chunks(data_generator, batch_size=100):
    print(f"正在上传 {len(ids_vectors_chunk)} 个向量...")
    pinecone_index.upsert(vectors=ids_vectors_chunk)

在 LlamaIndex 中构建 Pinecone 向量存储

最后，我们将通过 LlamaIndex 构建 Pinecone 向量存储并进行查询。

from llama_index.core import VectorStoreIndex
from IPython.display import Markdown, display

index = VectorStoreIndex.from_vector_store(vector_store=vector_store)

执行索引查询！

query_engine = index.as_query_engine(
    vector_store_query_mode="hybrid", similarity_top_k=8
)
response = query_engine.query("如何使用 llama_index 的 SimpleDirectoryReader")

display(Markdown(f"<b>{response}</b>"))

输出:

<b>要使用llama_index中的SimpleDirectoryReader，需要先从llama_index库导入该类。导入后，通过提供目录路径作为参数创建SimpleDirectoryReader实例，然后对该实例调用`load_data()`方法即可加载指定目录中的文档。</b>