大规模文本嵌入与句子相似性检索

关键要点

在本文中，我们介绍了如何利用 Amazon SageMaker JumpStart 来进行文本嵌入和句子相似性检索。通过使用大型语言模型LLMs，生成文本的向量表示，应用于相似性搜索、信息检索、推荐系统等场景。此外，SageMaker JumpStart 提供了无缝部署预训练模型的方法，简化了文档嵌入的过程。

文本向量或嵌入是文本的数值向量表示，它们由大型语言模型LLMs生成。经过充分预训练或从不同任务例如文本补全、问答和翻译微调后，这些文本嵌入能够捕捉输入文本的语义信息。文本嵌入使得相似性搜索、信息检索、推荐和个性化、多语言翻译等多种下游应用成为可能。

在构建智能应用之前，企业和组织需要将现有文档嵌入，这可能既昂贵又技术复杂。 Amazon SageMaker JumpStart 是一个机器学习ML中心，帮助加速这一旅程。借助 SageMaker JumpStart，您可以访问来自多个模型提供商如 Hugging Face、AI 21 Labs、Cohere 和 Meta AI的预训练、尖端文本嵌入模型。您可以通过 SageMaker JumpStart 用户界面或 SDK 将这些模型无缝部署到生产环境。此外，您的数据不会用于训练底层模型。由于所有数据都经过加密，并且不会离开其 VPC，您可以信任数据的私密性和机密性。

在本文中，我们展示了如何使用 SageMaker Python SDK 进行文本嵌入和句子相似性检索。句子相似性涉及在将文本转换为嵌入后评估两段文本之间的相似度，这对于 Retrieval Augmented Generation (RAG) 等应用是基础步骤。我们将演示如何：

在从 SageMaker JumpStart 部署的文本嵌入模型上运行推断使用自己的数据集找到输入句子的最近邻在大型文档上运行批量转换以降低成本

所有代码均可在 GitHub 获取。

通过 SageMaker JumpStart 部署文本嵌入模型

要在 Amazon SageMaker 上托管模型，第一步是设置和验证 AWS 服务的使用。在 Amazon SageMaker Studio 中，我们使用与笔记本实例相关的执行角色。以下是相关代码：

pythonimport sagemaker boto3 jsonfrom sagemakersession import Session

sagemakersession = Session()awsrole = sagemakersessiongetcalleridentityarn()awsregion = boto3Session()regionnamesess = sagemakerSession()

在 Hugging Face 上，Massive Text Embedding Benchmark (MTEB) 被提供作为各种文本嵌入任务的 排行榜。目前，共提供 129 个基准数据集，涵盖 8 种不同任务以及 113 种语言。来自 MTEB 排行榜的顶级文本嵌入模型可以在 SageMaker JumpStart 中找到，包括 bge、gte、e5 等。在本篇中，我们以 huggingfacesentencesimilaritybgelargeen 为例。我们可以使用 SageMaker SDK 部署这个最先进的文本嵌入模型：

pythonfrom sagemakerjumpstartmodel import JumpStartModel

modelid = huggingfacesentencesimilaritybgelargeentextembeddingmodel = JumpStartModel(modelid=modelid)predictor = textembeddingmodeldeploy()

文本嵌入模型查询

接下来我们详细介绍文本嵌入模型的查询。

文本到嵌入

如果您之前已经部署过 SageMaker 端点，可以通过以下方式恢复 predictor：

pythonfrom sagemakerpredictor import Predictorfrom sagemakerdeserializers import JSONDeserializerfrom sagemakerserializers import IdentitySerializer

predictor = Predictor( endpointname= deserializer=JSONDeserializer() serializer=IdentitySerializer())predictorcontenttype = application/xtext

在模型成功部署后，您可以使用 JSON 有效负载查询一批输入文本：

