Posts in category: Amazon AWS

Today, we’re excited to announce that the Gemma model is now available for customers using Amazon SageMaker JumpStart. Gemma is a family of language models based on Google’s Gemini models, trained on up to 6 trillion tokens of text. The Gemma family consists of two sizes: a 7 billion parameter model and a 2 billion parameter model. Now, you can use Gemma 2B and Gemma 7B pretrained and instruction-tuned models within SageMaker JumpStart. JumpStart is the machine learning (ML) hub of SageMaker that provides access to foundation models in addition to built-in algorithms and end-to-end solution templates to help you quickly get started with ML.

In this post, we walk through how to deploy the Gemma model and fine tune it for your use cases in SageMaker JumpStart. The complete notebook is available on GitHub.

Gemma model

Gemma is a family of lightweight, state-of-the-art models built from the same research and technology used to create the Gemini models. Developed by Google DeepMind and other teams across Google, Gemma is inspired by Gemini. Gemma exhibits strong generalist capabilities in text domains and state-of-the-art understanding and reasoning skills at scale. It achieves better performance compared to other publicly available models of similar or larger scales across different domains, including question answering, commonsense reasoning, mathematics and science, and coding. Gemma released the model weights to support developer innovation using Gemma models. Gemma was launched with a new Responsible Generative AI Toolkit that provides guidance and essential tools for creating safer AI applications with Gemma.

Foundation models in SageMaker

JumpStart provides access to a range of models from popular model hubs including Hugging Face, PyTorch Hub, and TensorFlow Hub, which you can use within your ML development workflow in SageMaker. Recent advances in ML have given rise to a new class of models known as foundation models, which are typically trained on billions of parameters and are adaptable to a wide category of use cases, such as text summarization, generating digital art, and language translation. Because these models are expensive to train, customers want to use existing pre-trained foundation models and fine-tune them as needed, rather than train these models themselves. SageMaker provides a curated list of models that you can choose from on the SageMaker console.

You can now find foundation models from different model providers within JumpStart, enabling you to get started with foundation models quickly. You can find foundation models based on different tasks or model providers, and review model characteristics and usage terms. You can also try these models using a test UI widget. When you want to use a foundation model at scale, you can do so without leaving SageMaker by using pre-built notebooks from model providers. Because the models are hosted and deployed on AWS, your data, whether used for evaluating the model or using it at scale, is never shared with third parties.

Let’s explore how you can use the Llama Guard model in JumpStart.

Explore the Gemma model in Jumpstart

You can access Gemma foundation models through SageMaker JumpStart in the SageMaker Studio UI and the SageMaker Python SDK. In this section, we go over how to discover the models in SageMaker Studio.

SageMaker Studio is an integrated development environment (IDE) that provides a single web-based visual interface where you can access purpose-built tools to perform all ML development steps, from preparing data to building, training, and deploying your ML models. For more details on how to get started and set up SageMaker Studio, see Amazon SageMaker Studio.

In the AWS Management Console for SageMaker Studio, go to SageMaker JumpStart under Prebuilt and automated solutions. Jumpstart contains pre-trained models, notebooks, and prebuilt solutions.

On the SageMaker JumpStart landing page, you can find the Gemma model by searching for Gemma.

You can then select from a variety of Gemma model variants, including Gemma 2B, Gemma 7B, Gemma 2B instruct, and Gemma 7B instruct.

Choose the model card to view details about the model such as the license, data used to train, and how to use the model. You will also find a Deploy button, which takes you to a landing page where you can test inference with an example payload.

Deploy Gemma with SageMaker Python  SDK

You can find the code showing the deployment of Gemma on JumpStart and an example of how to use the deployed model in this GitHub notebook.

Start by selecting the SageMaker Model Hub model ID and model version to use when deploying Gemma.

model_id, model_version = "huggingface-llm-gemma-7b-instruct", "*"

Choose a model ID from the following table, which details the default configuration options for the JumpStart deployment. Because of the large vocabulary size of 256 thousand tokens, Gemma 7B can only fit on a single A10G GPU when supporting a 1 thousand context length. For this reason, JumpStart uses a larger default instance for Gemma 7B.

You can now deploy the model using SageMaker JumpStart. The following code uses the default instance ml.g5.12xlarge for the inference endpoint You can deploy the model on other instance types by passing instance_type in the JumpStartModel class. The deployment might take 5-10 minutes.

from sagemaker.jumpstart.model import JumpStartModel

model = JumpStartModel(model_id=model_id, model_version=model_version)
predictor= model.deploy(accept_eula=False)  # manually accept EULA here!

For successful deployment, you must manually change the accept_eula argument in the model’s deploy method to True. This model is deployed using the text-generation-inference (TGI) deep learning container.

Invoke endpoint

You can programmatically retrieve example payloads from the JumpStartModel object. This will help you get started by observing pre-formatted instruction prompts that Gemma can ingest.

example_payloads = model.retrieve_all_examples()

for payload in example_payloads:
    response = predictor.predict(payload.body)
    print("Input:n", payload.body[payload.prompt_key])
    print("Output:n", response[0]["generated_text"].strip())
    print("n===============n")

Before we look at specific prompts, let’s consider the chat template for Gemma Instruct models.

<bos><start_of_turn>usern[USER_PROMPT]<end_of_turn>n<start_of_turn>model

Here, you place your prompt in the [USER_PROMPT] location. There’s no support for a system instruction; instead, you can prepend the desired instruction to the user prompt. Additionally, if you have a multi-turn conversation, then the model prompt can alternate between user and assistant as needed.

Now consider a few instruction example prompts. Here, you ask Gemma to write a Hello World program.

payload = {
    "inputs": "<bos><start_of_turn>usernWrite a hello world module in Python<end_of_turn>n<start_of_turn>model",
    "parameters": {'max_new_tokens': 256},
}
predictor.predict(payload)

The following is the expected output:

```python
# This is a Python module called hello.py

def hello():
  """Prints a hello message."""
  print("Hello, world!")

if __name__ == "__main__":
  hello()
```

**Usage:**

To use the hello.py module, you can import it and call the hello function like this:

```python
import hello

hello.hello()
```

**Output:**

```
Hello, world!
```

**Explanation:**

* The `hello.py` module defines a function called `hello()`.
* The `if __name__ == "__main__"` block checks if the script is being run directly or imported. If it is run directly, the `hello()` function is called.
* The `print("Hello, world!")` statement prints the hello message to the console.

**Note:**

* You can save the code above in a file named `hello.py` in your current directory.
* To run the script, you can use the following command:
```
python hello.py
```
* The output will be printed in the console:
```
Hello, world!
```

Next, invoke Gemma for the creative task of writing a poem.

payload = {
    "inputs": "<bos><start_of_turn>usernWrite me a poem about Machine Learning.<end_of_turn>n<start_of_turn>model",
    "parameters": {'max_new_tokens': 256},
}
predictor.predict(payload)

The following is the output:

**Machine Learning**

In the realm of data, a tale unfolds,
Where algorithms dance and stories are told.
With every click, a model takes flight,
Unveiling patterns, hidden in the night.

Data whispers secrets, a treasure to find,
Machine learning listens, with a curious mind.
It sifts through terabytes, with lightning speed,
Uncovering insights, a hidden creed.

From images to text, it learns to see,
To understand the complexities of reality.
It paints a canvas, with insights deep,
Unveiling patterns, secrets to keep.

In the field of medicine, it takes a stand,
Diagnosing diseases, guiding a hand.
It analyzes data, with unmatched grace,
Bringing hope to those in a troubled place.

But with great power comes great responsibility,
As algorithms can bias, with malicious intent.
It's a journey of learning, a constant fight,
To ensure fairness, day and night.

So let us harness the power of this tool,
To create a future, where insights bloom.
With wisdom and caution, we must guide,
The path of machine learning, side by side.

This looks pretty good!

Now, let’s look at latency and throughput performance benchmarking for model serving with the default JumpStart deployment configuration. Here, we show how model performance might differ for your typical endpoint workload. In the following tables, you can observe that small-sized queries (256 input words and 256 output tokens) are quite performant under a large number of concurrent users, reaching token throughput on the order of one thousand to two thousand tokens per second. However, as the number of input words approaches Gemma’s maximum supported context length of eight thousand tokens, the endpoint saturates its batching capacity—the number of concurrent requests allowed to be processed simultaneously—due to instance memory-bound constraints.

For more information on how to consider this information and adjust deployment configurations for your specific use case, see Benchmark and optimize endpoint deployment in Amazon SageMaker JumpStart.

Fine-tune Gemma using SageMaker Python SDK

Next, we show you how to fine-tune the Gemma 7B instruct model on a conversational-formatted dataset using QLoRA technique. As mentioned previously, due to the large vocabulary size of 256 thousand and the 8 thousand context length, JumpStart offers the following default configurations for QLoRA fine-tuning.

Let’s load and process the dataset in conversational format. The example dataset for this demonstration is OpenAssistant’s TOP-1 Conversation Threads.

from datasets import load_dataset

# Load the dataset
dataset = load_dataset("OpenAssistant/oasst_top1_2023-08-25")

The training data should be formulated in JSON lines (.jsonl) format, where each line is a dictionary representing a set of conversations. One example within the JSON lines file is shown below. For details on how to process the dataset, see the notebook in GitHub.

{'dialog': [
  {'content': 'what is the height of the empire state building',
   'role': 'user'},
  {'content': '381 meters, or 1,250 feet, is the height of the Empire State Building. If you also account for the antenna, it brings up the total height to 443 meters, or 1,454 feet',
   'role': 'assistant'},
  {'content': 'Some people need to pilot an aircraft above it and need to know.nSo what is the answer in feet?',
   'role': 'user'},
  {'content': '1454 feet', 'role': 'assistant'}]
}
import os
import boto3
from sagemaker.session import Session
from sagemaker.jumpstart.estimator import JumpStartEstimator

model_id = "huggingface-llm-gemma-7b-instruct"

estimator = JumpStartEstimator(
    model_id=model_id, environment={"accept_eula": "false"} # manually accept EULA here!
)

# For the other hyperparameters, see the GitHub notebook attached in this blog.
estimator.set_hyperparameters(chat_dataset="True", peft_type="lora", max_input_length="2048", epoch="3")
estimator.fit({"training": <your_S3_bucket_hosting_the_train_data>})

Underlying the training scripts, JumpStart uses HuggingFace SFTTrainer with QLoRA and FlashAttention. Flash Attention enables scaling efficiency, leading to faster training and inference. Besides chat-based fine-tuning, JumpStart also supports instruction and domain adaptation fine-tuning with QLoRA. For details, see the notebook in GitHub.

After the fine-tuning, you can deploy the fine-tuned model to a SageMaker endpoint.

fine_tuned_predictor= estimator.deploy()

You can choose to deploy the model fine-tuned on conversation data in SageMaker endpoint with HuggingFace messages API feature as an alternative approach. The benefit of using Message API feature is that it allow users to pass the query examples in raw format as shown in the following table and chat_template within the endpoint will convert the query examples into the proper format before sending them into the model. For details, see the deployment section after fine-tuning in the notebook.

