Link Search Menu Expand Document Documentation Menu
文档

在 Amazon SageMaker 中使用模型进行语义搜索

本教程展示了如何使用 Amazon SageMaker 中的嵌入模型在 Amazon OpenSearch Service 中实现语义搜索。有关更多信息，请参阅语义搜索。

如果使用 Python，您可以创建 Amazon SageMaker 连接器并使用 opensearch-py-ml 客户端 CLI 测试模型。该 CLI 自动化了许多配置步骤，使设置更快并减少了出错的可能性。有关使用 CLI 的更多信息，请参阅 CLI 文档。

如果使用自托管 OpenSearch 而非 Amazon OpenSearch Service，请使用蓝图在 Amazon SageMaker 中创建模型的连接器。有关创建连接器的更多信息，请参阅连接器。

本教程不涵盖如何将模型部署到 Amazon SageMaker。有关部署的更多信息，请参阅实时推理。

在 Amazon OpenSearch Service 中设置嵌入模型的最简单方法是使用 AWS CloudFormation。另外，您可以使用 AIConnectorHelper 笔记本设置嵌入模型。

将以 your_ 为前缀的占位符替换为您自己的值。

模型输入和输出要求

确保您的 Amazon SageMaker 模型输入遵循默认预处理函数所需的格式。

模型输入必须是字符串数组

["hello world", "how are you"]

此外，确保模型输出遵循默认后处理函数所需的格式。模型输出必须是数组的数组，其中每个内部数组对应于一个输入字符串的嵌入

[
  [
    -0.048237994,
    -0.07612697,
    ...
  ],
  [
    0.32621247,
    0.02328475,
    ...
  ]
]

如果您的模型输入/输出与所需的默认值不同，您可以使用 Painless 脚本构建自己的预处理/后处理函数。

示例：Amazon Bedrock Titan 嵌入模型

例如，Amazon Bedrock Titan 嵌入模型（蓝图）输入如下

{ "inputText": "your_input_text" }

OpenSearch 需要以下输入格式

{ "text_docs": [ "your_input_text1", "your_input_text2"] }

要将 text_docs 转换为 inputText，您必须定义以下预处理函数

"pre_process_function": """
    StringBuilder builder = new StringBuilder();
    builder.append("\"");
    String first = params.text_docs[0];// Get the first doc, ml-commons will iterate all docs
    builder.append(first);
    builder.append("\"");
    def parameters = "{" +"\"inputText\":" + builder + "}"; // This is the Bedrock Titan embedding model input
    return  "{" +"\"parameters\":" + parameters + "}";"""

默认的 Amazon Bedrock Titan 嵌入模型输出具有以下格式

{
  "embedding": <float_array>
}

然而，OpenSearch 需要以下格式

{
  "name": "sentence_embedding",
  "data_type": "FLOAT32",
  "shape": [ <embedding_size> ],
  "data": <float_array>
}

要将 Amazon Bedrock Titan 嵌入模型输出转换为 OpenSearch 所需的格式，您必须定义以下后处理函数

"post_process_function": """
      def name = "sentence_embedding";
      def dataType = "FLOAT32";
      if (params.embedding == null || params.embedding.length == 0) {
        return params.message;
      }
      def shape = [params.embedding.length];
      def json = "{" +
                 "\"name\":\"" + name + "\"," +
                 "\"data_type\":\"" + dataType + "\"," +
                 "\"shape\":" + shape + "," +
                 "\"data\":" + params.embedding +
                 "}";
      return json;
    """

先决条件：创建 OpenSearch 集群

前往 Amazon OpenSearch Service 控制台并创建一个 OpenSearch 域。

记下域 Amazon Resource Name (ARN)；您将在后续步骤中使用它。

步骤 1：创建 IAM 角色以调用 Amazon SageMaker 中的模型

要在 Amazon SageMaker 中调用模型，您必须创建具有相应权限的 AWS 身份和访问管理 (IAM) 角色。连接器将使用此角色来调用模型。

