用OpenAI和Pinecone构建一个语义搜索引擎

作者简介：

August Wester，奥胡斯大学计算机科学学士（2017-2020），哥本哈根信息技术大学计算机科学硕士（2020-2022）。一个对研究有热情的深度学习爱好者。

原文链接：

https://sigmoidprime.com/post/searchthearxiv/

源码地址：

https://github.com/augustwester/searchthearxiv

推荐理由：

1.OpenAI 的应用小实例，易上手，是实践者精简案例；

2.从动手思考中体会，为什么OpenAI 与微软合作后，必应火了，Google担忧了。

在这篇文章中，我们将介绍如何使用OpenAI嵌入模型和Pinecone向量数据库来构建一个简单的语义搜索引擎。更具体地说，我们将看到如何构建 searchthearxiv.com，这是一个语义搜索引擎，使学生和研究人员能够使用自然语言在arXiv上搜索超过25万篇ML论文。所涉及的原则将足够普遍，你可以将同样的技术应用于你自己的数据集，这样你就可以在你自己的文件集上提高搜索效率。

概述

在我们动手之前，让我们先把我们要解决的问题分解成四个步骤，我们可以逐一进行解决：

我们需要数据。这可能是原始的PDF文件，一个SQL数据库，甚至是一个原始的JSON文件。在我们的案例中，我们将使用Kaggle上的arXiv元数据集，它包含了JSON格式的arXiv上每篇论文的元数据（如标题和摘要）。
我们需要嵌入。一旦我们有了数据，我们需要嵌入每个条目（在我们的例子中，每篇论文），这样每个条目就由一个高维向量表示。
我们需要一个向量索引。一旦我们有了我们的嵌入，我们需要能够有效地在它们之间进行搜索。像Pinecone这样的服务可以存储数以百万计的嵌入，并允许你在给定一个查询嵌入的情况下进行快速的余弦相似性搜索。
我们需要一个接口。最后一步是建立一个界面，用户可以在其中输入查询并使用自然语言搜索我们的数据库。通常情况下，这将是某种网络前端；在这里，我们将使用一个简单的命令行界面来代替。

（一）数据

Kaggle上的arXiv数据集由康奈尔大学维护，每周更新一次。它包含了所有STEM领域在arXiv上发表的每一篇论文。我们将使用这个数据集的一个子集，只保留ML论文。最新版本可以在数据集页面上手动下载，或者使用Kaggle CLI下载，效果更好：

kaggle datasets download -d Cornell-University/arxiv && unzip arxiv.zip

这将创建一个名为arxiv-metadata-oai-snapshot.json的JSON文件。关于我们如何加载和预处理数据的细节是相当无聊的，而且几乎不可能适用于你自己的数据集。(如果你仍然好奇，可以去GitHub repo查看代码。)你只需要知道我们将每个JSON对象转换为自定义的Paper表示，并过滤掉2012年以前发表的和/或不属于任何ML类别的文章。

JSON_FILE_PATH = "arxiv-metadata-oai-snapshot.json"
CATEGORIES = ["cs.cv", "cs.lg", "cs.cl", "cs.ai", "cs.ne", "cs.ro"]
START_YEAR = 2012

print("Loading data...")
papers = list(load_data(JSON_FILE_PATH, CATEGORIES, START_YEAR))

我们指定的类别大致对应于数据集中的所有ML论文，与Andrej Karpathy的arxiv-sanity-lite所使用的相同。

（二）嵌入

现代语义搜索的神奇之处在于嵌入的概念，即对其基础文本的语义进行编码的高维向量。

为了创建我们数据集中的论文的嵌入，我们将使用一个名为text-embedding-ada-002的OpenAI嵌入模型。每1K标记为0.0004美元，这是OpenAI最便宜的嵌入模型。作为参考，我们的数据集中的标记总数约为7000万个，这意味着嵌入整个东西将花费约30美元。

窥视Paper类，我们发现以下帮助程序方法：

@property
def embedding_text(self):
    text = ["Title: " + self.title,
            "By: " + self.authors_string,
            "From: " + str(self.year),
            "Abstract: " + self.abstract]
    return ". ".join(text)

因此，在嵌入论文时，我们将使用由论文标题、作者列表和出版年份组成的“增强摘要”。我们不嵌入原始摘要的原因是有机会为查询返回有用的结果，例如 yoshua Bengio 的论文（尽管这种类型的查询在实践中不会给出最有用的结果）。