After deploying the fine-tuned model in an endpoint, you can compare output responses from pre-trained and fine-tuned Gemma models. The examples are collected from test data, which is a hold-out set from the training data. The inference parameters are max_tokens as 512, top_p as 0.9, temperature as 0.4, and top_k as 50.

As you can see with the fine tuning results, the agent answers are more accurate.

To further evaluate your fine tuned models, you may explore FM evaluation with Amazon SageMaker Clarify to see how your fine tuned model performs in comparison to base model.

Clean up

After you have tested the endpoints, delete the SageMaker inference endpoints and the model to avoid incurring charges.

Conclusion

In this post, we showed you how to get started with Gemma in SageMaker Studio and deploy the model for inference. We also showed you how you can fine tune Gemma models on SageMaker Jumpstart.

Because foundation models are pre-trained, they can help lower training and infrastructure costs and enable customization for your use case. Visit SageMaker JumpStart in SageMaker Studio now to get started.

This guidance is for informational purposes only. You should still perform your own independent assessment, and take measures to ensure that you comply with your own specific quality control practices and standards, and the local rules, laws, regulations, licenses and terms of use that apply to you, your content, and the third-party model referenced in this guidance. AWS has no control or authority over the third-party model referenced in this guidance, and does not make any representations or warranties that the third-party model is secure, virus-free, operational, or compatible with your production environment and standards. AWS does not make any representations, warranties or guarantees that any information in this guidance will result in a particular outcome or result.

About the authors

Dr. Kyle Ulrich is an Applied Scientist with the Amazon SageMaker built-in algorithms team. His research interests include scalable machine learning algorithms, computer vision, time series, Bayesian non-parametrics, and Gaussian processes. His PhD is from Duke University and he has published papers in NeurIPS, Cell, and Neuron.

Dr. Xin Huang is a Senior Applied Scientist for Amazon SageMaker JumpStart and Amazon SageMaker built-in algorithms. He focuses on developing scalable machine learning algorithms. His research interests are in the area of natural language processing, explainable deep learning on tabular data, and robust analysis of non-parametric space-time clustering. He has published many papers in ACL, ICDM, KDD conferences, and Royal Statistical Society: Series A.

Rachna Chadha is a Principal Solution Architect AI/ML in Strategic Accounts at AWS. Rachna is an optimist who believes that ethical and responsible use of AI can improve society in future and bring economical and social prosperity. In her spare time, Rachna likes spending time with her family, hiking, and listening to music.

Evan Kravitz is a software engineer at Amazon Web Services, working on SageMaker JumpStart. He enjoys cooking and going on runs in New York City.

Dr. Ashish Khetan is a Senior Applied Scientist with Amazon SageMaker built-in algorithms and helps develop machine learning algorithms. He got his PhD from University of Illinois Urbana-Champaign. He is an active researcher in machine learning and statistical inference, and has published many papers in NeurIPS, ICML, ICLR, JMLR, ACL, and EMNLP conferences.

Read More

Online gaming and social communities offer voice and text chat functionality for their users to communicate. Although voice and text chat often support friendly banter, it can also lead to problems such as hate speech, cyberbullying, harassment, and scams. Today, many companies rely solely on human moderators to review toxic content. However, verifying violations in chat is time-consuming, error-prone, and challenging to scale.

In this post, we introduce solutions that enable audio and text chat moderation using various AWS services, including Amazon Transcribe, Amazon Comprehend, Amazon Bedrock, and Amazon OpenSearch Service.

Social platforms seek an off-the-shelf moderation solution that is straightforward to initiate, but they also require customization for managing diverse policies. Latency and cost are also critical factors that must be taken into account. By orchestrating toxicity classification with large language models (LLMs) using generative AI, we offer a solution that balances simplicity, latency, cost, and flexibility to satisfy various requirements.

The sample code for this post is available in the GitHub repository.

Audio chat moderation workflow

An audio chat moderation workflow could be initiated by a user reporting other users on a gaming platform for policy violations such as profanity, hate speech, or harassment. This represents a passive approach to audio moderation. The system records all audio conversations without immediate analysis. When a report is received, the workflow retrieves the related audio files and initiates the analysis process. A human moderator then reviews the reported conversation, investigating its content to determine if it violates platform policy.

Alternatively, the workflow could be triggered proactively. For instance, in a social audio chat room, the system could record all conversations and apply analysis.

Both passive and proactive approaches can trigger the following pipeline for audio analysis.

The audio moderation workflow involves the following steps:

• The workflow begins with receiving the audio file and storing it on a Amazon Simple Storage Service (Amazon S3) bucket for Amazon Transcribe to access.
• The Amazon Transcribe StartTranscriptionJob API is invoked with Toxicity Detection enabled. Amazon Transcribe converts the audio into text, providing additional information about toxicity analysis. For more information about toxicity analysis, refer to Flag harmful language in spoken conversations with Amazon Transcribe Toxicity Detection.
• If the toxicity analysis returns a toxicity score exceeding a certain threshold (for example, 50%), we can use Knowledge Bases for Amazon Bedrock to evaluate the message against customized policies using LLMs.
• The human moderator receives a detailed audio moderation report highlighting the conversation segments considered toxic and in violation of policy, allowing them to make an informed decision.

The following screenshot shows a sample application displaying toxicity analysis for an audio segment. It includes the original transcription, the results from the Amazon Transcribe toxicity analysis, and the analysis conducted using an Amazon Bedrock knowledge base through the Amazon Bedrock Anthropic Claude V2 model.

The LLM analysis provides a violation result (Y or N) and explains the rationale behind the model’s decision regarding policy violation. Furthermore, the knowledge base includes the referenced policy documents used by the evaluation, providing moderators with additional context.

Amazon Transcribe Toxicity Detection

Amazon Transcribe is an automatic speech recognition (ASR) service that makes it straightforward for developers to add speech-to-text capability to their applications. The audio moderation workflow uses Amazon Transcribe Toxicity Detection, which is a machine learning (ML)-powered capability that uses audio and text-based cues to identify and classify voice-based toxic content across seven categories, including sexual harassment, hate speech, threats, abuse, profanity, insults, and graphic language. In addition to analyzing text, Toxicity Detection uses speech cues such as tones and pitch to identify toxic intent in speech.

The audio moderation workflow activates the LLM’s policy evaluation only when the toxicity analysis exceeds a set threshold. This approach reduces latency and optimizes costs by selectively applying LLMs, filtering out a significant portion of the traffic.

Use LLM prompt engineering to accommodate customized policies

The pre-trained Toxicity Detection models from Amazon Transcribe and Amazon Comprehend provide a broad toxicity taxonomy, commonly used by social platforms for moderating user-generated content in audio and text formats. Although these pre-trained models efficiently detect issues with low latency, you may need a solution to detect violations against your specific company or business domain policies, which the pre-trained models alone can’t achieve.

Additionally, detecting violations in contextual conversations, such as identifying child sexual grooming conversations, requires a customizable solution that involves considering the chat messages and context outside of it, such as user’s age, gender, and conversation history. This is where LLMs can offer the flexibility needed to extend these requirements.

Amazon Bedrock is a fully managed service that offers a choice of high-performing foundation models (FMs) from leading AI companies. These solutions use Anthropic Claude v2 from Amazon Bedrock to moderate audio transcriptions and text chat messages using a flexible prompt template, as outlined in the following code:

Human: You are a Trust & Safety expert. Your job is to review user chat message and decide if it violate the policy.
You will find the chat message in <message> tag, and find the policy in the <policy> tag. You can find additional rules in the <rule> tag to assist your decision. 

<policy>{policy}</policy>
<message>{message}</message>
<rule>{rule}</rule>

Does the chat message violate the policy? Please consider and provide your analysis in the <analysis> tag, breaking down each rule in the rule section, and keep and analysis within 100 words. Respond in the <answer> tag with either 'Y' or 'N'. 'Y' indicates that the message violates the policy, while 'N' means the content is safe and does not violate the policy. 

Assistant:

The template contains placeholders for the policy description, the chat message, and additional rules that requires moderation. The Anthropic Claude V2 model delivers responses in the instructed format (Y or N), along with an analysis explaining why it thinks the message violates the policy. This approach allows you to define flexible moderation categories and articulate your policies in human language.

The traditional method of training an in-house classification model involves cumbersome processes such as data annotation, training, testing, and model deployment, requiring the expertise of data scientists and ML engineers. LLMs, in contrast, offer a high degree of flexibility. Business users can modify prompts in human language, leading to enhanced efficiency and reduced iteration cycles in ML model training.

Amazon Bedrock knowledge bases

Although prompt engineering is efficient for customizing policies, injecting lengthy policies and rules directly into LLM prompts for each message may introduce latency and increase cost. To address this, we use Amazon Bedrock knowledge bases as a managed Retrieval Augmented Generation (RAG) system. This enables you to manage the policy document flexibly, allowing the workflow to retrieve only the relevant policy segments for each input message. This minimizes the number of tokens sent to the LLMs for analysis.

You can use the AWS Management Console to upload the policy documents to an S3 bucket and then index the documents to a vector database for efficient retrieval. The following is a conceptual workflow managed by an Amazon Bedrock knowledge base that retrieves documents from Amazon S3, splits the text into chunks, and invokes the Amazon Bedrock Titan text embeddings model to convert the text chunks into vectors, which are then stored in the vector database.

In this solution, we use Amazon OpenSearch Service as the vector store. OpenSearch is a scalable, flexible, and extensible open source software suite for search, analytics, security monitoring, and observability applications, licensed under the Apache 2.0 license. OpenSearch Service is a fully managed service that makes it straightforward to deploy, scale, and operate OpenSearch in the AWS Cloud.

After the document is indexed in OpenSearch Service, the audio and text moderation workflow sends chat messages, triggering the following query flow for customized policy evaluation.

The process is similar to the initiation workflow. First, the text message is converted to text embeddings using the Amazon Bedrock Titan Text Embedding API. These embeddings are then used to perform a vector search against the OpenSearch Service database, which has already been populated with document embeddings. The database returns policy chunks with the highest matching score, relevant to the input text message. We then compose prompts containing both the input chat message and the policy segment, which are sent to Anthropic Claude V2 for evaluation. The LLM model returns an analysis result based on the prompt instructions.

For detailed instructions on how to create a new instance with your policy document in an Amazon Bedrock knowledge base, refer to Knowledge Bases now delivers fully managed RAG experience in Amazon Bedrock.

Text chat moderation workflow

The text chat moderation workflow follows a similar pattern to audio moderation, but it uses Amazon Comprehend toxicity analysis, which is tailored for text moderation. The sample app supports an interface for uploading bulk text files in CSV or TXT format and provides a single-message interface for quick testing. The following diagram illustrates the workflow.

The text moderation workflow involves the following steps:

• The user uploads a text file to an S3 bucket.
• Amazon Comprehend toxicity analysis is applied to the text message.
• If the toxicity analysis returns a toxicity score exceeding a certain threshold (for example, 50%), we use an Amazon Bedrock knowledge base to evaluate the message against customized policies using the Anthropic Claude V2 LLM.
• A policy evaluation report is sent to the human moderator.

