txtai:基于 Transformer 的人工智能搜索引擎

自然语言处理领域正在迅速发展,出现了许多新的进展。大规模的通用语言模型是一种令人兴奋的新能力,使我们能够在有限的计算和人力的情况下快速添加惊人的功能。创新仍在继续,新的模型和进步似乎每周都有。本文将对 txtai 进行介绍,这是一个基于人工智能的搜索引擎,可以在任何应用程序中实现基于自然语言理解(Natural Language Understanding,NLU)的搜索。

txtai 介绍

txtai 在文本部分上建立了一个基于人工智能的索引。txtai 支持建立文本索引来执行相似度搜索,并创建基于抽取(extractive)的问答系统。tatai 是开源的,可以在 GitHub 上获得:http://github.com/neuml/txtai

txtai 是构建在以下技术栈上的:

txtai 背后的概念已用于支持下列自然语言处理应用程序:

  • cord19q:COVID-19 文献分析
  • paperai:用于医学/科学论文的人工智能文献发现和评论引擎
  • neuspo:一个以事实为导向的实时体育赛事和新闻网站
  • codequestion:直接从终端询问编码问题

安装并运行 txtai

下面的代码段展示了如何安装 txtai 并创建一个嵌入模型。

复制代码

pip install txtai

接下来,我们可以创建一个简单的内存模型,其中包含一些示例记录来尝试 txtai。

import numpy as np
from txtai.embeddings import Embeddings
# Create embeddings model, backed by sentence-transformers & transformers
embeddings = Embeddings({"method": "transformers", "path": "sentence-transformers/bert-base-nli-mean-tokens"})
sections = ["US tops 5 million confirmed virus cases",
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
"The National Park Service warns against sacrificing slower friends in a bear attack",
"Maine man wins $1M from $25 lottery ticket",
"Make huge profits without work, earn up to $100,000 a day"]
print("%-20s %s" % ("Query", "Best Match"))
print("-" * 50)
for query in ("feel good story", "climate change", "health", "war", "wildlife", "asia",
"north america", "dishonest junk"):
# Get index of best section that best matches query
uid = np.argmax(embeddings.similarity(query, sections))
print("%-20s %s" % (query, sections[uid]))

运行上面的代码将打印以下内容:

上面的示例显示,对于几乎所有的查询,实际文本并没有存储在文本部分列表中。这就是 Transformer 模型相对于基于令牌的搜索的真正威力。

建立嵌入索引

对于较小的文本列表,上述方法是有效的。但是对于较大的文档存储库,对每个查询进行标记和转换的所有嵌入是没有意义的。txtai 支持建立预计算索引,从而显著提高性能。

在上一个示例的基础上,下面的示例运行索引方法来构建和存储文本嵌入。在这种情况下,每次搜索只将查询转换为嵌入向量。

# Create an index for the list of sections
embeddings.index([(uid, text, None) for uid, text in enumerate(sections)])
print("%-20s %s" % ("Query", "Best Match"))
print("-" * 50)
# Run an embeddings search for each query
for query in ("feel good story", "climate change", "health", "war", "wildlife", "asia",
"north america", "dishonest junk"):
# Extract uid of first result
# search result format: (uid, score)
uid = embeddings.search(query, 1)[0][0]
# Print section
print("%-20s %s" % (query, sections[uid]))

再次返回相同的结果,唯一不同之处在于嵌入是预计算的。

保存并加载嵌入索引

嵌入索引可以保存到磁盘并重新加载。此事,索引并不是以增量方式创建的,需要完全重建才能合并新数据。

embeddings.save("index")
embeddings = Embeddings()
embeddings.load("index")
uid = embeddings.search("climate change", 1)[0][0]
print(sections[uid])

以上代码运行的结果:

Canada’s last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg

通过有限的代码,我们就能构建一个对自然语言有深刻理解的系统。来自 Transformer 模型的知识量是惊人的。

句子嵌入

txtai 构建句子嵌入来执行相似性搜索。txtai 获取每个文本记录条目,将其进行标记化并构建该记录的嵌入表示。在搜索时,查询被转换为文本嵌入,然后与文本嵌入的存储库进行比较。

txtai 支持两种创建文本嵌入的方法:句子转换器和词嵌入向量。这两种方法各有优点,如下所示。

Sentence Transformers

GitHub 项目地址:

https://github.com/huggingface/transformers

  • 通过 Transformer 库生成的向量的平均池化来创建单个嵌入向量。
  • 支持模型存储在 Hugging Face 模型中心或本地存储。
  • 有关如何创建自定义模型的详细信息,请参阅 Sentence Transformers ,这些模型可以保存在本地,也可以上传到 Hugging Face 模型中心。
  • 基本模型需要强大的算力(首选 GPU)。可以构建更小 / 更轻的模型,以牺牲正确率来换取速度。

