mirrors/aws--amazon-sagemaker-examples
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
43ec1feedeba352cf701fc069ac1b23d45ddeee7
aws--amazon-sagemaker-examples/ generative_ai/sm-text_embedding_custom_processing.ipynb
T
wehub-resource-sync 68ca4df005 chore: import upstream snapshot with attribution
2026-07-13 13:40:10 +08:00

554 lines
24 KiB
Plaintext
Raw Blame History

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Pre/Post-Processing for Hugging Face (HF) Text Embeddings Inference (TEI)\n",
"\n",
"---\n",
"\n",
"This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.\n",
"\n",
"![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-2/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"---"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Overview\n",
"\n",
"In this demo notebook, we demonstrate how to implement pre/post-processing logic for [HF TEI](https://huggingface.co/docs/text-embeddings-inference/en/index) use cases. While it provides improved performance and convenient SageMaker integrations, the [SageMaker TEI image](https://docs.aws.amazon.com/sagemaker/latest/dg/pre-built-containers-support-policy.html#pre-built-containers-support-policy-dlc) does not support custom pre/post-processing logic out-of-the-box. For customers looking to customize their logic, we show how to leverage the [SageMaker Transformers Image](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-inference-containers) to load and invoke a HF TEI model with a custom InferenceSpec so that customers can still pre/post-process their results from their SageMaker endpoint(s).\n",
"\n",
"## Prerequisites\n",
"\n",
"This notebook was tested in region `us-west-2` with kernel `conda_python3 (3.10.15 | packaged by conda-forge | (main, Sep 20 2024, 16:37:05) [GCC 13.3.0])` and uses the following (versioned) resources:\n",
"\n",
"| Resource | Value |\n",
"| :-------- | :----- |\n",
"| TEI Model | jinaai/jina-embeddings-v2-small-en |\n",
"| TEI Image (as a control for performance testing) | 246618743249.dkr.ecr.us-west-2.amazonaws.com/tei-cpu:2.0.1-tei1.4.0-cpu-py310-ubuntu22.04 |\n",
"| Transformers Image | 763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference:2.1.0-transformers4.37.0-cpu-py310-ubuntu22.04 |\n",
"| Instance Type | ml.m5.4xlarge |\n",
"\n",
"- The accounts must match the account from which the region is pulled from. This varies based on resource. For example, `246618743249` is the account providing the TEI Image for the `us-west-2` region. Please contact the resource providers for more info about target resources. \n",
"- The images used are based on `Python3.10`. The **image versions must match the PySDK version** for reasons related to [pickling](https://docs.python.org/3/library/pickle.html) performed by the PySDK.\n",
"- The images used are for target instance type `ml.m5.4xlarge` i.e. `CPU` . Please be sure to use the appropriate TEI and Transformer images for target hardware.\n",
"- The TEI model used must fit on the target instance type; otherwise, the inferences will fail (even if the endpoint still reports \"online\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Demonstration"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Dependencies\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!pip install 'sagemaker<3.0' numpy transformers datasets --upgrade"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Role\n",
"\n",
"To host on Amazon SageMaker, we need to set up and authenticate the use of AWS services. Here, we use the execution role associated with the current notebook as the AWS account role with SageMaker access."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from sagemaker import get_execution_role\n",
"\n",
"ROLE = get_execution_role()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Constants for our Intended Resources"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"HF_TEI_MODEL = \"jinaai/jina-embeddings-v2-small-en\"\n",
"INSTANCE_TYPE = \"ml.m5.4xlarge\"\n",
"TEI_IMAGE = (\n",
" \"246618743249.dkr.ecr.us-west-2.amazonaws.com/tei-cpu:2.0.1-tei1.4.0-cpu-py310-ubuntu22.04\"\n",
")\n",
"TRANSFORMERS_IMAGE = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference:2.1.0-transformers4.37.0-cpu-py310-ubuntu22.04\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Standard Helper Functions to Deploy Endpoints\n",
"\n",
"This notebook will deploy two different endpoints:\n",
"1. Endpoint using standard TEI image (as a control to compare against)\n",
"2. Endpoint using Transformers image\n",
"\n",
"To simplify this, we define a few helper functions "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"##### **deploy()**\n",
"\n",
"To deploy a SageMaker endpoint, we need to first build a SageMaker Model. Since we will later use [InferenceSpec](https://sagemaker.readthedocs.io/en/v2.208.0/api/inference/model_builder.html#sagemaker.serve.spec.inference_spec.InferenceSpec), we should use the SageMaker Python SDK constrcut `ModelBuilder` to build a Model (rather than directly use the `Model` construct, which does not support defining a custom InferenceSpec)\n",
"\n",
"When defining a new model with `ModelBuilder`, it requires a `SchemaBuilder` to understand the input and output data types. This is essential for properly (de)serializing the data. `SchemaBuilder` is able to infer the kind of (de)serializers to use by providing example inputs. For the TEI case, the input and output look like:\n",
"- Input: a (stringified) JSON object with the sentences to get embeddings for\n",
"- Output: a (nested) lists, which represents the embeddings. \n",
"\n",
"The exact shapes and names used within these samples are not critical since these are just used to infer the (de)serializers and not reshape or fit actual inputs/outputs into.\n",
"\n",
"Below, we use some Python shorthand involving `**` to define default arguments and allow for overrides if they are passed."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"from sagemaker.serve.builder.schema_builder import SchemaBuilder\n",
"import sagemaker\n",
"\n",
"sagemaker.model.FrameworkModel\n",
"\n",
"from sagemaker.serve.builder.model_builder import ModelBuilder\n",
"\n",
"\n",
"def deploy(model_builder_kwargs={}):\n",
" model = ModelBuilder(\n",
" **{\n",
" **dict(\n",
" role_arn=ROLE,\n",
" schema_builder=SchemaBuilder(\n",
" json.dumps({\"inputs\": [\"hello\", \"world\"]}), [[1, 2, 3], [4, 5, 6]]\n",
" ),\n",
" env_vars={\n",
" \"TS_DISABLE_TOKEN_AUTHORIZATION\": \"true\" # See https://github.com/pytorch/serve/blob/master/docs/README.md\n",
" },\n",
" ),\n",
" **model_builder_kwargs, # allow for overridding these arguments if desired\n",
" }\n",
" ).build()\n",
" endpoint = model.deploy(\n",
" initial_instance_count=1,\n",
" instance_type=INSTANCE_TYPE,\n",
" )\n",
" return (model, endpoint)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"##### **clean()**\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def clean(model, endpoint):\n",
" try:\n",
" endpoint.delete_endpoint()\n",
" except Exception as e:\n",
" print(e)\n",
" pass\n",
"\n",
" try:\n",
" model.delete_model()\n",
" except Exception as e:\n",
" print(e)\n",
" pass"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Endpoint-Specific Usage of Helpers"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"##### TEI Image (as-is, used as control to compare against)\n",
"\n",
"The SageMaker Python SDK and SageMaker TEI image make it convenient to simply specify the intended `HF_TEI_MODEL` with the SageMaker `TEI_IMAGE`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import time\n",
"\n",
"\n",
"def deploy_tei():\n",
" return deploy(\n",
" model_builder_kwargs=dict(\n",
" name=f\"tei-{int(time.time())}\",\n",
" image_uri=TEI_IMAGE,\n",
" model=HF_TEI_MODEL,\n",
" )\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"##### Transformers Image\n",
"\n",
"Since we will be defining custom logic using the Transformers Image, we need to provide our own [`InferenceSpec`](https://sagemaker.readthedocs.io/en/v2.208.0/api/inference/model_builder.html#sagemaker.serve.spec.inference_spec.InferenceSpec). This construct requires the following implementations:\n",
"- `load()` : how to load your intended model\n",
"- `invoke()` : how to call your loaded model\n",
"\n",
"To do this, we create our own subclass of `InferenceSpec` called `CustomerInferenceSpec`. Then, we provide our implementations of each by overriding the expected methods with the same signature in the base class. \n",
"\n",
"##### **load()**\n",
"\n",
"[We can load our TEI model `jinaai/jina-embeddings-v2-small-en` using the transformers library as documented by the HF model owners](https://huggingface.co/jinaai/jina-embeddings-v2-base-en#usage).\n",
"\n",
"##### **invoke()**\n",
"\n",
"[We can invoke our loaded TEI model `jinaai/jina-embeddings-v2-small-en` with a single `encode()` call as documented by the HF model owners](https://huggingface.co/jinaai/jina-embeddings-v2-base-en#usage). \n",
"\n",
"Note: The inputs to `encode()` must be extracted from the input data sent to the endpoint. This is done with a simple JSON load and reference to `inputs`, which is what the endpoint was invoked with.\n",
"\n",
"##### **Dependencies**\n",
"The SageMaker Python SDK can auto-detect dependencies at the class level of an InferenceSpec. It will install them in the intended serving container while provisioning our endpoint. \n",
"\n",
"For this demo, however, we would like a specific version of the HF `transformers` library that complies with the version of Python and other dependencies already in our target image. Therefore, we opt to use the `dependencies` parameter explicitly to specify a particular version of `transformers` to ensure compatibility."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from sagemaker.serve.spec.inference_spec import InferenceSpec\n",
"\n",
"\n",
"def deploy_transformer():\n",
" class CustomerInferenceSpec(InferenceSpec):\n",
" def load(self, model_dir):\n",
" from transformers import AutoModel\n",
"\n",
" return AutoModel.from_pretrained(HF_TEI_MODEL, trust_remote_code=True)\n",
"\n",
" def invoke(self, x, model):\n",
" return model.encode(x)\n",
"\n",
" def preprocess(self, input_data):\n",
" return json.loads(input_data)[\"inputs\"]\n",
"\n",
" def postprocess(self, predictions):\n",
" assert predictions is not None\n",
" return predictions\n",
"\n",
" return deploy(\n",
" dict(\n",
" name=f\"transformers-{int(time.time())}\",\n",
" image_uri=TRANSFORMERS_IMAGE,\n",
" inference_spec=CustomerInferenceSpec(),\n",
" dependencies={\n",
" \"custom\": [\n",
" \"transformers==4.38.0\" # so we don't override the DLC dependency versions\n",
" ],\n",
" },\n",
" )\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Running the Demo\n",
"\n",
"In order to run the demo, we will need to\n",
"1. deploy our (models and) endpoints\n",
"2. invoke against our (models and) endpoints\n",
"3. clean up our (models and) endpoints \n",
"\n",
"We can define a just few more helpers for better reuse and customizations if needed in the future"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Helper Functions to Run This Specific Demo"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"##### **invoke_many()**\n",
"\n",
"This is a synchronous function that will be sent to a thread using asyncio. \n",
"\n",
"It a simple invoker used to invoke against the endpoint given a list of samples to call with.\n",
"\n",
"We intentionally do not use batching here for a basic performance test later in this notebook.\n",
"\n",
"Note: This is tightly-coupled with our sample [dataset](https://huggingface.co/datasets/sentence-transformers/stsb). Specfically, the way we run invocations with `predict` against our endpoints is based on how the dataset is structured. Each sample has two sentences. See dataset page for more info.\n",
"\n",
"Note: The `initial_args` are needed only for the endpoint with the TEI Image."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def invoke_many(endpoint, samples): # intentionally not batching\n",
" results = []\n",
" for sample in samples:\n",
" start = time.perf_counter()\n",
" res = endpoint.predict(\n",
" json.dumps({\"inputs\": [sample[\"sentence1\"], sample[\"sentence2\"]]}),\n",
" initial_args=(\n",
" {\"ContentType\": \"application/json\"} if \"tei\" in endpoint.endpoint_name else None\n",
" ),\n",
" )\n",
" end = time.perf_counter()\n",
" results.append(\n",
" {\n",
" \"latency\": end - start,\n",
" \"embeddings\": res,\n",
" \"endpoint\": endpoint.endpoint_name,\n",
" \"sample\": sample,\n",
" }\n",
" )\n",
" return results"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"##### **basic_performance_test()**\n",
"\n",
"This is an async function that will leverage asyncio to create, invoke, and clean up endpoints.\n",
"\n",
"We use asyncio here for better concurrency since the endpoint (de)provisioning can take a long, indeterminate amount of time.\n",
"\n",
"In this demo, we will \n",
"1. Create two endpoints:\n",
" 1. One with TEI Image \n",
" 2. One with Transformers Image\n",
"2. Load sample sentences from the [`sentence-transformers/stsb` dataset](https://huggingface.co/datasets/sentence-transformers/stsb).\n",
"3. Invoke each endpoint with the samples.\n",
"4. Analyze the latencies of the sample invocations.\n",
"5. Clean up all of the resources."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from datasets import load_dataset\n",
"import numpy as np\n",
"import asyncio\n",
"import pprint\n",
"\n",
"\n",
"async def basic_performance_test():\n",
" deployments = []\n",
" try:\n",
" ######################################################\n",
" # Deploy\n",
" ######################################################\n",
" print(\"Deploying endpoints...\")\n",
" deployments = await asyncio.gather(\n",
" *[asyncio.to_thread(deploy_tei), asyncio.to_thread(deploy_transformer)]\n",
" )\n",
"\n",
" ######################################################\n",
" # Invoke\n",
" ######################################################\n",
" print(\"Invoking endpoints...\")\n",
" samples = load_dataset(\"sentence-transformers/stsb\", streaming=True, split=\"test\").take(500)\n",
"\n",
" results = await asyncio.gather(\n",
" *[asyncio.to_thread(invoke_many, endpoint, samples) for _, endpoint in deployments]\n",
" )\n",
"\n",
" ######################################################\n",
" # Analyze\n",
" ######################################################\n",
" print(\"Analyzing invocations...\")\n",
" for invocations in results:\n",
" latencies = np.array([invocation[\"latency\"] for invocation in invocations])\n",
" pprint.pp(\n",
" {\n",
" \"shape\": np.shape(latencies),\n",
" \"tm99\": np.mean(\n",
" np.sort(latencies)[: (len(latencies) - int(len(latencies) * 0.01))]\n",
" ),\n",
" \"p90\": np.percentile(latencies, 90),\n",
" \"avg\": np.mean(latencies),\n",
" \"max\": np.max(latencies),\n",
" \"min\": np.min(latencies),\n",
" \"endpoint\": invocations[0][\"endpoint\"],\n",
" },\n",
" indent=4,\n",
" )\n",
" finally:\n",
" ######################################################\n",
" # Clean\n",
" ######################################################\n",
" print(\"Cleaning resources...\")\n",
" errors = await asyncio.gather(\n",
" *[asyncio.to_thread(clean, model, endpoint) for model, endpoint in deployments],\n",
" return_exceptions=True\n",
" )\n",
" for error in errors:\n",
" if error:\n",
" print(error)\n",
"\n",
" print(\"Complete!\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Run the Demo\n",
"\n",
"Note: To see live output from CloudWatch, please ensure your role has `logs:FilterLogEvents` permissions for the created endpoints\n",
"\n",
"Note: The logging output from the endpoints will be in <span style=\"background-color:#ffdddd\"> **red** because it is logged to stderr </span> by default"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"await basic_performance_test()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Results\n",
"\n",
"This demo demonstrates:\n",
"1. How to achieve custom pre/post-processing using the Transformers image alternative and customizing InferenceSpec\n",
"2. Reduced, yet comparable performance versus the SageMaker TEI image while using the Transformers image. If pre/post-processing is needed for the TEI endpoint, then this is may be a viable alternative for production endpoints"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Notebook CI Test Results\n",
"\n",
"This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
"\n",
"\n",
"![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-1/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-2/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-1/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ca-central-1/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/sa-east-1/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-1/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-2/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-3/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-central-1/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-north-1/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-1/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-2/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-1/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-2/generative_ai|sm-text_embedding_custom_processing.ipynb)\n",
"\n",
"![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-south-1/generative_ai|sm-text_embedding_custom_processing.ipynb)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "conda_python3",
"language": "python",
"name": "conda_python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.15"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Reference in New Issue View Git Blame Copy Permalink