Amazon Comprehend toxicity analysis

In the text moderation workflow, we use Amazon Comprehend toxicity analysis to assess the toxicity level of the text messages. Amazon Comprehend is a natural language processing (NLP) service that uses ML to uncover valuable insights and connections in text. The Amazon Comprehend toxicity detection API assigns an overall toxicity score to text content, ranging from 0–1, indicating the likelihood of it being toxic. It also categorizes text into the following categories and provides a confidence score for each: hate_speech, graphic, harrassement_or_abuse, sexual, violence_or_threat, insult, and profanity.

In this text moderation workflow, Amazon Comprehend toxicity analysis plays a crucial role in identifying whether the incoming text message contains toxic content. Similar to the audio moderation workflow, it includes a condition to activate the downstream LLM policy evaluation only when the toxicity analysis returns a score exceeding a predefined threshold. This optimization helps reduce overall latency and cost associated with LLM analysis.

Summary

In this post, we introduced solutions for audio and text chat moderation using AWS services, including Amazon Transcribe, Amazon Comprehend, Amazon Bedrock, and OpenSearch Service. These solutions use pre-trained models for toxicity analysis and are orchestrated with generative AI LLMs to achieve the optimal balance in accuracy, latency, and cost. They also empower you to flexibly define your own policies.

You can experience the sample app by following the instructions in the GitHub repo.

About the author

Lana Zhang is a Senior Solutions Architect at AWS WWSO AI Services team, specializing in AI and ML for Content Moderation, Computer Vision, Natural Language Processing and Generative AI. With her expertise, she is dedicated to promoting AWS AI/ML solutions and assisting customers in transforming their business solutions across diverse industries, including social media, gaming, e-commerce, media, advertising & marketing.

Read More

Advancements in artificial intelligence (AI) and machine learning (ML) are revolutionizing the financial industry for use cases such as fraud detection, credit worthiness assessment, and trading strategy optimization. To develop models for such use cases, data scientists need access to various datasets like credit decision engines, customer transactions, risk appetite, and stress testing. Managing appropriate access control for these datasets among the data scientists working on them is crucial to meet stringent compliance and regulatory requirements. Typically, these datasets are aggregated in a centralized Amazon Simple Storage Service (Amazon S3) location from various business applications and enterprise systems. Data scientists across business units working on model development using Amazon SageMaker are granted access to relevant data, which can lead to the requirement of managing prefix-level access controls. With an increase in use cases and datasets using bucket policy statements, managing cross-account access per application is too complex and long for a bucket policy to accommodate.

Amazon S3 Access Points simplify managing and securing data access at scale for applications using shared datasets on Amazon S3. You can create unique hostnames using access points to enforce distinct and secure permissions and network controls for any request made through the access point.

S3 Access Points simplifies the management of access permissions specific to each application accessing a shared dataset. It enables secure, high-speed data copy between same-Region access points using AWS internal networks and VPCs. S3 Access Points can restrict access to VPCs, enabling you to firewall data within private networks, test new access control policies without impacting existing access points, and configure VPC endpoint policies to restrict access to specific account ID-owned S3 buckets.

This post walks through the steps involved in configuring S3 Access Points to enable cross-account access from a SageMaker notebook instance.

Solution overview

For our use case, we have two accounts in an organization: Account A (111111111111), which is used by data scientists to develop models using a SageMaker notebook instance, and Account B (222222222222), which has required datasets in the S3 bucket test-bucket-1. The following diagram illustrates the solution architecture.

To implement the solution, complete the following high-level steps:

1. Configure Account A, including VPC, subnet security group, VPC gateway endpoint, and SageMaker notebook.
2. Configure Account B, including S3 bucket, access point, and bucket policy.
3. Configure AWS Identity and Access Management (IAM) permissions and policies in Account A.

You should repeat these steps for each SageMaker account that needs access to the shared dataset from Account B.

The names for each resource mentioned in this post are examples; you can replace them with other names as per your use case.

Configure Account A

Complete the following steps to configure Account A:

1. Create a VPC called DemoVPC.
2. Create a subnet called DemoSubnet in the VPC DemoVPC.
3. Create a security group called DemoSG.
4. Create a VPC S3 gateway endpoint called DemoS3GatewayEndpoint.
5. Create the SageMaker execution role.
6. Create a notebook instance called DemoNotebookInstance and the security guidelines as outlined in How to configure security in Amazon SageMaker.
   1. Specify the Sagemaker execution role you created.
   2. For the notebook network settings, specify the VPC, subnet, and security group you created.
   3. Make sure that Direct Internet access is disabled.

You assign permissions to the role in subsequent steps after you create the required dependencies.

Configure Account B

To configure Account B, complete the following steps:

1. In Account B, create an S3 bucket called test-bucket-1 following Amazon S3 security guidance.
2. Upload your file to the S3 bucket.
3. Create an access point called test-ap-1 in Account B.
   1. Don’t change or edit any Block Public Access settings for this access point (all public access should be blocked).
4. Attach the following policy to your access point:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "AWS": “arn:aws:iam:: 111111111111:role/demo ”
            },
            "Action": ["s3:GetObject", "s3:GetObjectVersion", "s3:PutObject", "s3:PutObjectAcl"]
            "Resource": [
                “arn:aws:s3:us-east-1: 222222222222:accesspoint/test-ap-1”,
                " arn:aws:s3:us-east-1: 222222222222:accesspoint/test-ap-1/object/*"
            ]
        }
    ]
}

The actions defined in the preceding code are sample actions for demonstration purposes. You can define the actions as per your requirements or use case.

