Pinecone RAG

使用 Pinecone 向量数据库实现 RAG

🌐 查看英文原文 | 源码 Notebook


使用 Pinecone 实现检索增强生成

本教程演示了如何通过检索增强生成(RAG)技术将 Claude 与 Pinecone 向量数据库中的数据连接起来。我们将涵盖以下步骤:

  1. 使用 Voyage AI 的 embedding 模型对数据集进行 embedding
  2. 将 embedding 上传到 Pinecone 索引
  3. 从向量数据库中检索信息
  4. 使用 Claude 根据数据库中的信息回答问题

环境配置

首先,让我们安装必要的库并设置本教程中需要使用的 API keys。我们需要获取 Claude API key、免费的 Pinecone API key 和免费的 Voyage AI API key。

%pip install anthropic datasets pinecone-client voyageai
# Insert your API keys here
ANTHROPIC_API_KEY = "<YOUR_ANTHROPIC_API_KEY>"
PINECONE_API_KEY = "<YOUR_PINECONE_API_KEY>"
VOYAGE_API_KEY = "<YOUR_VOYAGE_API_KEY>"

下载数据集

现在让我们下载包含超过 1 万条亚马逊商品描述的亚马逊商品数据集,并将其加载到 DataFrame 中。

import pandas as pd

# Download the JSONL file
!wget  https://www-cdn.anthropic.com/48affa556a5af1de657d426bcc1506cdf7e2f68e/amazon-products.jsonl

data = []
with open("amazon-products.jsonl") as file:
    for line in file:
        try:
            data.append(eval(line))  # noqa: S307
        except (SyntaxError, ValueError):
            # Skip malformed lines in the dataset
            pass

df = pd.DataFrame(data)
display(df.head())
len(df)

向量数据库

要创建我们的向量数据库,我们首先需要一个免费的 Pinecone API key。获取密钥后,我们可以按如下方式初始化数据库:

from pinecone import Pinecone

pc = Pinecone(api_key=PINECONE_API_KEY)

接下来,我们设置索引规范,用于定义我们想要部署索引的云服务提供商和区域。你可以在这里找到所有可用的提供商和区域列表。

from pinecone import ServerlessSpec

spec = ServerlessSpec(cloud="aws", region="us-west-2")

然后,我们初始化索引。我们将使用 Voyage 的”voyage-2”模型来创建 embedding,因此我们将维度设置为 1024。

import time

index_name = "amazon-products"
existing_indexes = [index_info["name"] for index_info in pc.list_indexes()]

# check if index already exists (it shouldn't if this is first time)
if index_name not in existing_indexes:
    # if does not exist, create index
    pc.create_index(
        index_name,
        dimension=1024,  # dimensionality of voyage-2 embeddings
        metric="dotproduct",
        spec=spec,
    )
    # wait for index to be initialized
    while not pc.describe_index(index_name).status["ready"]:
        time.sleep(1)

# connect to index
index = pc.Index(index_name)
time.sleep(1)
# view index stats
index.describe_index_stats()

我们应该会看到新的 Pinecone 索引的 total_vector_count 为 0,因为我们还没有添加任何向量。

Embedding

要开始使用 Voyage 的 embedding,请前往这里获取 API key。

现在让我们设置 Voyage 客户端,并演示如何使用 embed 方法创建 embedding。要了解有关将 Voyage embedding 与 Claude 结合使用的更多信息,请参见此教程。

import voyageai

vo = voyageai.Client(api_key=VOYAGE_API_KEY)

texts = ["Sample text 1", "Sample text 2"]

result = vo.embed(texts, model="voyage-2", input_type="document")
print(result.embeddings[0])
print(result.embeddings[1])

将数据上传到 Pinecone 索引

embedding 模型设置好后,我们就可以对商品描述进行 embedding,并将 embedding 上传到 Pinecone 索引。

from time import sleep

from tqdm.auto import tqdm

descriptions = df["text"].tolist()
batch_size = 100  # how many embeddings we create and insert at once