pythonsentences = [ # 宠物 你的狗真可爱。 你的狗太可爱了！ 你有那么可爱的狗！ # 城市 悉尼是我工作的地方。 我在悉尼工作。 # 颜色 你最喜欢什么颜色？ 你最喜欢的颜色是什么？]

predictorpredict(jsondumps(sentences)encode(utf8))

下面的图展示了这些句子的嵌入之间的相关性。

正如前面的图所示，相同主题内的内容高度相关，包括 宠物、城市 和 颜色；而不同主题之间则差异很大。这表明 LLM 生成的嵌入在这种情况下，为 bge能够准确代表语义信息。

在本篇文章中，我们使用了上述示例，并比较当前可用的不同句子嵌入模型在 SageMaker JumpStart 上的延迟。延迟指的是从用户发送请求到应用程序指示请求已完成的时间。下表中的数字表示在 mlg52xlarge 和 mlc6ixlarge 实例上使用相同输入文本批量的 100 个请求的平均延迟。

模型g52xlarge 平均延迟 (毫秒)c6ixlarge 平均延迟 (毫秒)语言支持allMiniLML6v2195279英语BGE Base En212114英语BGE Small En283456英语BGE Large En347337英语Multilingual E5 Base221118多语言Multilingual E5 Large398360多语言E5 Base256117英语E5 Base V2252123英语E5 Large322339英语E5 Large V2325331英语GTE Base222112英语GTE Small19746英语GTE Large397347英语

获取最近邻

从 SageMaker JumpStart 部署的模型还可以方便地识别查询与语料库之间的最近邻关系。当提供查询和语料库时，模型将生成 corpusid，表示在输入语料库列表中相关条目的位置，以及表示与查询的接近程度的分数。它使用以下参数：

corpus 提供要寻找最近邻的输入列表queries 提供要从语料库寻找最近邻的输入列表topk 从语料库中查找最近邻的数量mode 设置为 nncorpus 以获取输入查询在语料库中的最近邻

以下是相关代码：

pythoncorpus = [ Amazon SageMaker 是一个全面托管的服务，用于准备数据并构建、训练和部署机器学习ML模型，以满足各种用例，配备全面管理的基础设施、工具和工作流程。 Amazon SageMaker 将代码存储在 ML 存储卷中，受到安全组保护，并可选择加密存储。 Amazon SageMaker 提供了完整的端到端工作流，但您可以继续使用您现有的工具与 SageMaker 协同工作。您可以根据业务需求轻松将每个阶段的结果进出 SageMaker。]queries = [ Amazon SageMaker 是什么？ Amazon SageMaker 如何保护我的代码？ 如果我有自己的笔记本、训练或托管环境在我自己的业务环境中怎么办？]

payloadnearestneighbor = {corpus corpus queries queries topk 3 mode nncorpus}queryresponse = predictorpredict(payloadnearestneighbor)

我们得到以下输出：

python[ [ {corpusid 0 score 08992230892181396} {corpusid 2 score 08664969205856323} {corpusid 1 score 08456423282623291} ] [ {corpusid 1 score 08919335603713989} {corpusid 0 score 0840064525604248} {corpusid 2 score 08145401477813721} ] [ {corpusid 2 score 07712811231613159} {corpusid 1 score 07564010620117188} {corpusid 0 score 07525666356086731} ]]

这意味着第一个查询与第一个语料库最相似，第二个查询与第二个语料库更接近，依此类推。在这个例子中得到了正确的匹配。

exp加速器下载

我们还使用了上述示例，并比较了 SageMaker JumpStart 中当前可用的不同句子嵌入模型的延迟。下表中的数字表示在 mlg52xlarge 和 mlc6ixlarge 实例上使用相同有效负载的 100 个请求的平均延迟。

模型g52xlarge 平均延迟 (毫秒)c6ixlarge 平均延迟 (毫秒)语言支持allMiniLML6v2217691英语BGE Base En291372英语BGE Small En292124英语BGE Large En4721240英语Multilingual E5 Base30389多语言Multilingual E5 Large4711380多语言E5 Base304373英语E5 Base V231409英语E5 Large4591230英语E5 Large V24961220英语GTE Base303375英语GTE Small285129英语GTE Large4661320英语

在大数据集上获取最近邻

在向 SageMaker 调用端点时，有效负载限制为约 5MB，且请求超时设置为 1 分钟。如果语料库大小超过这些限制，您可以使用 SageMaker 训练作业来生成大型数据集的嵌入，并将其与模型一起保留在 SageMaker 端点中，因此不必作为调用有效负载的一部分传递。查找最近邻的过程是使用 SentenceTransformer 及其实用函数进行的。最近邻是基于输入句子嵌入与训练作业期间预计算句子嵌入之间的余弦相似度。

以下示例中，我们获取和准备了 AmazonSageMakerFAQs 数据集，以便在查找输入问题的最近邻时使用：

python!aws s3 cp s3//jumpstartcacheproduswest2/trainingdatasets/AmazonSageMakerFAQs/AmazonSageMakerFAQscsv AmazonSageMakerFAQscsv

import pandas as pd

data = pdreadcsv(AmazonSageMakerFAQscsv names=[Questions Answers])data[id] = dataindexdatareq = data[[id Answers]]datareqtocsv(datacsv index=False header=False)

outputbucket = sessdefaultbucket()outputprefix = jumpstartexamplesstraining

s3outputlocation = fs3//{outputbucket}/{outputprefix}/outputtrainingdatasets3path = fs3//{outputbucket}/{outputprefix}/data/datacsv

!aws s3 cp datacsv {trainingdatasets3path}

可以获取或覆盖特定算法的训练超参数：

pythonfrom sagemaker import hyperparameters

hyperparameters = hyperparametersretrievedefault(modelid=modelid modelversion=)hyperparameters[batchsize] = 64print(hyperparameters)

gtgtgt {maxseqlength None batchsize 64 storetextwithembedding True}

SageMaker 训练分为两个步骤：创建估计器对象并启动训练作业。输出是一个预打包的模型，包含用作训练数据的嵌入，您可以将其部署以获取任何输入句子的最近邻。代码如下：

pythonfrom sagemakerjumpstartestimator import JumpStartEstimator

estimator = JumpStartEstimator( modelid=modelid hyperparameters=hyperparameters outputpath=s3outputlocation)

estimatorfit( {training fs3//{outputbucket}/{outputprefix}/data})predictor = estimatordeploy()

将文本转换为嵌入的查询语法与之前相同。但获取最近邻的代码可以简化如下：

pythonpayloadnearestneighbour = { queries [R 是否受支持于 Amazon SageMaker？] topk 1 mode nntraindata}

response = predictorpredict(payloadnearestneighbour)

gtgtgt [[{id 9 score 09240573048591614}]]

data[Answers]iloc[int(response[0][0][id])]

gtgtgt 是的，R 在 Amazon SageMaker 中受到支持。您可以在 SageMaker 笔记本实例中使用 R，这些实例包括预安装的 R 内核和 reticulate 库。Reticulate 提供了 Amazon SageMaker Python SDK 的 R 接口，使 ML 从业者能够构建、训练、调整和处理 R 模型。

我们还可以使用 AmazonSageMakerFAQs 数据集中的问题查询端点，并比较返回的正确对应答案数量。下面的示例中，我们计算 top3 准确率，这意味着如果正确答案作为 top3 返回之一，则被视为正确查询。

pythontotalcorrectanswers = 0

for i in range(len(data)) question = data[Questions]iloc[i] payloadnearestneighbor = { queries [question] topk 3 mode nntraindata } response = predictorpredict(payloadnearestneighbor) responseids = [int(res[id]) for res in response[0]]

if i in responseids    totalcorrectanswers = 1else    predanswer = [data[Answers]iloc[responseid] for responseid in responseids]

print(totalcorrectanswers 100 / len(data))

gtgtgt

8116883116883118

运行批量转换以获取大数据集的嵌入

对于具有大量历史文档的企业和组织，超过单个端点实例内存的内容，您可以使用 SageMaker 批量转换以节省成本。启动批量转换作业时，SageMaker 启动必要的计算资源来处理数据。在作业执行过程中，SageMaker 自动配置和管理计算资源。当批量转换作业完成时，这些资源会自动清理，从而最大限度地降低成本。通过将大型数据集分割成较小的块并使用更多实例，您可以以相似的成本扩展计算，实现更快的推断，而无须管理基础设施。批量转换的最大有效负载为 100 MB，最大超时为 1 小时。

我们的批量转换作业的输入格式为 JSONL 文件，条目为 JSON 的一行，包含 id 和 textinputs。请参考以下代码：

pythontestdatafilename = testjsonltestdata = []

for i in range(len(data)) answer = dataloc[i Answers] payload = {id i textinputs answer} testdataappend(payload)

with open(testdatafilename w) as outfile for entry in testdata outfilewrite(f{jsondumps