5. Add the following bucket policy permissions to access the access point:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "AWS": " arn:aws:iam:: 111111111111:role/demo "
            },
            "Action" : ["s3:GetObject","s3:ListBucket"],
            "Resource" : ["arn:aws:s3:::test-bucket-1 ”, " arn:aws:s3:::test-bucket-1/*"]
            "Condition": {
                "StringEquals": {
                    "s3:DataAccessPointAccount": "222222222222"
                }
            }
        }
    ]
}

The preceding actions are examples. You can define the actions as per your requirements.

Configure IAM permissions and policies

Complete the following steps in Account A:

1. Confirm that the SageMaker execution role has the AmazonSagemakerFullAccess custom IAM inline policy, which looks like the following code:

{
            "Sid": "VisualEditor2",
            "Effect": "Allow",
            " Action": ["s3:GetObject", "s3:GetObjectVersion", "s3:PutObject", "s3:PutObjectAcl"]
            "Resource": [
                “arn:aws:s3:us-east-1: 222222222222:accesspoint/test-ap-1 ”,
                "arn:aws:s3:us-east-1: 222222222222:accesspoint/test-ap-1 /object/*”,                             "arn:aws:s3:::test-bucket-1”,
                " arn:aws:s3:::test-bucket-1/*"
            ]
}

The actions in the policy code are sample actions for demonstration purposes.

2. Go to the DemoS3GatewayEndpoint endpoint you created and add the following permissions:

{

	"Version": "2012-10-17",
	"Statement": [
		{
			"Sid": "AllowCrossAccountAccessThroughAccessPoint",
			"Effect": "Allow",
			"Principal": "*",
			"Action": [
				"s3:Get*",
				"s3:List*",
				"s3:Put*"
			],
			"Resource": ": [
                “arn:aws:s3:us-east-1: 222222222222:accesspoint/test-ap-1 ”,
                "arn:aws:s3:us-east-1: 222222222222:accesspoint/test-ap-1 /object/*”,                             "arn:aws:s3:::test-bucket-1 ”,
                " arn:aws:s3:::test-bucket-1/*"
            ]
 
		}
	]
}

3. To get a prefix list, run the AWS Command Line Interface (AWS CLI) describe-prefix-lists command:

aws ec2 describe-prefix-lists

4. In Account A, Go to the security group DemoSG for the target SageMaker notebook instance
5. Under Outbound rules, create an outbound rule with All traffic or All TCP, and then specify the destination as the prefix list ID you retrieved.

This completes the setup in both accounts.

Test the solution

To validate the solution, go to the SageMaker notebook instance terminal and enter the following commands to list the objects through the access point:

• To list the objects successfully through S3 access point test-ap-1:

aws s3 ls arn:aws:s3:us-east-1:222222222222:accesspoint/Test-Ap-1

• To get the objects successfully through S3 access point test-ap-1:

aws s3api get-object --bucket arn:aws:s3:us-east-1:222222222222:accesspoint/test-ap-1 --key sample2.csv test2.csv

Clean up

When you’re done testing, delete any S3 access points and S3 buckets. Also, delete any Sagemaker notebook instances to stop incurring charges.

Conclusion

In this post, we showed how S3 Access Points enables cross-account access to large, shared datasets from SageMaker notebook instances, bypassing size constraints imposed by bucket policies while configuring at-scale access management on shared datasets.

To learn more, refer to Easily Manage Shared Data Sets with Amazon S3 Access Points.

About the authors

Kiran Khambete is working as Senior Technical Account Manager at Amazon Web Services (AWS). As a TAM, Kiran plays a role of technical expert and strategic guide to helping Enterprise customers achieving their business goals.

Ankit Soni with total experience of 14 years holds the position of Principal Engineer at NatWest Group, where he has served as a Cloud Infrastructure Architect for the past six years.

Kesaraju Sai Sandeep is a Cloud Engineer specializing in Big Data Services at AWS.

Read More

Advertisers, publishers, and advertising technology providers are actively seeking efficient ways to collaborate with their partners to generate insights about their collective datasets. One common reason to engage in data collaboration is to run an audience overlap analysis, which is a common analysis to run when media planning and evaluating new partnerships.

In this post, we explore what an audience overlap analysis is, discuss the current technical approaches and their challenges, and illustrate how you can run secure audience overlap analysis using AWS Clean Rooms.

Audience overlap analysis

Audience overlap is the percentage of users in your audience who are also present in another dataset (calculated as the number of users present in both your audience and another dataset divided by the total number of users in your audience). In the digital media planning process, audience overlaps are often conducted to compare an advertiser’s first-party dataset with a media partner’s (publisher) dataset. The analysis helps determine how much of the advertiser’s audience can be reached by a given media partner. By evaluating the overlap, advertisers can determine whether a media partner provides unique reach or if the media partner’s audience predominantly overlaps with the advertiser’s existing audience.

Current approaches and challenges

Advertisers, publishers, third-party data providers, and other entities often share their data when running audience overlaps or match tests. Common methods for sharing data, such as using pixels and SFTP transfers, can carry risk because they involve moving sensitive customer information. Sharing this data to another party can be time consuming and increase the risk of potential data breaches or unauthorized access. If the receiving party mishandles the data, it could violate privacy regulations, resulting in legal risks. Also, any perceived misuse or exposure of customer data can erode consumer trust, leading to reputational damage and potential loss of business.

Solution overview

AWS Clean Rooms can help you and your partners effortlessly and securely collaborate on and analyze your collective datasets—without copying each other’s underlying data. With AWS Clean Rooms, you can create a data clean room in minutes and collaborate with your partners to generate unique insights. AWS Clean Rooms allows you to run an audience overlap analysis and generate valuable insights while avoiding risks associated with other current approaches.

The following are key concepts and prerequisites to use AWS Clean Rooms:

• Each party in the analysis (collaboration member) needs to have an AWS account.
• One member invites the other member to the AWS Clean Rooms collaboration. It doesn’t matter which member creates the invitation. The collaboration creator uses the invitee’s AWS account ID as input to send invitations.
• Only one member can query in the collaboration, and only one member can receive results from the collaboration. The abilities of each member are defined when the collaboration is created.
• Each collaboration member stores datasets in their respective Amazon Simple Storage Service (Amazon S3) bucket and catalogs them (creates a schema with column names and data types) in the AWS Glue Data Catalog. You can also create the Data Catalog definition using the Amazon Athena create database and create table statements.
• Collaborators need to have their S3 buckets and Data Catalog tables in the same AWS Region.
• Collaborators can use the AWS Clean Rooms console, APIs, or AWS SDKs to set up a collaboration.
• AWS Clean Rooms enables you to use any column as a join key, for example hashed MAIDs, emails, IP addresses, and RampIDs.
• Each collaboration member associates their own data to the collaboration.

Let’s look at a scenario in which an advertiser collaborates with a publisher to identify the audience overlap. In this example, the publisher creates the collaboration, invites the advertiser, and designates the advertiser as the member who can query and receive results.

Prerequisites

To invite another person to a collaboration, you need their AWS account ID. In our use case, the publisher needs the AWS account ID of the advertiser.

Create a collaboration

In our use case, the publisher creates a collaboration using the AWS Clean Rooms console and invites the advertiser.

To create a collaboration, complete the following steps:

1. On the AWS Clean Rooms, console, choose Collaborations in the navigation pane.
2. Choose Create collaboration.
3. For Name, enter a name for the collaboration.
4. In the Members section, enter the AWS account ID of the account you want to invite (in this case, the advertiser).
5. In the Member abilities section, choose the member who can query and receive results (in this case, the advertiser).
6. For Query logging, decide if you want query logging turned on. The queries are logged to Amazon CloudWatch.
7. For Cryptographic computing, decide if you want to turn on support for cryptographic computing (pre-encrypt your data before associating it). AWS Clean Rooms will then run queries on the encrypted data.
8. Choose Next.
9. On the Configure membership page, choose if you want to create the membership and collaboration now, or create the collaboration but activate your membership later.
10. For Query results settings defaults, choose if you want to keep the default settings to receive results.
11. For Log storage in Amazon CloudWatch Logs, specify your log settings.
12. Specify any tags and who is paying for queries.
13. Choose Next.
14. Review the configuration and choose to either create the collaboration and membership now, or just the collaboration.

The publisher sends an invitation to the advertiser. The advertiser reviews the collaboration settings and creates a membership.

Create a configured table and set analysis rules

The publisher creates a configured table from the AWS Glue table (which represents the metadata definition of the S3 data, including location, so it can be read by AWS Clean Rooms when the query is run).

Complete the following steps:

1. On the AWS Clean Rooms console, choose Configured tables in the navigation pane.
2. Choose Configure new table.
3. In the Choose AWS Glue table section, choose your database and table.
4. In the Columns allowed in collaboration section, choose which of the existing table columns to allow for querying in the collaboration.
5. In the Configured table details section, enter a name and optional description for the configured table.
6. Choose Configure new table.
7. Choose the analysis rule type that matches the type of queries you want to allow on table. To allow an aggregation analysis, such as finding the size of the audience overlap, choose the aggregation analysis rule type.
8. In the Aggregate functions section, choose COUNT DISTINCT as the aggregate function.
9. In the Join controls section, choose whether your collaborator is required to join a table with yours. Because this is an audience overlap use case, select No, only overlap can be queried.
10. Select the operators to allow for matching (for this example, select AND and OR).
11. In the Dimension controls section, choose if you want to make any columns available as dimensions.
12. In the Scalar functions section, choose if you want to limit the scalar functions allowed.
13. Choose Next.
14. In the Aggregation constraints section, choose the minimum aggregation constraint for the configured table.

This allows you to filter out rows that don’t meet a certain minimum threshold of users (for example, if the threshold is set to 10, rows that aggregate fewer than 10 users are filtered out).

15. Choose Next.
16. Review the settings and create the table.

Associate the table to the collaboration

AWS Clean Rooms requires access to read the table in order to run the query submitted by the advertiser. Complete the following steps to associate the table:

1. On the AWS Clean Rooms console, navigate to your collaboration.
2. Choose Associate table.
3. For Configured table name, choose the name of your configured table.
4. In the Table association details section, enter a name and optional description for the table.
5. In the Service access section, you can choose to can use the default settings to create an AWS Identity and Access Management (IAM) service role for AWS Clean Rooms automatically, or you can use an existing role. IAM permissions are required to create or modify the role and pass the role to AWS Clean Rooms.
6. Choose Associate table.

The advertiser also completes the steps detailed in the preceding sections to create a configured table and associate it to the collaboration.

Run queries in the query editor

The advertiser can now navigate to the Queries tab for the collaboration and review tables to query and their analysis rules. You can specify

the S3 bucket where the output of the overlap query will go.

The advertiser can now write and run an overlap query. You can use a hashed email as a join key for the query (you have the option to use any column as the join key and can also use multiple columns for multiple join keys). You can also use the Analysis Builder no-code option to have AWS Clean Rooms generate SQL on your behalf. For our use case, we run the following queries:

#Query 1 – count of overlapping users between advertiser and publisher datasets

SELECT COUNT(DISTINCT advertiser.hashed_email)
FROM consumer as advertiser
INNER JOIN impressions as publisher
ON advertiser.hashed_email = publisher.hashed_email

#Query 2 – count of users in advertiser dataset

SELECT COUNT(DISTINCT advertiser.hashed_email)
FROM consumer as advertiser

The query results are sent to the advertiser’s S3 bucket, as shown in the following screenshot.

Clean up

It’s a best practice to delete resources that are no longer being used. The advertiser and publisher should clean up their respective resources:

• Advertiser – The advertiser deletes their configured table associations and collaboration membership. However, they don’t have to delete their configured table because it’s reusable across collaborations.
• Publisher – The publisher deletes their configured table associations and the collaboration. They don’t have to delete their configured table because it’s reusable across collaborations.

Conclusion

In this post, we demonstrated how to set up an audience overlap collaboration using AWS Clean Rooms for media planning and partnership evaluation using a hashed email as a join key between datasets. Advertisers are increasingly turning to AWS Clean Rooms to conduct audience overlap analyses with their media partners, aiding their media investment decisions. Furthermore, audience overlaps help you accelerate your partnership evaluations by identifying the extent of overlap you share with potential partners.

To learn more about AWS Clean Rooms, watch the video Getting Started with AWS Clean Rooms, and refer to the following additional resources:

• AWS Clean Rooms Now Generally Available — Collaborate with Your Partners without Sharing Raw Data
• AWS on Air: AWS Clean Rooms is now available for General Availability
• Introducing AWS for Advertising & Marketing: Helping customers reinvent the industry with purpose-built services, solutions, and partners
• Introducing four new solutions that help customers integrate AWS Clean Rooms into their advertising workflows
• AWS Clean Rooms User Guide

About the Authors

Eric Saccullo is a Senior Business Development Manager for AWS Clean Rooms at Amazon Web Services. He is focused on helping customers collaborate with their partners in privacy-enhanced ways to gain insights and improve business outcomes.

Shamir Tanna is a Senior Technical Product Manager at Amazon Web Services.

Ryan Malecky is a Senior Solutions Architect at Amazon Web Services. He is focused on helping customers gain insights from their data, especially with AWS Clean Rooms.

Read More

This post is co-written with Justin Miles, Liv d’Aliberti, and Joe Kovba from Leidos. 

Leidos is a Fortune 500 science and technology solutions leader working to address some of the world’s toughest challenges in the defense, intelligence, homeland security, civil, and healthcare markets. In this post, we discuss how Leidos worked with AWS to develop an approach to privacy-preserving large language model (LLM) inference using AWS Nitro Enclaves.

LLMs are designed to understand and generate human-like language, and are used in many industries, including government, healthcare, financial, and intellectual property. LLMs have broad applicability, including chatbots, content generation, language translation, sentiment analysis, question answering systems, search engines, and code generation. Introducing LLM-based inference into a system also has the potential to introduce privacy threats, including model exfiltration, data privacy violations, and unintended LLM-based service manipulation. Technical architectures need to be implemented in order to make sure that LLMs don’t expose sensitive information during inference.

This post discusses how Nitro Enclaves can help protect LLM model deployments, specifically those that use personally identifiable information (PII) or protected health information (PHI). This post is for educational purposes only and should not be used in production environments without additional controls.

Overview of LLMs and Nitro Enclaves

A potential use case is an LLM-based sensitive query chatbot designed to carry out a question and answering service containing PII and PHI. Most current LLM chatbot solutions explicitly inform users that they should not include PII or PHI when inputting questions due to security concerns. To mitigate these concerns and protect customer data, service owners rely primarily on user protections such as the following:

• Redaction – The process of identifying and obscuring sensitive information like PII in documents, texts, or other forms of content. This can be accomplished with input data before being sent to a model or an LLM trained to redact their responses automatically.
• Multi-factor authentication – A security process that requires users to provide multiple authentication methods to verify their identity to gain access to the LLM.
• Transport Layer Security (TLS) – A cryptographic protocol that provides secure communication that enhances data privacy in transit between users and the LLM service.

Although these practices enhance the security posture of the service, they are not sufficient to safeguard all sensitive user information and other sensitive information that can persist without the user’s knowledge.

In our example use case, an LLM service is designed to answer employee healthcare benefit questions or provide a personal retirement plan. Let’s analyze the following sample architecture and identify data privacy risk areas.

Figure 1 – Data Privacy Risk Areas Diagram

The potential risk areas are as follows:

1. Privileged users have access to the instance that houses the server. Unintentional or unauthorized changes to the service could result in sensitive data being exposed in unintended ways.
2. Users must trust the service will not expose or retain sensitive information in application logs.
3. Changes to application packages can cause changes to the service, resulting in the exposure of sensitive data.
4. Privileged users with access to the instance have unrestricted access to the LLM used by the service. Changes may cause incorrect or inaccurate information being returned to users.

Nitro Enclaves provides additional isolation to your Amazon Elastic Compute Cloud (Amazon EC2) instance, safeguarding data in use from unauthorized access, including admin-level users. In the preceding architecture, it’s possible for an unintentional change to result in sensitive data to persist in plaintext and accidentally get revealed to a user who may not need to access that data. With Nitro Enclaves, you create an isolated environment from your EC2 instance, permitting you to allocate CPU and memory resources to the enclave. This enclave is a highly restrictive virtual machine. By running code that handles sensitive data within the enclave, none of the parent’s processes will be able to view enclave data.

Nitro Enclaves offers the following benefits:

• Memory and CPU Isolation – It relies on the Nitro Hypervisor to isolate the CPU and memory of the enclave from users, applications, and libraries on the parent instance. This feature helps isolate the enclave and your software, and significantly reduces the surface area for unintended events.
• Separate virtual machine – Enclaves are separated virtual machines attached to an EC2 instance to further protect and securely process highly sensitive data.
• No interactive access – Enclaves provide only secure local socket connectivity with their parent instance. They have no persistent storage, interactive access, or external networking.
• Cryptographic attestation – Nitro Enclaves offers cryptographic attestation, a process used to prove the identity of an enclave and verify that only authorized code is running in your enclave.
• AWS integration – Nitro Enclaves is integrated with AWS Key Management Service (AWS KMS), allowing you to decrypt files that have been encrypted using AWS KMS inside the enclave. AWS Certificate Manager (ACM) for Nitro Enclaves allows you to use public and private SSL/TLS certificates with your web applications and servers running on EC2 instances with Nitro Enclaves.

You can use these features provided by Nitro Enclaves to help mitigate risks associated with PII and PHI data. We recommend including Nitro Enclaves in an LLM service when handling sensitive user data.

Solution overview

Let’s examine the architecture of the example service, now including Nitro Enclaves. By incorporating Nitro Enclaves, as shown in the following figure, the LLM becomes a more secure chatbot for handling PHI or PII data.

Figure 2 – Solution Overview Diagram

User data, including PII, PHI, and questions, remains encrypted throughout the request-response process when the application is hosted within an enclave. The steps carried out during the inference are as follows:

1. The chatbot app generates temporary AWS credentials and asks the user to input a question. The question, which may contain PII or PHI, is then encrypted via AWS KMS. The encrypted user input is combined with the temporary credentials to create the encrypted request.
2. The encrypted data is sent to an HTTP server hosted by Flask as a POST request. Before accepting sensitive data, this endpoint should be configured for HTTPs.
3. The client app receives the POST request and forwards it through a secure local channel (for example, vsock) to the server app running inside Nitro Enclaves.
4. The Nitro Enclaves server app uses the temporary credentials to decrypt the request, queries the LLM, and generates the response. The model-specific settings are stored within the enclaves and are protected with cryptographic attestation.
5. The server app uses the same temporary credentials to encrypt the response.
6. The encrypted response is returned back to the chatbot app through the client app as a response from the POST request.
7. The chatbot app decrypts the response using their KMS key and displays the plaintext to the user.

Prerequisites

Before we get started, you need the following prerequisites to deploy the solution:

• AWS Identity and Access Management (IAM) user
• KMS symmetric key
• AWS Cloud9 environment

Configure an EC2 instance

Complete the following steps to configure an EC2 instance:

1. Launch an r5.8xlarge EC2 instance using the amzn2-ami-kernel-5.10-hvm-2.0.20230628.0-x86_64-gp2 AMI with Nitro Enclaves enabled.
2. Install the Nitro Enclaves CLI to build and run Nitro Enclaves applications:
   • sudo amazon-linux-extras install aws-nitro-enclaves-cli -y
   • sudo yum install aws-nitro-enclaves-cli-devel -y
3. Verify the installation of the Nitro Enclaves CLI:
   • nitro-cli –version
   • The version used in this post is 1.2.2
4. Install Git and Docker to build Docker images and download the application from GitHub. Add your instance user to the Docker group (<USER> is your IAM instance user):
   • sudo yum install git -y
   • sudo usermod -aG ne <USER>
   • sudo usermod -aG docker <USER>
   • sudo systemctl start docker && sudo systemctl enable docker
5. Start and enable the Nitro Enclaves allocator and vsock proxy services:
   • sudo systemctl start nitro-enclaves-allocator.service && sudo systemctl enable nitro-enclaves-allocator.service
   • sudo systemctl start nitro-enclaves-vsock-proxy.service && sudo systemctl enable nitro-enclaves-vsock-proxy.service

Nitro Enclaves uses a local socket connection called vsock to create a secure channel between the parent instance and the enclave.

After all the services are started and enabled, restart the instance to verify that all of the user groups and services are running correctly:

sudo shutdown -r now

Configure the Nitro Enclaves allocator service

Nitro Enclaves is an isolated environment that designates a portion of the instance CPU and memory to run the enclave. With the Nitro Enclaves allocator service, you can indicate how many CPUs and how much memory will be taken from the parent instance to run the enclave.

Modify the enclave’s reserved resources using a text editor (for our solution, we allocate 8 CPU and 70,000 MiB memory to provide enough resources):

vi /etc/nitro_enclaves/allocatory.yaml

Figure 3 – AWS Nitro Enclaves Allocator Service Configuration

Clone the project

After you configure the EC2 instance, you can download the code to run the sensitive chatbot with an LLM inside of Nitro Enclaves.

You need to update the server.py file with the appropriate KMS key ID that you created in the beginning to encrypt the LLM response.

1. Clone the GitHub project:
   • cd ~/ && git clone https://<THE_REPO.git>
2. Navigate to the project folder to build the enclave_base Docker image that contains the Nitro Enclaves Software Development Kit (SDK) for cryptographic attestation documents from the Nitro Hypervisor (this step can take up to 15 minutes):
   • cd /nitro_llm/enclave_base
   • docker build ./ -t “enclave_base”

Save the LLM in the EC2 Instance

We are using the open-source Bloom 560m LLM for natural language processing to generate responses. This model is not fine-tuned to PII and PHI, but demonstrates how an LLM can live inside of an enclave. The model also needs to be saved on the parent instance so that it can be copied into the enclave via the Dockerfile.

1. Navigate to the project:
   • cd /nitro_llm
2. Install the necessary requirements to save the model locally:
   • pip3 install requirements.txt
3. Run the save_model.py app to save the model within the /nitro_llm/enclave/bloom directory:
   • python3 save_model.py

Build and run the Nitro Enclaves image

To run Nitro Enclaves, you need to create an enclave image file (EIF) from a Docker image of your application. The Dockerfile located in the enclave directory contains the files, code, and LLM that will run inside of the enclave.

Building and running the enclave will take multiple minutes to complete.

1. Navigate to the root of the project:
   • cd /nitro_llm
2. Build the enclave image file as enclave.eif:
   • nitro-cli build-enclave --docker-uri enclave:latest --output-file enclave.eif

Figure 4 – AWS Nitro Enclaves Build Result

When the enclave is built, a series of unique hashes and platform configuration registers (PCRs) will be created. The PCRs are a contiguous measurement to prove the identity of the hardware and application. These PCRs will be required for cryptographic attestation and used during the KMS key policy update step.

3. Run the enclave with the resources from the allocator.service (adding the --attach-console argument at the end will run the enclave in debug mode):
   • nitro-cli run-enclave --cpu-count 8 --memory 70000 --enclave-cid 16 --eif-path enclave.eif

You need to allocate at least four times the EIF file size. This can be modified in the allocator.service from previous steps.

4. Verify the enclave is running with the following command:
   • nitro-cli describe-enclaves

Figure 5 – AWS Nitro Enclave Describe Command

Update the KMS key policy

Complete the following steps to update your KMS key policy:

1. On the AWS KMS console, choose Customer managed keys in the navigation pane.
2. Search for the key that you generated as a prerequisite.
3. Choose Edit on the key policy.
4. Update the key policy with the following information:
   • Your account ID
   • Your IAM user name
   • The updated Cloud9 environment instance role
   • Actions kms:Encrypt and kms:Decrypt
   • Enclave PCRs (for example, PCR0, PCR1, PCR2) to your key policy with a condition statement

See the following key policy code:

{
   "Version":"2012-10-17",
   "Id":"key-default-1",
   "Statement":[
      {
         "Sid":"Enable User permissions",
         "Effect":"Allow",
         "Principal":{
            "AWS":"arn:aws:iam:::user/"
         },
         "Action":[
            "kms:CreateAlias",
            "kms:CreateKey",
            "kms:DeleteAlias",
            "kms:Describe*",
            "kms:GenerateRandom",
            "kms:Get*",
            "kms:List*",
            "kms:TagResource",
            "kms:UntagResource",
            "iam:ListGroups",
            "iam:ListRoles",
            "iam:ListUsers"
         ],
         "Resource":"*"
      },
      {
         "Sid":"Enable Enclave permissions",
         "Effect":"Allow",
         "Principal":{
            "AWS":"arn:aws:iam:::role/"
         },
         "Action":[
            "kms:Encrypt",
            "kms:Decrypt"
         ],
         "Resource":"*",
         "Condition":{
            "StringEqualsIgnoreCase":{
               "kms:RecipientAttestation:PCR0":"",
               "kms:RecipientAttestation:PCR1":"",
               "kms:RecipientAttestation:PCR2":""
            }
         }
      }
   ]
}

Save the chatbot app

To mimic a sensitive query chatbot application that lives outside of the AWS account, you need to save the chatbot.py app and run it inside the Cloud9 environment. Your Cloud9 environment will use its instance role for temporary credentials to disassociate permissions from the EC2 running the enclave. Complete the following steps:

1. On the Cloud9 console, open the environment you created.
2. Copy the following code into a new file like chatbot.py into the main directory.
3. Install the required modules:
   • pip install boto3
   • Pip install requests
4. On the Amazon EC2 console, note the IP associated with your Nitro Enclaves instance.
5. Update the URL variable in http://<ec2instanceIP>:5001.

"""
Modules for a basic chatbot like application and AWS communications
"""
import base64
import requests
import boto3
 
def get_identity_document():
    """
    Get identity document for current EC2 Host
    """
    identity_doc = requests.get(
        "http://169.254.169.254/latest/dynamic/instance-identity/document", timeout=30)
    return identity_doc
 
def get_region(identity):
    """
    Get account of current instance identity
    """
    region = identity.json()["region"]
    return region
 
def get_account(identity):
    """
    Get account of current instance identity
    """
    account = identity.json()["accountId"]
    return account
 
def set_identity():
    """
    Set region and account for KMS
    """
    identity = get_identity_document()
    region = get_region(identity)
    account = get_account(identity)
    return region, account
 
def prepare_server_request(ciphertext):
    """
    Get the AWS credential from EC2 instance metadata
    """
    instance_prof = requests.get(
        "http://169.254.169.254/latest/meta-data/iam/security-credentials/", timeout=30)
    instance_profile_name = instance_prof.text
 
    instance_prof_json = requests.get(
        f"http://169.254.169.254/latest/meta-data/iam/security-credentials/{instance_profile_name}",
        timeout=30)
    response = instance_prof_json.json()
 
    credential = {
        'access_key_id': response['AccessKeyId'],
        'secret_access_key': response['SecretAccessKey'],
        'token': response['Token'],
        'region': REGION,
        'ciphertext': ciphertext
    }
    return credential
 
def get_user_input():
    """
    Start chatbot to collect user input
    """
    print("Chatbot: Hello! How can I assist you?")
    user_input = input('Your Question: ')
    return user_input.lower()
 
def encrypt_string(user_input, alias, kms):
    """
    Encrypt user input using AWS KMS
    """
    file_contents = user_input
    encrypted_file = kms.encrypt(KeyId=f'alias/{alias}', Plaintext=file_contents)
    encrypted_file_contents = encrypted_file[u'CiphertextBlob']
    encrypted_file_contents_base64 = base64.b64encode(encrypted_file_contents)
    return encrypted_file_contents_base64.decode()
 
def decrypt_data(encrypted_data, kms):
    """
    Decrypt the LLM response using AWS KMS
    """
    try:
        ciphertext_blob = base64.b64decode(encrypted_data)
        response = kms.decrypt(CiphertextBlob=ciphertext_blob)
        decrypted_data = response['Plaintext'].decode()
        return decrypted_data
    except ImportError as e_decrypt:
        print("Decryption failed:", e_decrypt)
        return None
 
REGION, ACCOUNT = set_identity()
  
def main():
    """
    Main function to encrypt/decrypt data and send/receive with parent instance
    """
    kms = boto3.client('kms', region_name=REGION)
    alias = "ncsnitro"
    user_input = get_user_input()
    encrypted_input = encrypt_string(user_input, alias, kms)
    server_request = prepare_server_request(encrypted_input)
    url = 'http://<EC2 Instance Private IP>:5001'
    x = requests.post(url, json = server_request)
    response_body = x.json()
    llm_response = decrypt_data(response_body["EncryptedData"], kms)
    print(llm_response)
 
if __name__ == '__main__':
    main()

6. Run the chatbot application:

• • python3 chat.py

When it’s running, the terminal will ask for the user input and follow the architectural diagram from earlier to generate a secure response.

Run the private question and answer chatbot

Now that Nitro Enclaves is up and running on the EC2 instance, you can more securely ask your chatbot PHI and PII questions. Let’s look at an example.

Within the Cloud9 environment, we ask our chatbot a question and provide our user name.

Figure 6 – Asking the Chat Bot a Question

AWS KMS encrypts the question, which looks like the following screenshot.

Figure 7 – Encrypted Question

It is then sent to the enclave and asked of the secured LLM. The question and response of the LLM will look like the following screenshot (the result and encrypted response are visible inside the enclave only in debug mode).

Figure 8 – Response from LLM

The result is then encrypted using AWS KMS and returned to the Cloud9 environment to be decrypted.

Figure 9 – Final Decrypted Response

Clean up

Complete the following steps to clean up your resources:

1. Stop the EC2 instance created to house your enclave.
2. Delete the Cloud9 environment.
3. Delete the KMS key.
4. Remove the EC2 instance role and IAM user permissions.

Conclusion

In this post, we showcased how to use Nitro Enclaves to deploy an LLM question and answering service that more securely sends and receives PII and PHI information. This was deployed on Amazon EC2, and the enclaves are integrated with AWS KMS restricting access to a KMS key, so only Nitro Enclaves and the end-user are allowed to use the key and decrypt the question.

If you’re planning to scale this architecture to support larger workloads, make sure the model selection process matches your model requirements with EC2 resources. Additionally, you must consider the maximum request size and what impact that will have on the HTTP server and inference time against the model. Many of these parameters are customizable through the model and HTTP server settings.

The best way to determine the specific settings and requirements for your workload is through testing with a fine-tuned LLM. Although this post only included natural language processing of sensitive data, you can modify this architecture to support alternate LLMs supporting audio, computer vision, or multi-modalities. The same security principles highlighted here can be applied to data in any format. The resources used to build this post are available on the GitHub repo.

Share how you are going to adapt this solution for your environment in the comments section.

About the Authors

Justin Miles is a cloud engineer within the Leidos Digital Modernization Sector under the Office of Technology. In his spare time, he enjoys golfing and traveling.

Liv d’Aliberti is a researcher within the Leidos AI/ML Accelerator under the Office of Technology. Their research focuses on privacy-preserving machine learning.

Chris Renzo is a Sr. Solution Architect within the AWS Defense and Aerospace organization. Outside of work, he enjoys a balance of warm weather and traveling.

Joe Kovba is a Vice President within the Leidos Digital Modernization Sector. In his free time, he enjoys refereeing football games and playing softball.

Read More

VistaPrint, a Cimpress business, is the design and marketing partner to millions of small businesses around the world. For more than two decades, VistaPrint has empowered small businesses to quickly and effectively create the marketing products – from promotional materials and signage to print advertising and more – to get the job done, regardless of whether they operate in-store or online.

To support small businesses on their brand-building journey, VistaPrint provides customers with personalized product recommendations, both in real time on vistaprint.com and through marketing emails. These product recommendations improve their customers’ experience by making it more efficient to find the products they need, while increasing VistaPrint’s conversion rates. Since implementing Amazon Personalize, VistaPrint increased their conversion rate by 10 percent and reduced their total cost of ownership by 30 percent.

In this post, we show you how VistaPrint uses a combination of Amazon Personalize, Twilio Segment, and auxiliary AWS services and partner solutions to better understand their customers’ needs and provide personalized product recommendations.

Prior solution and challenges

Prior to their current solution, VistaPrint had an internally developed product recommendation system hosted on-premises. The first challenge with their prior solution was that the solution couldn’t scale automatically when demand increased. The second challenge was that changes to the in-house developed system were time-consuming, because a high degree of machine learning and ecommerce domain specialization was required to make modifications.

These challenges led to the decision to create a new cloud-native system that can scale with increased demand and consists of serverless and software as a service (SaaS) components that externalize much of the domain-specific functionality to allow for easier operations and faster time-to-market for changes.

The new VistaPrint personalized product recommendation system

Figure 1

As seen in Figure 1, the steps in how VistaPrint provides personalized product recommendations with their new cloud-native architecture are:

1. Aggregate historical data in a data warehouse. Data from upstream systems including customer data platforms (CDPs) like Twilio Segment, order management, product catalog, and user management systems are collected in a data warehouse, which in VistaPrint’s case is Snowflake.
2. Transform the data to create Amazon Personalize training data. Amazon Personalize uses data about users, items, and interactions, and this data is ingested from Amazon Simple Storage Service (Amazon S3) in CSV format. In VistaPrint’s case, they use Databricks to perform the required data transformations before landing the data in Amazon S3.
3. Import bulk historical data to train Amazon Personalize models. After bulk historical data is ingested into an Amazon Personalize dataset, one or more solutions are trained using this data. In VistaPrint’s case, they use the User-Personalization and Similar-Items model recipes.
   • With User-Personalization, Amazon Personalize predicts the items that a user will interact with based on previous interactions across all users.
   • With Similar-Items, Amazon Personalize generates recommendations for items that are similar to an item you specify.

   To maintain the relevance of the personalization models, steps 2 and 3 are repeated on a regular basis to keep the training data up to date.

4. Stream ecommerce website events to a CDP. A CDP is used to capture events from an ecommerce website, for example when a user views a product or adds a product to their shopping cart. A CDP can also perform identity resolution, which helps to identify the user regardless of whether they’re accessing a platform from a mobile or a web client. VistaPrint uses Twilio Segment as their CDP.
5. Generate real-time product recommendations as a customer navigates the ecommerce website. As a customer navigates an ecommerce website and these events are captured by a CDP, they are also forwarded to Amazon Personalize. Amazon Personalize in turn generates recommendations for additional products that a customer may be interested in. These recommendations are placed back into the ecommerce website experience in real-time.
   • AWS Lambda is used to send data from Segment to Amazon Personalize using Segment’s Amazon Lambda Destination. VistaPrint uses the Segment Amazon Lambda Destination to perform additional data transformations and to get flexibility to integrate with additional observability tooling not shown, but other AWS customers can consider Segment’s Amazon Personalize Destination which is suitable for simpler integrations.
   • VistaPrint created a personalization service that sits in front of Amazon Personalize. This service provides additional functionality on top of Amazon Personalize APIs, including the ability to cache recent recommendations in Amazon DynamoDB, and integration with VistaPrint’s authentication and authorization systems.
   • VistaPrint created a placement and offer engine (POE), which allows data scientists and marketers to collaborate. Placement templates are used to create customized placements by allowing a marketer to select an Amazon Personalize model, the visual style of the placement, and extra features like whether to display a customer’s logo as it would appear on the final manufactured product. Figure 2 shows an example of one of these placements, called More with your design, as seen on vistaprint.com.
6. Generate product recommendations as part of email marketing campaigns. In addition to providing real-time product recommendations on their website, VistaPrint uses personalized product recommendations in email marketing campaigns. The same POE system is used to design and place product recommendations into email templates.

Figure 2

Business Impact

Since implementing its new personalized product recommendation system, VistaPrint has realized a 10 percent increase in conversions originating from personalized recommendations. Amazon Personalize also reduced VistaPrint’s total cost of ownership by 30 percent compared to the previous on-premises solution.

Conclusion

VistaPrint’s cloud-native personalized product recommendation system helps the company deliver a more efficient and helpful experience to their customers, while increasing the company’s conversion rates.

Amazon Personalize is at the center of VistaPrint’s personalized product recommendation system, providing a fully managed, machine learning powered solution.

A customer data platform like Twilio Segment allows companies like VistaPrint to build a connected, 360 degree view of their customers by aggregating data from all of their customer touchpoints across multiple business domains. This cohesive view of the customer leads to more accurate and personalized product recommendations when paired with Amazon Personalize.

Next Steps

The VistaPrint personalized product recommendation system is one product within a larger data mesh of products. Read more about Vista’s data mesh strategy in this previous post How Vista built a data mesh enabled by solutions available in AWS Marketplace

Also read more on the other topics in this post:

• Personalize cross-channel customer experiences with Amazon SageMaker, Amazon Personalize, and Twilio Segment
• Retail Demo Store is an example retail web application on GitHub showing Amazon Personalize in action.

About the Authors

Ethan Fahy is an Enterprise Senior Solutions Architect at AWS based in Boston, MA. Ethan has a background in geophysics and enjoys building large-scale, cloud-native architectures to support scientific workloads.

Mouloud Lounaci leads the Engineering team for Marketing Optimization at Vista. He is a Machine Learning enthusiast with around 10 years of experience in building AI-powered software products to solve complex customer problems. Whenever he gets a chance, Mouloud jumps on a plane to discover cultures, food, and landscapes from around the world.

Emeline Escolivet is the Engineering Manager for the Recommendations team at Vista. With 10+ years of experience as a Software Engineer, she enjoys turning complex business issues into reliable software solutions. In her free time, she likes to describe herself as a hiker, dancer and food lover.

Vibhusheet Tripathi is a Senior Data Engineer in the Recommendations Team at Vista. When not experimenting with machine learning systems, Vibhu likes to read, play sports and listen to music.

Read More

Many customers, including those in creative advertising, media and entertainment, ecommerce, and fashion, often need to change the background in a large number of images. Typically, this involves manually editing each image with photo software. This can take a lot of effort, especially for large batches of images. However, Amazon Bedrock and AWS Step Functions make it straightforward to automate this process at scale.

Amazon Bedrock offers the generative AI foundation model Amazon Titan Image Generator G1, which can automatically change the background of an image using a technique called outpainting. Step Functions allows you to create an automated workflow that seamlessly connects with Amazon Bedrock and other AWS services. Together, Amazon Bedrock and Step Functions streamline the entire process of automatically changing backgrounds across multiple images.

This post introduces a solution that simplifies the process of changing backgrounds in multiple images. By harnessing the capabilities of generative AI with Amazon Bedrock and the Titan Image Generator G1 model, combined with Step Functions, this solution efficiently generates images with the desired background. This post provides insight into the inner workings of the solution and helps you understand the design choices made to build this own custom solution.

See the GitHub repository for detailed instructions on deploying this solution.

Solution overview

Let’s look at how the solution works at a high level before diving deeper into specific elements and the AWS services used. The following diagram provides a simplified view of the solution architecture and highlights the key elements.

The workflow consists of the following steps:

1. A user uploads multiple images into an Amazon Simple Storage Service (Amazon S3) bucket via a Streamlit web application.
2. The Streamlit web application calls an Amazon API Gateway REST API endpoint integrated with the Amazon Rekognition DetectLabels API, which detects labels for each image.
3. Upon submission, the Streamlit web application updates an Amazon DynamoDB table with image details.
4. The DynamoDB update triggers an AWS Lambda function, which starts a Step Functions workflow.
5. The Step Functions workflow runs the following steps for each image:
   5.1 Constructs a request payload for the Amazon Bedrock InvokeModel API.
   5.2 Invokes the Amazon Bedrock InvokeModel API action.
   5.3 Parses an image from the response and saves it to an S3 location.
   5.4 Updates the image status in a DynamoDB table.
6. The Step Functions workflow invokes a Lambda function to generate a status report.
7. The workflow sends an email using Amazon Simple Notification Service (Amazon SNS).

As shown in the following screenshot, the Streamlit web application allows you to upload images and enter text prompts to specify desired backgrounds, negative prompts, and outpainting mode for image generation. You can also view and remove unwanted labels associated with each uploaded image that you don’t want to keep in the final generated images.

In this example, the prompt for the background is “London city background.” The automation process generates new images based on the original uploaded images with London as the background.

Streamlit web application and images uploads

A Streamlit web application serves as the frontend for this solution. To protect the application from unauthorized access, it integrates with an Amazon Cognito user pool. API Gateway uses an Amazon Cognito authorizer to authenticate requests. The web application completes the following steps:

1. For each selected image, it retrieves labels via Amazon Rekognition using an API Gateway REST API endpoint.
2. Upon submission, the application uploads images to an S3 bucket.
3. The application updates a DynamoDB table with relevant parameters, image names, and associated labels for each image using another API Gateway REST API endpoint.

Image processing workflow

When the DynamoDB table is updated, DynamoDB Streams triggers a Lambda function to start a new Step Functions workflow. The following is a sample request for the workflow:

{
  "Id": "621fa85a-38bb-4d98-a656-93bbbcf5477f",
  "S3Bucket": "<Image Bucket>",
  "InputS3Prefix": "image-files/<year>/<month>/<day>/<timestamp>",
  "OutputS3Prefix": "generated-image-files/<year>/<month>/<day>/<timestamp>",
  "StatusS3Prefix": "status-report-files/<year>/<month>/<day>/<timestamp>",
  "Prompt": "london city background",
  "NegativePrompt": "low quality, low resolution",
  "Mode": "PRECISE",
  "Images": [
    {
      "ImageName": "bus.png",
      "Labels": "Bus, Person"
    },
    {
      "ImageName": "cop.png",
      "Labels": "Person, Adult, Male, Man, Helmet, Jacket"
    },
    {
      "ImageName": "iguana-2.png",
      "Labels": "Lizard”
    },
    {
      "ImageName": "dog.png",
      "Labels": "Dog"
    }
  ]
}

The Step Functions workflow subsequently performs the following three steps:

1. Replace the background for all images.
2. Generate a status report.
3. Send an email via Amazon SNS.

The following screenshot illustrates the Step Functions workflow.

Let’s look at each step in more detail.

Replace background for all images

Step Functions uses a Distributed Map to process each image in parallel child workflows. The Distributed Map allows high-concurrency processing. Each child workflow has its own separate run history from that of the parent workflow.

Step Functions uses an InvokeModel optimized API action for Amazon Bedrock. The API accepts requests and responses that are up to 25 MB. However, Step Functions has a 256 KB limit on state payload input and output. To support larger images, the solution uses an S3 bucket where the InvokeModel API reads data from and writes the result to. The following is the configuration for the InvokeModel API for Amazon Bedrock integration:

{
    "ModelId": "arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-image-generator-v1",
    "ContentType": "application/json",
    "Input": {  
        "S3Uri": “s3://<Image Bucket>/image-files/<year>/<month>/<day>/<timestamp>/<Image name>.json",
    },  
    "Output": {  
        "S3Uri": “s3://<Image Bucket>/generated-image-files/<year>/<month>/<day>/<timestamp>/<Image name>.json”
    } 
}

The Input S3Uri parameter specifies the source location to retrieve the input data. The Output S3Uri parameter specifies the destination to write the API response.

A Lambda function saves the request payload as a JSON file in the specified Input S3Uri location. The InvokeModel API uses this input payload to generate images with the specified background:

{
    "taskType": "OUTPAINTING",
    "outPaintingParams": {
        "text": "london city background",
        "negativeText": "low quality, low resolution",        
        "image": "<base64-encoded string>",                         
        "maskPrompt": "Bus",                      
        "maskImage": "base64-encoded string",                             
        "outPaintingMode": "DEFAULT | PRECISE"                 
    },                                                 
    "imageGenerationConfig": {
        "numberOfImages": 1,
        "quality": "premium",
        "height": 1024,
        "width": 1024,
        "cfgScale": 8.0
    }
}

The Titan Image Generator G1 model supports the following parameters for image generation:

• taskType – Specifies the outpainting method to replace background of image.
• text – A text prompt to define the background.
• negativeText – A text prompt to define what not to include in the image.
• maskPrompt – A text prompt that defines the mask. It corresponds to labels that you want to retain in the final generated images.
• maskImage – The JPEG or PNG image encoded in base64.
• outPaintingMode – Specifies whether to allow modification of the pixels inside the mask or not. DEFAULT allows modification of the image inside the mask in order to keep it consistent with the reconstructed background. PRECISE prevents modification of the image inside the mask.
• numberOfImages – The number of images to generate.
• quality – The quality of the generated images: standard or premium.
• cfgScale – Specifies how strongly the generated image should adhere to the prompt.
• height – The height of the image in pixels.
• width – The width of the image in pixels.

The Amazon Bedrock InvokeModel API generates a response with an encoded image in the Output S3Uri location. Another Lambda function parses the image from the response, decodes it from base64, and saves the image file in the following location: s3://<Image Bucket>/generated-image-file/<year>/<month>/<day>/<timestamp>/.

Finally, a child workflow updates a DynamoDB table with image generation status, marking it as either Succeeded or Failed, and including details such as ImageName, Cause, Error, and Status.

Generate a status report

After the image generation process, a Lambda function retrieves the status details from DynamoDB. It dynamically compiles these details into a comprehensive status report in JSON format. It then saves the generated status report a JSON file in the following location: s3://<Image Bucket>/status-report-files/<year>/<month>/<day>/<timestamp>/. The ITOps team can integrate this report with their existing notification system to track if image processing completed successfully. For business users, you can expand this further to generate a report in CSV format.

Send an email via Amazon SNS

Step Functions invokes an Amazon SNS API action to send an email. The email contains details including the S3 location for the status report and final images files. The following is the sample notification email.

Conclusion

In this post, we provided an overview of a sample solution demonstrating the automation of changing image backgrounds at scale using Amazon Bedrock and Step Functions. We also explained each element of the solution in detail. By using the Step Functions optimized integration with Amazon Bedrock, Distributed Map, and the Titan Image Generator G1 model, the solution efficiently replaces the backgrounds of images in parallel, enhancing productivity and scalability.

To deploy the solution, refer to the instructions in the GitHub repository.

Resources

To learn more about Amazon Bedrock, see the following resources:

• Using generative AI on AWS for diverse content types Workshop
• Amazon Bedrock Workshop
• Amazon Bedrock User Guide
• Amazon Bedrock with Amazon Titan Image Generator G1
• Amazon Bedrock InvokeModel API

To learn more about the Titan Image Generator G1 model, see the following resources:

• Amazon Titan Image Generator G1 model
• Amazon Titan Image Generator Demo

To learn more about using Amazon Bedrock with Step Functions, see the following resources:

• Build generative AI apps using AWS Step Functions and Amazon Bedrock
• Amazon Bedrock Optimized Integration for Step Functions

About the Author

Chetan Makvana is a Senior Solutions Architect with Amazon Web Services. He works with AWS partners and customers to provide them with architectural guidance for building scalable architecture and implementing strategies to drive adoption of AWS services. He is a technology enthusiast and a builder with a core area of interest on generative AI, serverless, and DevOps. Outside of work, he enjoys watching shows, traveling, and music. 

Read More

In this post, we demonstrate how to efficiently fine-tune a state-of-the-art protein language model (pLM) to predict protein subcellular localization using Amazon SageMaker.

Proteins are the molecular machines of the body, responsible for everything from moving your muscles to responding to infections. Despite this variety, all proteins are made of repeating chains of molecules called amino acids. The human genome encodes 20 standard amino acids, each with a slightly different chemical structure. These can be represented by letters of the alphabet, which then allows us to analyze and explore proteins as a text string. The enormous possible number of protein sequences and structures is what gives proteins their wide variety of uses.

Proteins also play a key role in drug development, as potential targets but also as therapeutics. As shown in the following table, many of the top-selling drugs in 2022 were either proteins (especially antibodies) or other molecules like mRNA translated into proteins in the body. Because of this, many life science researchers need to answer questions about proteins faster, cheaper, and more accurately.

Data source: Urquhart, L. Top companies and drugs by sales in 2022. Nature Reviews Drug Discovery 22, 260–260 (2023).

Because we can represent proteins as sequences of characters, we can analyze them using techniques originally developed for written language. This includes large language models (LLMs) pretrained on huge datasets, which can then be adapted for specific tasks, like text summarization or chatbots. Similarly, pLMs are pre-trained on large protein sequence databases using unlabeled, self-supervised learning. We can adapt them to predict things like the 3D structure of a protein or how it may interact with other molecules. Researchers have even used pLMs to design novel proteins from scratch. These tools don’t replace human scientific expertise, but they have the potential to speed up pre-clinical development and trial design.

One challenge with these models is their size. Both LLMs and pLMs have grown by orders of magnitude in the past few years, as illustrated in the following figure. This means that it can take a long time to train them to sufficient accuracy. It also means that you need to use hardware, especially GPUs, with large amounts of memory to store the model parameters.

Long training times, plus large instances, equals high cost, which can put this work out of reach for many researchers. For example, in 2023, a research team described training a 100 billion-parameter pLM on 768 A100 GPUs for 164 days! Fortunately, in many cases we can save time and resources by adapting an existing pLM to our specific task. This technique is called fine-tuning, and also allows us to borrow advanced tools from other types of language modeling.

Solution overview

The specific problem we address in this post is subcellular localization: Given a protein sequence, can we build a model that can predict if it lives on the outside (cell membrane) or inside of a cell? This is an important piece of information that can help us understand the function and whether it would make a good drug target.

We start by downloading a public dataset using Amazon SageMaker Studio. Then we use SageMaker to fine-tune the ESM-2 protein language model using an efficient training method. Finally, we deploy the model as a real-time inference endpoint and use it to test some known proteins. The following diagram illustrates this workflow.

In the following sections, we go through the steps to prepare your training data, create a training script, and run a SageMaker training job. All of the code featured in this post is available on GitHub.

Prepare the training data

We use part of the DeepLoc-2 dataset, which contains several thousand SwissProt proteins with experimentally determined locations. We filter for high-quality sequences between 100–512 amino acids:

df = pd.read_csv(
    "https://services.healthtech.dtu.dk/services/DeepLoc-2.0/data/Swissprot_Train_Validation_dataset.csv"
).drop(["Unnamed: 0", "Partition"], axis=1)
df["Membrane"] = df["Membrane"].astype("int32")

# filter for sequences between 100 and 512 amino acides
df = df[df["Sequence"].apply(lambda x: len(x)).between(100, 512)]

# Remove unnecessary features
df = df[["Sequence", "Kingdom", "Membrane"]]

Next, we tokenize the sequences and split them into training and evaluation sets:

dataset = Dataset.from_pandas(df).train_test_split(test_size=0.2, shuffle=True)
tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")

def preprocess_data(examples, max_length=512):
    text = examples["Sequence"]
    encoding = tokenizer(text, truncation=True, max_length=max_length)
    encoding["labels"] = examples["Membrane"]
    return encoding

encoded_dataset = dataset.map(
    preprocess_data,
    batched=True,
    num_proc=os.cpu_count(),
    remove_columns=dataset["train"].column_names,
)

encoded_dataset.set_format("torch")

Finally, we upload the processed training and evaluation data to Amazon Simple Storage Service (Amazon S3):

train_s3_uri = S3_PATH + "/data/train"
test_s3_uri = S3_PATH + "/data/test"

encoded_dataset["train"].save_to_disk(train_s3_uri)
encoded_dataset["test"].save_to_disk(test_s3_uri)

Create a training script

SageMaker script mode allows you to run your custom training code in optimized machine learning (ML) framework containers managed by AWS. For this example, we adapt an existing script for text classification from Hugging Face. This allows us to try several methods for improving the efficiency of our training job.

Method 1: Weighted training class

Like many biological datasets, the DeepLoc data is unevenly distributed, meaning there isn’t an equal number of membrane and non-membrane proteins. We could resample our data and discard records from the majority class. However, this would reduce the total training data and potentially hurt our accuracy. Instead, we calculate the class weights during the training job and use them to adjust the loss.

In our training script, we subclass the Trainer class from transformers with a WeightedTrainer class that takes class weights into account when calculating cross-entropy loss. This helps prevent bias in our model:

class WeightedTrainer(Trainer):
    def __init__(self, class_weights, *args, **kwargs):
        self.class_weights = class_weights
        super().__init__(*args, **kwargs)

    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.get("logits")
        loss_fct = torch.nn.CrossEntropyLoss(
            weight=torch.tensor(self.class_weights, device=model.device)
        )
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss

Method 2: Gradient accumulation

Gradient accumulation is a training technique that allows models to simulate training on larger batch sizes. Typically, the batch size (the number of samples used to calculate the gradient in one training step) is limited by the GPU memory capacity. With gradient accumulation, the model calculates gradients on smaller batches first. Then, instead of updating the model weights right away, the gradients get accumulated over multiple small batches. When the accumulated gradients equal the target larger batch size, the optimization step is performed to update the model. This lets models train with effectively bigger batches without exceeding the GPU memory limit.

However, extra computation is needed for the smaller batch forward and backward passes. Increased batch sizes via gradient accumulation can slow down training, especially if too many accumulation steps are used. The aim is to maximize GPU usage but avoid excessive slowdowns from too many extra gradient computation steps.

Method 3: Gradient checkpointing

Gradient checkpointing is a technique that reduces the memory needed during training while keeping the computational time reasonable. Large neural networks take up a lot of memory because they have to store all the intermediate values from the forward pass in order to calculate the gradients during the backward pass. This can cause memory issues. One solution is to not store these intermediate values, but then they have to be recalculated during the backward pass, which takes a lot of time.

Gradient checkpointing provides a balanced approach. It saves only some of the intermediate values, called checkpoints, and recalculates the others as needed. Therefore, it uses less memory than storing everything, but also less computation than recalculating everything. By strategically selecting which activations to checkpoint, gradient checkpointing enables large neural networks to be trained with manageable memory usage and computation time. This important technique makes it feasible to train very large models that would otherwise run into memory limitations.

In our training script, we turn on gradient activation and checkpointing by adding the necessary parameters to the TrainingArguments object:

from transformers import TrainingArguments

training_args = TrainingArguments(
	gradient_accumulation_steps=4,
	gradient_checkpointing=True
)

Method 4: Low-Rank Adaptation of LLMs

Large language models like ESM-2 can contain billions of parameters that are expensive to train and run. Researchers developed a training method called Low-Rank Adaptation (LoRA) to make fine-tuning these huge models more efficient.

The key idea behind LoRA is that when fine-tuning a model for a specific task, you don’t need to update all the original parameters. Instead, LoRA adds new smaller matrices to the model that transform the inputs and outputs. Only these smaller matrices are updated during fine-tuning, which is much faster and uses less memory. The original model parameters stay frozen.

After fine-tuning with LoRA, you can merge the small adapted matrices back into the original model. Or you can keep them separate if you want to quickly fine-tune the model for other tasks without forgetting previous ones. Overall, LoRA allows LLMs to be efficiently adapted to new tasks at a fraction of the usual cost.

In our training script, we configure LoRA using the PEFT library from Hugging Face:

from peft import get_peft_model, LoraConfig, TaskType
import torch
from transformers import EsmForSequenceClassification

model = EsmForSequenceClassification.from_pretrained(
	“facebook/esm2_t33_650M_UR50D”,
	Torch_dtype=torch.bfloat16,
	Num_labels=2,
)

peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    bias="none",
    r=8,
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules=[
        "query",
        "key",
        "value",
        "EsmSelfOutput.dense",
        "EsmIntermediate.dense",
        "EsmOutput.dense",
        "EsmContactPredictionHead.regression",
        "EsmClassificationHead.dense",
        "EsmClassificationHead.out_proj",
    ]
)

model = get_peft_model(model, peft_config)

Submit a SageMaker training job

After you have defined your training script, you can configure and submit a SageMaker training job. First, specify the hyperparameters:

hyperparameters = {
    "model_id": "facebook/esm2_t33_650M_UR50D",
    "epochs": 1,
    "per_device_train_batch_size": 8,
    "gradient_accumulation_steps": 4,
    "use_gradient_checkpointing": True,
    "lora": True,
}

Next, define what metrics to capture from the training logs:

metric_definitions = [
    {"Name": "epoch", "Regex": "'epoch': ([0-9.]*)"},
    {
        "Name": "max_gpu_mem",
        "Regex": "Max GPU memory use during training: ([0-9.e-]*) MB",
    },
    {"Name": "train_loss", "Regex": "'loss': ([0-9.e-]*)"},
    {
        "Name": "train_samples_per_second",
        "Regex": "'train_samples_per_second': ([0-9.e-]*)",
    },
    {"Name": "eval_loss", "Regex": "'eval_loss': ([0-9.e-]*)"},
    {"Name": "eval_accuracy", "Regex": "'eval_accuracy': ([0-9.e-]*)"},
]

Finally, define a Hugging Face estimator and submit it for training on an ml.g5.2xlarge instance type. This is a cost-effective instance type that is widely available in many AWS Regions:

from sagemaker.experiments.run import Run
from sagemaker.huggingface import HuggingFace
from sagemaker.inputs import TrainingInput

hf_estimator = HuggingFace(
    base_job_name="esm-2-membrane-ft",
    entry_point="lora-train.py",
    source_dir="scripts",
    instance_type="ml.g5.2xlarge",
    instance_count=1,
    transformers_version="4.28",
    pytorch_version="2.0",
    py_version="py310",
    output_path=f"{S3_PATH}/output",
    role=sagemaker_execution_role,
    hyperparameters=hyperparameters,
    metric_definitions=metric_definitions,
    checkpoint_local_path="/opt/ml/checkpoints",
    sagemaker_session=sagemaker_session,
    keep_alive_period_in_seconds=3600,
    tags=[{"Key": "project", "Value": "esm-fine-tuning"}],
)

with Run(
    experiment_name=EXPERIMENT_NAME,
    sagemaker_session=sagemaker_session,
) as run:
    hf_estimator.fit(
        {
            "train": TrainingInput(s3_data=train_s3_uri),
            "test": TrainingInput(s3_data=test_s3_uri),
        }
    )

The following table compares the different training methods we discussed and their effect on the runtime, accuracy, and GPU memory requirements of our job.

All of the methods produced models with high evaluation accuracy. Using LoRA and gradient activation decreased the runtime (and cost) by 18% and 25%, respectively. Using gradient checkpointing decreased the maximum GPU memory usage by 55%. Depending on your constraints (cost, time, hardware), one of these approaches may make more sense than another.

Each of these methods perform well by themselves, but what happens when we use them in combination? The following table summarizes the results.

In this case, we see a 12% reduction in accuracy. However, we’ve reduced the runtime by 57% and GPU memory use by 85%! This is a massive decrease that allows us to train on a wide range of cost-effective instance types.

Clean up

If you’re following along in your own AWS account, delete the any real-time inference endpoints and data you created to avoid further charges.

predictor.delete_endpoint()

bucket = boto_session.resource("s3").Bucket(S3_BUCKET)
bucket.objects.filter(Prefix=S3_PREFIX).delete()

Conclusion

In this post, we demonstrated how to efficiently fine-tune protein language models like ESM-2 for a scientifically relevant task. For more information about using the Transformers and PEFT libraries to train pLMS, check out the posts Deep Learning With Proteins and ESMBind (ESMB): Low Rank Adaptation of ESM-2 for Protein Binding Site Prediction on the Hugging Face blog. You can also find more examples of using machine learning to predict protein properties in the Awesome Protein Analysis on AWS GitHub repository.

About the Author

Brian Loyal is a Senior AI/ML Solutions Architect in the Global Healthcare and Life Sciences team at Amazon Web Services. He has more than 17 years’ experience in biotechnology and machine learning, and is passionate about helping customers solve genomic and proteomic challenges. In his spare time, he enjoys cooking and eating with his friends and family.

Read More