for i in tqdm(range(0, len(descriptions), batch_size)):
    # find end of batch
    i_end = min(len(descriptions), i + batch_size)
    descriptions_batch = descriptions[i:i_end]
    # create embeddings (try-except added to avoid RateLimitError. Voyage currently allows 300/requests per minute.)
    done = False
    while not done:
        try:
            res = vo.embed(descriptions_batch, model="voyage-2", input_type="document")
            done = True
        except Exception:
            sleep(5)

    embeds = [record for record in res.embeddings]
    # create unique IDs for each text
    ids_batch = [f"description_{idx}" for idx in range(i, i_end)]

    # Create metadata dictionaries for each text
    metadata_batch = [{"description": description} for description in descriptions_batch]

    to_upsert = list(zip(ids_batch, embeds, metadata_batch, strict=False))

    # upsert to Pinecone
    index.upsert(vectors=to_upsert)

进行查询

索引填充完成后,我们就可以开始进行查询以获取结果了。我们可以获取一个自然语言问题,对其进行 embedding,然后在索引中查询以返回语义相似的商品描述。

USER_QUESTION = (
    "I want to get my daughter more interested in science. What kind of gifts should I get her?"
)

question_embed = vo.embed([USER_QUESTION], model="voyage-2", input_type="query")
results = index.query(vector=question_embed.embeddings, top_k=5, include_metadata=True)
results
{'matches': [{'id': 'description_1771',
              'metadata': {'description': 'Product Name: Scientific Explorer '
                                          'My First Science Kids Science '
                                          'Experiment Kit\n'
                                          '\n'
                                          'About Product: Experiments to spark '
                                          'creativity and curiosity | Grow '
                                          'watery crystals, create a rainbow '
                                          'in a plate, explore the science of '
                                          'color and more | Represents STEM '
                                          '(Science, Technology, Engineering, '
                                          'Math) principles – open ended toys '
                                          'to construct, engineer, explorer '
                                          'and experiment | Includes cross '
                                          'linked polyacrylamide, 3 color '
                                          'tablets, 3 mixing cups, 3 test '
                                          'tubes, caps and stand, pipette, '
                                          'mixing tray, magnifier and '
                                          'instructions | Recommended for '
                                          'children 4 years of age and older '
                                          'with adult supervision\n'
                                          '\n'
                                          'Categories: Toys & Games | Learning '
                                          '& Education | Science Kits & Toys'},
              'score': 0.772703767,
              'values': []},
             {'id': 'description_3133',
              'metadata': {'description': 'Product Name: Super Science Magnet '
                                          'Kit.\n'
                                          '\n'
                                          'About Product: \n'
                                          '\n'
                                          'Categories: Toys & Games | Learning '
                                          '& Education | Science Kits & Toys'},
              'score': 0.765997052,
              'values': []},
             {'id': 'description_1792',
              'metadata': {'description': 'Product Name: BRIGHT Atom Model - '
                                          'Student\n'
                                          '\n'
                                          'About Product: \n'
                                          '\n'
                                          'Categories: Toys & Games | Learning '
                                          '& Education | Science Kits & Toys'},
              'score': 0.765654,
              'values': []},
             {'id': 'description_1787',
              'metadata': {'description': 'Product Name: Thames & Kosmos '
                                          'Biology Genetics and DNA\n'
                                          '\n'
                                          'About Product: Learn the basics of '
                                          'genetics and DNA. | Assemble a '
                                          'model to see the elegant '
                                          'double-stranded Helical structure '
                                          "of DNA. | A parents' Choice Gold "
                                          'award winner | 20 experiments in '
                                          'the 48 page full color experiment '
                                          'manual and learning guide\n'
                                          '\n'
                                          'Categories: Toys & Games | Learning '
                                          '& Education | Science Kits & Toys'},
              'score': 0.765174091,
              'values': []},
             {'id': 'description_120',
              'metadata': {'description': 'Product Name: Educational Insights '
                                          "Nancy B's Science Club Binoculars "
                                          'and Wildlife Activity Journal\n'
                                          '\n'
                                          'About Product: From bird search and '
                                          'ecosystem challenges to creative '
                                          'writing and drawing exercises, this '
                                          'set is perfect for the nature lover '
                                          'in your life! | Includes 4x '
                                          'magnification binoculars and '
                                          '22-page activity journal packed '
                                          'with scientific activities! | '
                                          'Binoculars are lightweight, yet '
                                          'durable. | Supports STEM learning, '
                                          'providing hands-on experience with '
                                          'a key scientific tool. | Great '
                                          'introductory tool for young '
                                          'naturalists on-the-go! | Part of '
                                          "the Nancy B's Science Club line, "
                                          'designed to encourage scientific '
                                          'confidence. | Winner of the '
                                          "Parents' Choice Recommended Award. "
                                          '| Scientific experience designed '
                                          'specifically for kids ages 8-11.\n'
                                          '\n'
                                          'Categories: Electronics | Camera & '
                                          'Photo | Binoculars & Scopes | '
                                          'Binoculars'},
              'score': 0.765075564,
              'values': []}],
 'namespace': '',
 'usage': {'read_units': 6}}