转到 IAM 控制台，创建一个名为 my_invoke_sagemaker_model_role 的新 IAM 角色，并添加以下信任策略和权限

  • 自定义信任策略
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "es.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}

  • 权限
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "sagemaker:InvokeEndpoint"
            ],
            "Resource": [
                "your_sagemaker_model_inference_endpoint_arn"
            ]
        }
    ]
}

记下角色 ARN；您将在后续步骤中使用它。

步骤 2：在 Amazon OpenSearch Service 中配置 IAM 角色

按照以下步骤在 Amazon OpenSearch Service 中配置 IAM 角色。

步骤 2.1：创建用于签署连接器请求的 IAM 角色

专门为签署您的创建连接器 API 请求生成一个新的 IAM 角色。

创建一个名为 my_create_sagemaker_connector_role 的 IAM 角色，并添加以下信任策略和权限

  • 自定义信任策略
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "AWS": "your_iam_user_arn"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}

您将在步骤 3.1 中使用 your_iam_user_arn IAM 用户来承担此角色。

  • 权限
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": "iam:PassRole",
            "Resource": "your_iam_role_arn_created_in_step1"
        },
        {
            "Effect": "Allow",
            "Action": "es:ESHttpPost",
            "Resource": "your_opensearch_domain_arn"
        }
    ]
}

记下此角色 ARN；您将在后续步骤中使用它。

步骤 2.2：映射后端角色

按照以下步骤映射后端角色

  1. 登录 OpenSearch Dashboards，并在顶部菜单中选择 Security（安全）。
  2. 选择 Roles（角色），然后选择 ml_full_access 角色。
  3. 在 ml_full_access 角色详情页面，选择 Mapped users（已映射用户），然后选择 Manage mapping（管理映射）。
  4. 在“后端角色”字段中输入在步骤 2.1 中创建的 IAM 角色 ARN，如下图所示。映射后端角色
  5. 选择 Map（映射）。

IAM 角色现已成功在您的 OpenSearch 集群中配置。

步骤 3：创建连接器

按照以下步骤为模型创建连接器。有关创建连接器的更多信息，请参阅连接器。

步骤 3.1：获取临时凭证

使用步骤 2.1 中指定的 IAM 用户的凭证来承担角色

aws sts assume-role --role-arn your_iam_role_arn_created_in_step2.1 --role-session-name your_session_name

从响应中复制临时凭证，并将其配置在 ~/.aws/credentials 中

[default]
AWS_ACCESS_KEY_ID=your_access_key_of_role_created_in_step2.1
AWS_SECRET_ACCESS_KEY=your_secret_key_of_role_created_in_step2.1
AWS_SESSION_TOKEN=your_session_token_of_role_created_in_step2.1

步骤 3.2：创建连接器

使用在 ~/.aws/credentials 中配置的临时凭证运行以下 Python 代码

import boto3
import requests 
from requests_aws4auth import AWS4Auth

host = 'your_amazon_opensearch_domain_endpoint'
region = 'your_amazon_opensearch_domain_region'
service = 'es'

credentials = boto3.Session().get_credentials()
awsauth = AWS4Auth(credentials.access_key, credentials.secret_key, region, service, session_token=credentials.token)


path = '/_plugins/_ml/connectors/_create'
url = host + path

payload = {
  "name": "Sagemaker embedding model connector",
  "description": "Connector for my Sagemaker embedding model",
  "version": "1.0",
  "protocol": "aws_sigv4",
  "credential": {
    "roleArn": "your_iam_role_arn_created_in_step1"
  },
  "parameters": {
    "region": "your_sagemaker_model_region",
    "service_name": "sagemaker"
  },
  "actions": [
    {
      "action_type": "predict",
      "method": "POST",
      "headers": {
        "content-type": "application/json"
      },
      "url": "your_sagemaker_model_inference_endpoint",
      "request_body": "${parameters.input}",
      "pre_process_function": "connector.pre_process.default.embedding",
      "post_process_function": "connector.post_process.default.embedding"
    }
  ]
}

headers = {"Content-Type": "application/json"}