Word Embeddings

  • 通过对每个单词成分进行 BM25 评分,创建单个嵌入向量。这篇参考文献对这种方法进行了详细的描述。
  • pymagnitude 库的支持下,经过预训练的词向量可从参考链接来安装。
  • 有关可以为自定义数据集构建词向量的代码,请参阅 vectors.py
  • 使用默认模型可以显著提高速度。对于较大的数据集,它在速度和准确性方面提供了很好的折衷。

大规模的相似性搜索

正如上面所讨论的,txtai 使用相似性搜索来将句子嵌入与存储库中的所有句子嵌入进行比较。我们可能会想到的第一个问题是,这如何扩展到数百万或数十亿条记录?答案是用近似最邻近搜索(Approximate Nearest Neighbor,ANN)算法。ANN 可以在大量数据的语料库上高效执行相似性查询。

Python 中有许多可用的健壮库,可支持 ANN 搜索。txtai 有一个可配置的索引后端,允许插入不同的 ANN 库。目前,txtai 支持以下这些库:

txtai 对上述每个库使用了合理的默认设置,以便尽可能轻松地启动和运行。缺省情况下,索引的选择是基于目标环境抽象的。

上面的库要么没有关联嵌入与记录 ID 的方法,要么假设 ID 是一个整数。txtai 会处理这个问题,并保留一个内部 ID 映射,允许任何 ID 类型。

每个受支持系统(以及其他系统)的基准测试可以帮助指导哪种 ANN 最适合给定的数据集。也有平台上的差异,比如,Faiss 只支持 Linux 和 macOS。

抽取式问答系统

除相似性搜索外,txtai 支持对返回结果进行抽取式问答。这一强大的功能可以让你针对搜索结果列表询问其他一系列问题。

这方面的示例用例之一是关于 Kaggle 上的 CORD-10 挑战。这项工作需要为一系列医疗查询创建摘要表,并为每个结果抽取额外的列。

下面展示了如何在 txtai 中创建一个 Extractive QA(抽取式问答)组件:

from txtai.embeddings import Embeddings
from txtai.extractor import Extractor
# Create embeddings model, backed by sentence-transformers & transformers
embeddings = Embeddings({"method": "transformers", "path": "sentence-transformers/bert-base-nli-mean-tokens"})
# Create extractor instance
extractor = Extractor(embeddings, "distilbert-base-cased-distilled-squad")

下一步是加载一组要提问的结果。下面的示例包含一系列竞赛的体育比分的文本片段:

sections = ["Giants hit 3 HRs to down Dodgers",
"Giants 5 Dodgers 4 final",
"Dodgers drop Game 2 against the Giants, 5-4",
"Blue Jays 2 Red Sox 1 final",
"Red Sox lost to the Blue Jays, 2-1",
"Blue Jays at Red Sox is over. Score: 2-1",
"Phillies win over the Braves, 5-0",
"Phillies 5 Braves 0 final",
"Final: Braves lose to the Phillies in the series opener, 5-0",
"Final score: Flyers 4 Lightning 1",
"Flyers 4 Lightning 1 final",
"Flyers win 4-1"]
# Add unique id to each section to assist with qa extraction
sections = [(uid, section) for uid, section in enumerate(sections)]
questions = ["What team won the game?", "What was score?"]
execute = lambda query: extractor(sections, [(question, query, question, False) for question in questions])
for query in ["Red Sox - Blue Jays", "Phillies - Braves", "Dodgers - Giants", "Flyers - Lightning"]:
print("----", query, "----")
for answer in execute(query):
print(answer)
print()
# Ad-hoc questions
question = "What hockey team won?"
print("----", question, "----")
print(extractor(sections, [(question, question, question, False)]))

上面示例代码的运行结果如下:

我们可以看到 Extractor(抽取器)能够理解上面部分的上下文,并且能够回答相关的问题。Extractor 组件可以使用 txtai Embeddings 索引以及外部数据存储。这种模块化允许我们选择使用 txtai 中的那些功能来创建自然语言感知的搜索系统。

延伸阅读

更为详细 txtai 示例和用例,可以在下面的 notebook 中找到。

结语

自然语言处理正在飞速发展,一年前都不可能实现的事情,现在已经成为可能。本文介绍了一个由人工智能驱动的搜索引擎 txtai,它可以快速整合强大的模型与对自然语言的深刻理解。Hugging Face 模型中心有很多基础模型和社区提供的模型,可以用来定制几乎所有数据集的搜索。可能性是无限的,我们很高兴看到人们在 txtai 之上可以建立什么!

作者介绍:

David Mezzetti,MeuML 创始人 / 首席执行官,专注于应用机器学习解决日常问题。曾与他人共同创立 Data Works,并将其打造成一家成功的 IT 服务公司。

原文链接:

https://towardsdatascience.com/introducing-txtai-an-ai-powered-search-engine-built-on-transformers-37674be252ec

  • 发表于:
  • 本文为 InfoQ 中文站特供稿件
  • 首发地址https://www.infoq.cn/article/HsxElG4ipxXRNKOZvYUF
  • 如有侵权,请联系 yunjia_community@tencent.com 删除。

扫码关注云+社区

领取腾讯云代金券