使用 Pinecone 进行检索增强生成

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

  1. 使用 Voyage AI 的嵌入模型嵌入数据集
  2. 将嵌入上传到 Pinecone 索引
  3. 从向量数据库检索信息
  4. 使用 Claude 回答来自数据库的信息

设置

首先,让我们安装必要的库并设置我们将在此笔记本中使用的 API 密钥。我们需要获取一个 Claude API 密钥、一个免费的 Pinecone API 密钥 和一个免费的 Voyage AI API 密钥

%pip install anthropic datasets pinecone-client voyageai
#在此处插入您的 API 密钥
ANTHROPIC_API_KEY="<YOUR_ANTHROPIC_API_KEY>"
PINECONE_API_KEY="<YOUR_PINECONE_API_KEY>"
VOYAGE_API_KEY="<YOUR_VOYAGE_API_KEY>"

下载数据集

现在让我们下载包含 10k 多个亚马逊产品描述的亚马逊产品数据集,并将其加载到 DataFrame 中。

import pandas as pd

# 下载 JSONL 文件
!wget https://www-cdn.anthropic.com/48affa556a5af1de657d426bcc1506cdf7e2f68e/amazon-products.jsonl

data = []
with open('amazon-products.jsonl', 'r') as file:
    for line in file:
        try:
            data.append(eval(line))
        except:
            pass

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

向量数据库

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

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”模型来创建嵌入,因此我们将维度设置为 1024。

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

# 检查索引是否已存在(如果是第一次,则不应存在)
if index_name not in existing_indexes:
    # 如果不存在,则创建索引
    pc.create_index(
        index_name,
        dimension=1024,  # voyage-2 嵌入的维度
        metric='dotproduct',
        spec=spec
    )
    # 等待索引初始化
    while not pc.describe_index(index_name).status['ready']:
        time.sleep(1)

# 连接到索引
index = pc.Index(index_name)
time.sleep(1)
# 查看索引统计信息
index.describe_index_stats()

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

嵌入

要开始使用 Voyage 的嵌入,请访问此处获取 API 密钥。

现在让我们设置我们的 Voyage 客户端,并演示如何使用 embed 方法创建嵌入。要了解有关将 Voyage 嵌入与 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 索引

设置好嵌入模型后,我们现在可以获取产品描述,嵌入它们,并将嵌入上传到 Pinecone 索引。

from tqdm.auto import tqdm
from time import sleep

descriptions = df["text"].tolist()
batch_size = 100  # 一次创建和插入多少个嵌入

for i in tqdm(range(0, len(descriptions), batch_size)):
    # 查找批次的结束
    i_end = min(len(descriptions), i+batch_size)
    descriptions_batch = descriptions[i:i_end]
    # 创建嵌入(添加 try-except 以避免 RateLimitError。Voyage 目前每分钟允许 300 个请求。)
    done = False
    while not done:
        try:
            res = vo.embed(descriptions_batch, model="voyage-2", input_type="document")
            done = True
        except:
            sleep(5)

    embeds = [record for record in res.embeddings]
    # 为每个文本创建唯一 ID
    ids_batch = [f"description_{idx}" for idx in range(i, i_end)]

    # 为每个文本创建元数据字典
    metadata_batch = [{'description': description} for description in descriptions_batch]

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

    # 插入到 Pinecone
    index.upsert(vectors=to_upsert)

进行查询

在填充了索引后,我们可以开始进行查询以获取结果。我们可以采用自然语言问题,对其进行嵌入,然后在索引中查询它,以返回语义上相似的产品描述。

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 客户端并创建好提示后,我们现在可以开始从问题中生成关键字。我们将以 JSON 对象的形式输出关键字,以便可以轻松地从 Claude 的输出中解析它们。

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

# 从 JSON 中提取关键字
data = json.loads(keyword_json)
keywords_list = data['keywords']
print(keywords_list)

现在我们有了关键字列表,让我们嵌入每个关键字,在索引中查询它们,并返回最相关的 3 个产品描述。

results_list = []
for keyword in keywords_list:
    # 获取关键字的嵌入
    query_embed = vo.embed([keyword], model="voyage-2", input_type="query")
    # 在 Pinecone 索引中搜索嵌入
    search_results = index.query(vector=query_embed.embeddings, top_k=3, include_metadata=True)
    # 将搜索结果附加到列表中
    for search_result in search_results.matches:
            results_list.append(search_result['metadata']['description'])
print(len(results_list))

使用 Claude 回答

现在我们有了一系列产品描述,让我们将它们格式化为 Claude 已经训练过的搜索模板,并将格式化的描述传递到另一个提示中。

# 格式化搜索结果
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)
 为了让您的女儿对科学更感兴趣,我建议给她买一套适合她年龄的科学工具包或套装,让她能够动手进行探索和实验。例如,对于年幼的孩子,您可以尝试入门级的化学套装、磁铁套装或水晶生长套件。对于年龄较大的孩子,可以寻找解决更高级科学原理的工具包,如物理学、工程学、机器人学等。关键是选择能够激发她天生的好奇心,并让她通过活动、观察和发现来积极探索概念的东西。用科学书籍、博物馆参观、纪录片以及关于她在日常生活中遇到的科学的谈话来补充这些工具包。让科学变得有趣和引人入胜是培养她兴趣的关键。