优化搜索

这些结果还不错,但我们还可以进一步优化。使用 Claude,我们可以从用户的问题中生成搜索关键词。这使我们能够在索引上进行更广泛、更多样化的搜索,从而获得更相关的商品描述。

import anthropic

client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)


def get_completion(prompt):
    completion = client.completions.create(
        model="claude-2.1",
        prompt=prompt,
        max_tokens_to_sample=1024,
    )
    return completion.completion
def create_keyword_prompt(question):
    return f"""\n\nHuman: Given a question, generate a list of 5 very diverse search keywords that can be used to search for products on Amazon.

The question is: {question}

Output your keywords as a JSON that has one property "keywords" that is a list of strings. Only output valid JSON.\n\nAssistant:{{"""

Anthropic 客户端设置完成且 prompt 创建好后,我们就可以开始从问题中生成关键词了。我们将关键词输出为 JSON 对象,以便从 Claude 的输出中轻松解析。

keyword_json = "{" + get_completion(create_keyword_prompt(USER_QUESTION))
print(keyword_json)
import json

# Extract the keywords from the JSON
data = json.loads(keyword_json)
keywords_list = data["keywords"]
print(keywords_list)

现在我们将关键词放入列表中,对每个关键词进行 embedding,在索引中查询,并返回最相关的前 3 个商品描述。

results_list = []
for keyword in keywords_list:
    # get the embeddings for the keywords
    query_embed = vo.embed([keyword], model="voyage-2", input_type="query")
    # search for the embeddings in the Pinecone index
    search_results = index.query(vector=query_embed.embeddings, top_k=3, include_metadata=True)
    # append the search results to the list
    for search_result in search_results.matches:
        results_list.append(search_result["metadata"]["description"])
print(len(results_list))

使用 Claude 回答

现在我们有了一个商品描述列表,让我们将它们格式化为 Claude 训练时使用的搜索模板,并将格式化后的描述传递到另一个 prompt 中。

# Formatting search results
def format_results(extracted: list[str]) -> str:
    result = "\n".join(
        [
            f'<item index="{i + 1}">\n<page_content>\n{r}\n</page_content>\n</item>'
            for i, r in enumerate(extracted)
        ]
    )

    return f"\n<search_results>\n{result}\n</search_results>"


def create_answer_prompt(results_list, question):
    return f"""\n\nHuman: {format_results(results_list)} Using the search results provided within the <search_results></search_results> tags, please answer the following question <question>{question}</question>. Do not reference the search results in your answer.\n\nAssistant:"""

最后,让我们提出原始用户的问题并从 Claude 获取答案。

answer = get_completion(create_answer_prompt(results_list, USER_QUESTION))
print(answer)
 To get your daughter more interested in science, I would recommend getting her an age-appropriate science kit or set that allows for hands-on exploration and experimentation. For example, for a younger child you could try a beginner chemistry set, magnet set, or crystal growing kit. For an older child, look for kits that tackle more advanced scientific principles like physics, engineering, robotics, etc. The key is choosing something that sparks her natural curiosity and lets her actively investigate concepts through activities, observations, and discovery. Supplement the kits with science books, museum visits, documentaries, and conversations about science she encounters in everyday life. Making science fun and engaging is crucial for building her interest.

Tip

中文读者提示:本章节的代码和输出为英文原文。如需理解具体实现细节,可参考上方中文导读和代码注释。如有疑问,欢迎在 GitHub Issues 讨论。