危险区域：嵌入数据不是免费的，因此请小心随意运行 OpenAI 嵌入代码。调用 API 并检索嵌入很容易，只是在程序终止后立即丢失它们。下面的帮助程序函数获取字符串列表（在本例中为“增强摘要”列表），并使用指定的 OpenAI 模型嵌入它们：

import openai

openai.api_key = os.environ["OPENAI_API_KEY"]

def get_embeddings(texts, model="text-embedding-ada-002"):
    embed_data = openai.Embedding.create(input=texts, model=model)
    return embed_data["data"]

如果你检查调用这个函数的结果为 n个增强的抽象，其返回值将是一个包含有 n个1536维的向量，每个抽象都有一个。下一步是将这些嵌入存储在一个安全的地方，以便进行有效的搜索。

(三)向量索引

为了存储嵌入物，我们将使用Pinecone。Pinecone的目的是持久地存储你的嵌入物，同时使你能够使用一个简单的API有效地搜索它们。当你注册并创建了一个索引后，你可以像这样连接到它：

import pinecone

pinecone.init(api_key=os.environ["PINECONE_API_KEY"])
index_name = os.environ["PINECONE_INDEX_NAME"]
index = pinecone.Index(index_name)

我们现在创建一个新的函数，接收我们的Paper对象列表、Pinecone索引名称和OpenAI嵌入模型的名称，嵌入这些论文并将它们上传到Pinecone：

def embed_and_upsert(papers, index_name, model, batch_size=100):
    with pinecone.Index(index_name, pool_threads=5) as index:
        for i in tqdm(range(0, len(papers), batch_size)):
            batch = papers[i:i+batch_size]
            texts = [paper.embedding_text for paper in batch]
            embed_data = get_embeddings(texts, model)
        
            pc_data = [(p.id, e["embedding"], p.metadata)
                       for p, e in zip(batch, embed_data)]
            index.upsert(pc_data)

这里，p.metadata是一个字典{"title": paper.title, "authors": paper.authors, "abstract": paper.abstract, "year": paper.year, "month": paper.month}。当从Pinecone获取搜索结果时，这将使我们能够向用户显示论文。你会注意到，每篇论文也有一个唯一的id，在这种情况下，是与原始数据集中每篇论文相关的arXiv id。

如果你调用这个函数，并在Pinecone控制台关注你的索引，你应该看到向量的数量随着嵌入的上传而增加。

（四）接口

第四步，也是最后一步，是在你最近创建的矢量索引中启用搜索。幸运的是，Pinecone让这一切变得异常简单。当收到一个查询时，我们只需使用我们用于数据集的相同嵌入模型来嵌入它。然后，我们将查询嵌入发送到Pinecone，Pinecone会识别出条目索引，这些条目与查询嵌入的余弦相似度最高（也就是语义上最相似的条目）。

要看到这个动作，请运行以下脚本：

import os
import openai
import pinecone

openai.api_key = os.environ["OPENAI_API_KEY"]

def get_embeddings(texts, model="text-embedding-ada-002"):
    embed_data = openai.Embedding.create(input=texts, model=model)
    return embed_data["data"]

pinecone.init(api_key=os.environ["PINECONE_API_KEY"])
index_name = os.environ["PINECONE_INDEX_NAME"]
index = pinecone.Index(index_name)

query = input("Enter your query: ")
embed = get_embeddings(query)[0]["embedding"]
response = index.query(vector=embed, top_k=5, include_metadata=True)
matches = response["matches"]

for i, match in enumerate(matches):
    metadata = match["metadata"]
    print(f"{i+1}: {metadata['title']}")

这将产生如下的输出：

输入你的查询：只使用注意力机制的模型

1: 注意力是你所需要的一切

2: 论训练注意力模型的动态性

3: 专注于重要的事情：人类姿势估计的自我注意模型

4: 语音识别的基于注意力的模型

5: 自我注意模型在以任务为导向的对话生成系统中的应用

结论：

由于语言模型的可用性增加，建立定制的语义搜索引擎--对于一个人来说曾经是一项几乎不可能完成的任务--现在实施起来极其简单。虽然启动像searchthearxiv.com这样的网站涉及到额外的挑战，例如托管、运行网络服务器、保持数据库的更新等等，但是您现在应该有能力利用自己的数据创建自己的语义搜索引擎。

如果您想进一步深入了解searchthearxiv.com如何工作的细节，请查看相关的GitHub repo，链接如下：

https://github.com/augustwester/searchthearxiv

内容中包含的图片若涉及版权问题，请及时与我们联系删除

用OpenAI和Pinecone构建一个语义搜索引擎

评论