r = requests.post(url, auth=awsauth, json=payload, headers=headers)
print(r.status_code)
print(r.text)

脚本将输出连接器 ID

{"connector_id":"tZ09Qo0BWbTmLN9FM44V"}

记下连接器 ID；您将在下一步中使用它。

步骤 4：创建并测试模型

登录 OpenSearch Dashboards，打开 DevTools 控制台，并运行以下请求来创建和测试模型。

  1. 创建模型组

     POST /_plugins/_ml/model_groups/_register
 {
     "name": "Sagemaker_embedding_model",
     "description": "Test model group for Sagemaker embedding model"
 }

    响应包含模型组 ID

     {
   "model_group_id": "MhU3Qo0BTaDH9c7tKBfR",
   "status": "CREATED"
 }

  2. 注册模型

     POST /_plugins/_ml/models/_register
 {
   "name": "Sagemaker embedding model",
   "function_name": "remote",
   "description": "test embedding model",
   "model_group_id": "MhU3Qo0BTaDH9c7tKBfR",
   "connector_id": "tZ09Qo0BWbTmLN9FM44V"
 }

    响应包含模型 ID

     {
   "task_id": "NhU9Qo0BTaDH9c7t0xft",
   "status": "CREATED",
   "model_id": "NxU9Qo0BTaDH9c7t1Bca"
 }

  3. 部署模型

     POST /_plugins/_ml/models/NxU9Qo0BTaDH9c7t1Bca/_deploy

    响应包含部署操作的任务 ID

     {
   "task_id": "MxU4Qo0BTaDH9c7tJxde",
   "task_type": "DEPLOY_MODEL",
   "status": "COMPLETED"
 }

  4. 测试模型

     POST /_plugins/_ml/models/NxU9Qo0BTaDH9c7t1Bca/_predict
 {
   "parameters": {
     "input": ["hello world", "how are you"]
   }
 }

    响应包含模型生成的嵌入

     {
   "inference_results": [
     {
       "output": [
         {
           "name": "sentence_embedding",
           "data_type": "FLOAT32",
           "shape": [
             384
           ],
           "data": [
             -0.034477264,
             0.031023195,
             0.0067349933,
             ...]
         },
         {
           "name": "sentence_embedding",
           "data_type": "FLOAT32",
           "shape": [
             384
           ],
           "data": [
             -0.031369038,
             0.037830487,
             0.07630822,
             ...]
         }
       ],
       "status_code": 200
     }
   ]
 }

按照以下步骤配置语义搜索。

步骤 5.1：创建摄取管道

首先，创建一个摄取管道，该管道使用 Amazon SageMaker 中的模型从输入文本创建嵌入

PUT /_ingest/pipeline/my_sagemaker_embedding_pipeline
{
    "description": "text embedding pipeline",
    "processors": [
        {
            "text_embedding": {
                "model_id": "your_sagemaker_embedding_model_id_created_in_step4",
                "field_map": {
                    "text": "text_knn"
                }
            }
        }
    ]
}

步骤 5.2：创建向量索引

接下来，创建一个用于存储输入文本和生成嵌入的向量索引

PUT my_index
{
  "settings": {
    "index": {
      "knn.space_type": "cosinesimil",
      "default_pipeline": "my_sagemaker_embedding_pipeline",
      "knn": "true"
    }
  },
  "mappings": {
    "properties": {
      "text_knn": {
        "type": "knn_vector",
        "dimension": your_sagemake_model_embedding_dimension
      }
    }
  }
}

步骤 5.3：摄取数据

将示例文档摄取到索引中

POST /my_index/_doc/1000001
{
    "text": "hello world."
}

步骤 5.4：搜索索引

运行向量搜索以从向量索引中检索文档

POST /my_index/_search
{
  "query": {
    "neural": {
      "text_knn": {
        "query_text": "hello",
        "model_id": "your_embedding_model_id_created_in_step4",
        "k": 100
      }
    }
  },
  "size": "1",
  "_source": ["text"]
}

剩余 350 字符

有问题？。

想做贡献？或。