文本向量或嵌入是由大型语言模型（LLM）生成的文本的数值向量表示。在LLM完全预训练于大型数据集或从不同任务（包括文本补全、问题回答和翻译）微调之后，文本嵌入捕捉输入文本的语义信息。文本嵌入使得包括相似性搜索、信息检索、推荐和个性化、多语言翻译等不同下游应用成为可能。
在利用嵌入构建智能应用之前，企业和组织必须将现有文档进行嵌入，这可能既昂贵又技术复杂。
Amazon SageMaker JumpStart是一个机器学习（ML）平台，可帮助加速这一过程。利用SageMaker JumpStart，您可以访问来自各种模型提供方（包括Hugging Face、AI 21 Labs、Cohere和Meta AI）的预训练、尖端文本嵌入模型。您可以通过SageMaker JumpStart用户界面或SDK轻松部署这些模型到生产环境。此外，所有数据都不用于训练底层模型。由于所有数据都经过加密，且不离开自己的VPC，您可以信任您的数据仍然是私密和保密的。
在本文中，我们演示了如何利用SageMaker Python SDK进行文本嵌入和句子相似度计算。句子相似度评估了LLM将其转换为嵌入的两段文本之间的相似程度，这是实现检索增强生成（RAG）等应用的基础步骤。我们演示了以下内容：
- 对由SageMaker JumpStart部署的文本嵌入模型进行推理
- 使用您自己的数据集找到输入句子的最近邻居
- 对大型文档运行批量转换以降低成本
所有代码都可以在GitHub上找到。
通过SageMaker JumpStart部署文本嵌入模型
要在Amazon SageMaker上托管模型，第一步是设置和验证使用AWS服务。在Amazon SageMaker Studio中，我们使用与笔记本实例关联的执行角色。请参阅以下代码：
import sagemaker, boto3, jsonfrom sagemaker.session import Sessionsagemaker_session = Session()aws_role = sagemaker_session.get_caller_identity_arn()aws_region = boto3.Session().region_namesess = sagemaker.Session()在Hugging Face上，提供了大规模文本嵌入基准（MTEB）作为各种文本嵌入任务的排行榜。目前，它提供了113种语言上8种不同任务的129个基准数据集。MTEB排行榜上的顶级文本嵌入模型可以从SageMaker JumpStart获得，包括bge、gte、e5等。在本文中，我们以huggingface-sentencesimilarity-bge-large-en为例。我们可以使用SageMaker SDK部署这个最先进的文本嵌入模型：
from sagemaker.jumpstart.model import JumpStartModelmodel_id = "huggingface-sentencesimilarity-bge-large-en"text_embedding_model = JumpStartModel(model_id=model_id)predictor = text_embedding_model.deploy()文本嵌入模型查询
让我们更详细地看一下文本嵌入模型查询。
文本转嵌入
如果您之前已经部署了SageMaker终端节点，则可以按如下恢复predictor：
from sagemaker.predictor import Predictorfrom sagemaker.deserializers import JSONDeserializerfrom sagemaker.serializers import IdentitySerializerpredictor = Predictor( endpoint_name=<YOUR_ENDPOINT_NAME>, deserializer=JSONDeserializer(), serializer=IdentitySerializer(),)predictor.content_type = "application/x-text"模型成功部署后，您可以在 JSON 负载中使用一批输入文本查询端点：
sentences = [ # 宠物 "你的狗真可爱。", "你的狗太可爱了！", "你有一只如此可爱的狗！", # 城市 "我在悉尼工作。", "我在悉尼工作。", # 颜色 "你最喜欢的颜色是什么？", "你最喜欢的颜色是什么？",]predictor.predict(json.dumps(sentences).encode('utf-8'))这些句子的嵌入相关性在下图中绘制出来。
如上图所示，相同主题之间高度相关，包括 Pets、Cities 和 Color；不同主题之间非常不相关。这表明 LLMs 生成的嵌入能够准确表示语义信息。
对于本文，我们使用上述示例比较了目前在 SageMaker JumpStart 中提供的不同句子嵌入模型的延迟。延迟是指从用户发送请求到应用程序表示请求已完成的时间量。下表中的数字表示在 ml.g5.2xlarge 和 ml.c6i.xlarge 实例上使用相同批量输入文本进行的总共 100 次请求的平均延迟。
| 模型 | g5.2xlarge 平均延迟 (ms) | c6i.xlarge 平均延迟 (ms) | 语言支持 |
|---|---|---|---|
| all-MiniLM-L6-v2 | 19.5 | 27.9 | 英语 |
| BGE Base En | 21.2 | 114 | 英语 |
| BGE Small En | 28.3 | 45.6 | 英语 |
| BGE Large En | 34.7 | 337 | 英语 |
| Multilingual E5 Base | 22.1 | 118 | 多语言 |
| Multilingual E5 Large | 39.8 | 360 | 多语言 |
| E5 Base | 25.6 | 117 | 英语 |
| E5 Base V2 | 25.2 | 123 | 英语 |
| E5 Large | 32.2 | 339 | 英语 |
| E5 Large V2 | 32.5 | 331 | 英语 |
| GTE Base | 22.2 | 112 | 英语 |
| GTE Small | 19.7 | 46 | 英语 |
| GTE Large | 39.7 | 347 | 英语 |
获取最近的邻居
来自SageMaker JumpStart的部署模型还可以促进在语料库中识别查询的最近邻居的过程。当提供查询和语料库时，模型将产生corpus_id，它表示输入语料库列表中相关语料库条目的位置，以及指示与查询的接近程度的分数。它使用以下参数：
- corpus – 提供要从中查找最近邻居的输入列表
- queries – 提供要从语料库中查找最近邻居的输入列表
- top_k – 从语料库中查找的最近邻居的数量
- mode – 设置为
nn_corpus以获取对输入查询在语料库中的最近邻居
请参阅以下代码：
corpus = [ "Amazon SageMaker是一项完全托管的服务，用于准备数据并构建、训练和部署适用于任何用例的机器学习（ML）模型，拥有完全托管的基础架构、工具和工作流程。", "Amazon SageMaker将代码存储在ML存储卷中，由安全组安全，并可选择在静止时进行加密。", "Amazon SageMaker提供完整的端到端工作流程，但您可以继续使用现有工具与SageMaker一起使用。您可以轻松地在SageMaker内部和外部的每个阶段之间传输结果，以适应您的业务要求。"]queries = [ "什么是Amazon SageMaker？", "Amazon SageMaker如何保护我的代码？", "如果我在自己的业务环境中拥有自己的笔记本、培训或托管环境，该怎么办？"]payload_nearest_neighbor = {"corpus": corpus, "queries": queries, "top_k": 3, "mode": "nn_corpus"}query_response = predictor.predict(payload_nearest_neighbor)我们得到以下输出：
[ [ {'corpus_id': 0, 'score': 0.8992230892181396}, {'corpus_id': 2, 'score': 0.8664969205856323}, {'corpus_id': 1, 'score': 0.8456423282623291} ], [ {'corpus_id': 1, 'score': 0.8919335603713989}, {'corpus_id': 0, 'score': 0.840064525604248}, {'corpus_id': 2, 'score': 0.8145401477813721} ], [ {'corpus_id': 2, 'score': 0.7712811231613159}, {'corpus_id': 1, 'score': 0.7564010620117188}, {'corpus_id': 0, 'score': 0.7525666356086731} ]]这个结果表示第一个查询最类似于第一个语料库，第二个查询更接近于第二个语料库，以此类推。在这个例子中是正确的匹配。
我们还对前面的示例进行了比较，比较了当前从SageMaker JumpStart获取的不同句子嵌入模型之间的延迟。以下表格中的数字表示在ml.g5.2xlarge和ml.c6i.xlarge实例上使用相同负载的100个请求的平均延迟。
| 模型 | g5.2xlarge平均延迟（毫秒） | c6i.xlarge平均延迟（毫秒） | 语言支持 |
|---|---|---|---|
| all-MiniLM-L6-v2 | 21.7 | 69.1 | 英语 |
| BGE Base En | 29.1 | 372 | 英语 |
| BGE Small En | 29.2 | 124 | 英语 |
| BGE Large En | 47.2 | 1240 | 英语 |
| Multilingual E5 Base | 30 | 389 | 多语言 |
| Multilingual E5 Large | 47.1 | 1380 | 多语言 |
| E5 Base | 30.4 | 373 | 英语 |
| E5 Base V2 | 31 | 409 | 英语 |
| E5 Large | 45.9 | 1230 | 英语 |
| E5 Large V2 | 49.6 | 1220 | 英语 |
| GTE Base | 30.3 | 375 | 英语 |
| GTE Small | 28.5 | 129 | 英语 |
| GTE Large | 46.6 | 1320 | 英语 |
在大型数据集上获取最近邻
当向SageMaker调用端点发出请求时，负载限制在约5MB左右，并且请求超时时间设置为1分钟。如果语料库大小超过这些限制，您可以使用SageMaker训练任务，该任务为您的大型数据集生成嵌入，并将其与模型一起存储在SageMaker端点内。因此，它们不需要作为调用负载的一部分传递。通过使用SentenceTransformer及其实用函数来执行查找最近邻的过程。最近邻基于输入句子嵌入和训练任务期间预计算的句子嵌入之间的余弦相似度。
在以下示例中，我们获取并准备Amazon_SageMaker_FAQs数据集，以便在找到输入问题的最近邻时使用：
!aws s3 cp s3://jumpstart-cache-prod-us-west-2/training-datasets/Amazon_SageMaker_FAQs/Amazon_SageMaker_FAQs.csv Amazon_SageMaker_FAQs.csvimport pandas as pddata = pd.read_csv("Amazon_SageMaker_FAQs.csv", names=["Questions", "Answers"])data["id"] = data.indexdata_req = data[["id", "Answers"]]data_req.to_csv("data.csv", index=False, header=False)output_bucket = sess.default_bucket()output_prefix = "jumpstart-example-ss-training"s3_output_location = f"s3://{output_bucket}/{output_prefix}/output"training_dataset_s3_path = f"s3://{output_bucket}/{output_prefix}/data/data.csv"!aws s3 cp data.csv {training_dataset_s3_path}对于特定算法的训练超参数，可以获取或覆盖SageMaker SDK：
from sagemaker import hyperparametershyperparameters = hyperparameters.retrieve_default(model_id=model_id, model_version = "*")hyperparameters["batch_size"] = "64"print(hyperparameters)>>> {'max_seq_length': 'None', 'batch_size': '64', 'store_text_with_embedding': 'True'}SageMaker训练包括两个步骤：创建估计器对象和启动训练任务。输出是一个带有您用作训练数据的大型数据集嵌入的预打包模型，可以部署用于获取任何输入句子的最近邻。请参阅以下代码：
from sagemaker.jumpstart.estimator import JumpStartEstimatorestimator = JumpStartEstimator( model_id=model_id, hyperparameters=hyperparameters, output_path=s3_output_location)estimator.fit( {"training": f"s3://{output_bucket}/{output_prefix}/data"})predictor = estimator.deploy()将文本转换为嵌入的查询语法与之前相同。然而，获取最近邻的代码可以简化如下：
payload_nearest_neighbour = { "queries": ["Is R supported with Amazon SageMaker?"], "top_k": 1, "mode": "nn_train_data",}response = predictor.predict(payload_nearest_neighbour)>>> [[{'id': '9', 'score': 0.9240573048591614}]]data["Answers"].iloc[int(response[0][0]["id"])]>>> "Yes, R is supported with Amazon SageMaker. You can use R within SageMaker notebook instances, which include a preinstalled R kernel and the reticulate library. Reticulate offers an R interface for the Amazon SageMaker Python SDK, enabling ML practitioners to build, train, tune, and deploy R models."我们还可以使用Amazon_SageMaker_FAQs数据集中的问题查询端点，并比较返回的正确对应答案的数量。在以下示例中，我们测量前3个准确率，因为可能存在类似的问题-答案对。这意味着如果正确答案作为前3个返回之一，它将被视为正确查询。
total_correct_answers = 0for i in range(len(data)): question = data["Questions"].iloc[i] payload_nearest_neighbor = { "queries": [question], "top_k": 3, "mode": "nn_train_data", } response = predictor.predict(payload_nearest_neighbor) response_ids = [int(res["id"]) for res in response[0]] if i in response_ids: total_correct_answers += 1 else: pred_answer = [data["Answers"].iloc[response_id] for response_id in response_ids]print(total_correct_answers*100/len(data))>>>81.16883116883118批量转换以在大数据集上获取嵌入
对于具有超出单个终端节点实例内存的大量历史文档的企业和组织，您可以使用SageMaker批量转换来节省成本。当您启动批量转换作业时，SageMaker会启动必要的计算资源来处理数据。作业期间，SageMaker会自动配置和管理计算资源。批量转换作业完成后，这些资源会自动清理，从而最小化成本。通过将大型数据集分割成较小的块并使用更多实例，您可以扩展计算，以实现更快的推断速度，并具有类似的成本，而无需管理基础架构。批量转换的最大有效载荷为100 MB，超时为1小时。
我们批量转换作业的输入格式是一个JSONL文件，其中每个条目是一个JSON行，其中包括id和text_inputs。请参考以下代码：
test_data_file_name = "test.jsonl"test_data = []for i in range(len(data)): answer = data.loc[i, "Answers"] payload = {"id": i, "text_inputs": answer} test_data.append(payload)with open(test_data_file_name, "w") as outfile: for entry in test_data: outfile.write(f"{json.dumps(entry)}\n")s3 = boto3.client("s3")s3.upload_file(test_data_file_name, output_bucket, f"{output_prefix}/batch_input/test.jsonl")当数据准备完毕并存储在亚马逊简单存储服务（Amazon S3）中后，您可以从SageMaker JumpStart模型创建批量转换对象，从而触发转换作业：
s3_input_data_path = f"s3://{output_bucket}/{output_prefix}/batch_input/"s3_output_data_path = f"s3://{output_bucket}/{output_prefix}/batch_output/"batch_transformer = text_embedding_model.transformer( instance_count=1, instance_type="ml.p3.2xlarge", output_path=s3_output_data_path, assemble_with="Line", accept="text/csv", max_payload=1,)batch_transformer.transform( s3_input_data_path, content_type="application/jsonlines", split_type="Line")batch_transformer.wait()批量转换作业完成后，您可以从亚马逊S3下载结果：
s3 = boto3.client("s3")s3.download_file( output_bucket, output_prefix + "/batch_output/" + "test.jsonl.out", "predict.jsonl")with open("predict.jsonl", "r") as json_file: json_list = list(json_file)结论
SageMaker JumpStart提供了一种简单的方式来使用最先进的大型语言基础模型进行文本嵌入和语义搜索。通过用户界面或仅几行代码，您可以部署高精度的文本嵌入模型，并在大规模数据集上快速找到语义匹配，同时具有成本效益。SageMaker JumpStart通过提供对MTEB领导者榜上基准测试的先进模型的即时访问，消除了实施语义搜索的障碍。企业和开发人员可以更快地构建智能搜索和推荐系统。
本文演示了如何找到语义上相似的问题和答案，并可应用于RAG用例、推荐和个性化、多语言翻译等。随着语言模型的不断进步以及SageMaker JumpStart的简易性，越来越多的组织可以将生成性AI能力融入其产品中。作为下一步，您可以尝试在自己的数据集上使用SageMaker JumpStart的文本嵌入模型，以测试和评估您的RAG用例的结果。