Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add example for batch inference. #149

Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
261 changes: 261 additions & 0 deletions 01_Generation/03_batch_generate.ipynb
Original file line number Diff line number Diff line change
@@ -0,0 +1,261 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "569bad05-d0cb-4460-9645-42bd1967bf71",
"metadata": {},
"source": [
"### Setting up the Bedrock SDK\n",
"\n",
"This section of code sets up the Bedrock SDK by:\n",
"\n",
"- Downloading and unzipping the preview Python SDK package from a cloudfront URL\n",
"- Installing the preview botocore and boto3 wheel packages from the extracted SDK directory \n",
"- Importing necessary libraries like boto3, botocore, and JSON utilities\n",
"\n",
"The Bedrock SDK provides an API client to interact with AWS Bedrock services for batch inference jobs. We set up the client and packages here to prepare for defining and invoking batch jobs later on.\n",
"\n",
"Some things to note:\n",
"\n",
"- The `!` prefix runs these as shell commands in the jupyter notebook \n",
"- We quietly install `-q` the wheel packages \n",
"- The preview SDK gives access to Bedrock services before general availability\n",
"- We'll use boto3 built on botocore to call Bedrock batch APIs with Python"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "157a66e0-d235-4158-851f-ed5f82563cef",
"metadata": {},
"outputs": [],
"source": [
"# Download and unzip preview SDK\n",
"!wget https://d2eo22ngex1n9g.cloudfront.net/Documentation/SDK/bedrock-python-sdk-reinvent.zip\n",
"!unzip -qq ./bedrock-python-sdk-reinvent.zip -d ./bedrock-python-sdk-reinvent\n",
"\n",
"# Install preview SDK packages\n",
"!pip install -q $(ls ./bedrock-python-sdk-reinvent/botocore-*.whl | head -1)\n",
"!pip install -q $(ls ./bedrock-python-sdk-reinvent/boto3-*.whl | head -1)"
]
},
{
"cell_type": "markdown",
"id": "91d1f9d2-d0a4-4d85-a8ce-011725d8e167",
"metadata": {},
"source": [
"### Import Python modules \n",
"\n",
"This code chunk imports Python libraries we'll need for using Bedrock and monitoring batch jobs:\n",
"\n",
"- `boto3` - AWS SDK library to call Bedrock services \n",
"- `botocore` - Underlying AWS SDK library boto3 builds on\n",
"- `Config` - Botocore config object for SDK behavior\n",
"- `time` - For sleep delays in loops monitoring job status\n",
"- `json` - For serializing data to JSON for Bedrock API input/output \n",
"- `uuid` - For generating unique job names \n",
"- `display`, `clear_output` - IPython display functions to refresh notebook output\n",
"\n",
"Key points:\n",
"\n",
"- We import AWS SDK clients and configs to call Bedrock batch APIs\n",
"- JSON helps pre-process data for batch input and parse results \n",
"- Time and IPython functions help display feedback on async job status\n",
"- We'll generate unique batch job names using the uuid module"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "14b56da1-8483-48f1-aa78-ef49a3efc561",
"metadata": {},
"outputs": [],
"source": [
"import boto3\n",
"import botocore\n",
"from botocore.client import Config\n",
"\n",
"import time\n",
"import json\n",
"import uuid\n",
"\n",
"from IPython.display import display, clear_output"
]
},
{
"cell_type": "markdown",
"id": "9219ee97-fe4e-4bba-8280-bb243f87a393",
"metadata": {},
"source": [
"### Upload input data to S3\n",
"\n",
"This code handles uploading our input data to an S3 bucket:\n",
"\n",
"- An S3 client is created using the boto3 SDK\n",
"- The S3 bucket and input/output keys are defined\n",
"- Sample input data is prepared in a list of dictionaries \n",
"- The data contains prompts and parameters for the Claude model\n",
"- The data is looped through and converted to JSON lines format\n",
"- The JSON lines output is uploaded to the S3 input key \n",
"\n",
"Key points:\n",
"\n",
"- Input data for Batch needs to be formatted and stored in S3\n",
"- Here we take care of uploading an input.jsonl file \n",
"- The file contains JSON records, one prompt/parameter input per line\n",
"- This will be consumed by the Bedrock batch job as model input "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3621d169-6692-44a4-97ca-d51539adb447",
"metadata": {},
"outputs": [],
"source": [
"# Set up S3 client\n",
"session = boto3.Session()\n",
"s3 = session.client('s3')\n",
"\n",
"# Bucket for input/output data \n",
"bucket = \"YOUR_BUCKET_NAME_HERE\"\n",
"input_key = \"input.jsonl\" \n",
"output_key = \"output/\"\n",
"\n",
"# Upload sample input data\n",
"data = [\n",
" {'recordId': '1', 'modelInput': {'prompt': 'Human: Write an email from Bob, Customer Service Manager, to the customer \"John Doe\" who provided negative feedback on the service provided by our customer support engineer Assistant:', 'max_tokens_to_sample':300}},\n",
"\n",
" {'recordId': '2', 'modelInput': {'prompt': 'Human: Craft an email response from Amanda, Sales Director, to a customer inquiry about pricing options for a large order Assistant:', 'max_tokens_to_sample':300}},\n",
"\n",
" {'recordId': '3', 'modelInput': {'prompt': 'Human: Write an email from the HR department to a new employee Daniel with orientation and onboarding information Assistant:', 'max_tokens_to_sample':300}},\n",
"\n",
" {'recordId': '4', 'modelInput': {'prompt': 'Human: Compose an email from IT support to employees about a planned network upgrade and potential service disruptions Assistant:', 'max_tokens_to_sample':300}},\n",
"\n",
" {'recordId': '5', 'modelInput': {'prompt': 'Human: Generate an email newsletter from a company CEO to employees highlighting recent company news and achievements Assistant:', 'max_tokens_to_sample':300}},\n",
"\n",
" {'recordId': '6', 'modelInput': {'prompt': 'Human: Draft an email from a manager to her team announcing an upcoming office holiday party with event details Assistant:', 'max_tokens_to_sample':300}},\n",
"\n",
" {'recordId': '7', 'modelInput': {'prompt': 'Human: Create an email template for the marketing team to use when sending lead nurturing campaigns to potential customers Assistant:', 'max_tokens_to_sample':300}},\n",
"\n",
" {'recordId': '8', 'modelInput': {'prompt': 'Human: Compose an out of office auto-reply email for John Smith, Accounting Controller, who will be on vacation next week Assistant:', 'max_tokens_to_sample':300}},\n",
"\n",
" {'recordId': '9', 'modelInput': {'prompt': 'Human: Write a follow up email from Amanda to existing customers announcing a special limited-time sale on products Assistant:', 'max_tokens_to_sample':300}},\n",
"\n",
" {'recordId': '10', 'modelInput': {'prompt': 'Human: Generate an email from the support team to users regarding an application outage that occurred yesterday Assistant:', 'max_tokens_to_sample':300}}\n",
"]\n",
"\n",
"# Process data and output to new lines\n",
"output_data = \"\"\n",
"for row in data:\n",
" output_data += json.dumps(row) + \"\\n\"\n",
"\n",
"s3.put_object(Body=output_data, Bucket=bucket, Key=input_key)"
]
},
{
"cell_type": "markdown",
"id": "bc37cbdf-41a3-4eba-820a-99f95f4c7dd5",
"metadata": {},
"source": [
"### Create and monitor batch job\n",
"\n",
"This code handles creating a Bedrock batch job and monitoring its status:\n",
"\n",
"- A Bedrock SDK client is created\n",
"- Input and output config point to the S3 buckets/keys \n",
"- A unique job name is generated using the UUID library\n",
"- The batch job is created using the bedrock.create_job API\n",
"- Parameters include the IAM role, Claude model ID, and config\n",
"- The returned job ARN is stored to track job status \n",
"- A loop checks the status every 5 seconds using the ARN\n",
"- IPython display functions clear and update the notebook output\n",
"- Once the job shows complete status, the output key is built\n",
"- Output JSON lines data is downloaded from S3 \n",
"- The output is printed line by line as formatted JSON\n",
"\n",
"Key points:\n",
"\n",
"- Bedrock APIs are used to submit and monitor async batch jobs \n",
"- A loop tracks status by periodically checking based on the job ARN\n",
"- Result data is stored back in the designated S3 output location\n",
"- We retrieve the output and print the finished JSON lines result"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c05cd717-fdb1-4ebb-bf23-b0486b1aeb04",
"metadata": {},
"outputs": [],
"source": [
"# Create Bedrock client \n",
"bedrock = boto3.client('bedrock')\n",
"\n",
"# Input and output config\n",
"input_config = {\"s3InputDataConfig\":{\"s3Uri\": f\"s3://{bucket}/{input_key}\"}} \n",
"output_config = {\"s3OutputDataConfig\":{\"s3Uri\": f\"s3://{bucket}/{output_key}\"}}\n",
"\n",
"#Let's create a unique job name\n",
"suffix = str(uuid.uuid4())[:8] \n",
"job_name = f\"claude-batch-test-{suffix}\" \n",
"\n",
"# Create batch job\n",
"response = bedrock.create_model_invocation_job(\n",
" roleArn=\"arn:aws:iam::450006286986:role/BedrockBatchInferenceRole\",\n",
" modelId=\"anthropic.claude-v2\",\n",
" jobName=job_name,\n",
" inputDataConfig=input_config,\n",
" outputDataConfig=output_config \n",
")\n",
"\n",
"job_arn = response['jobArn']\n",
"\n",
"# Monitor job status\n",
"status = 'Submitted'\n",
"while status == 'Submitted' or status == 'InProgress':\n",
" result = bedrock.get_model_invocation_job(jobIdentifier=job_arn)\n",
" status = result['status']\n",
" \n",
" clear_output(wait=True)\n",
" display(f\"Status: {status}\")\n",
" \n",
" time.sleep(5)\n",
" \n",
"print(\"Job completed\")\n",
"\n",
"#Get the output manifest result key\n",
"result_key = output_key + result['jobArn'].split('/')[1] + '/' + input_key.split('.')[0] + '.jsonl.out'\n",
"\n",
"# Get output data\n",
"obj = s3.get_object(Bucket=bucket, Key=result_key)\n",
"content = obj['Body'].read().decode('utf-8').strip().split('\\n')\n",
"\n",
"clear_output(wait=True)\n",
"\n",
"for line in content:\n",
" print(json.dumps(json.loads(line), indent=4))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "conda_python3",
"language": "python",
"name": "conda_python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
3 changes: 2 additions & 1 deletion 01_Generation/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ We will first generate text using a zero-shot prompt. The zero-shot prompt provi

## Audience

Architects and developer who want to learn how to use Amazon Bedrock LLMs to generate text.
Architects and developers who want to learn how to use Amazon Bedrock LLMs to generate text.
Some of the business use cases for text generation include:

- Generating product descriptions based on product features and benefits for marketing teams
Expand All @@ -21,6 +21,7 @@ We will generate an email response to a customer where the customer had provided
1. [Generate Email with Amazon Titan](./00_generate_w_bedrock.ipynb) - Invokes Amazon Titan large text model using Bedrock API to generate an email response to a customer. It uses a zero-shot prompt without context as instruction to the model.
2. [Zero-shot Text Generation with Anthropic Claude](01_zero_shot_generation.ipynb) - Invokes Anthropic's Claude Text model using the LangChain framework integration with Bedrock to generate an email to a customer. It uses a zero-shot prompt without context as instruction to the model.
3. [Contextual Text Generation using LangChain](./02_contextual_generation.ipynb) - We provide additional context in the prompt which includes the original email from the customer that we would like the model to generate a response for. The example includes a custom prompt template in LangChain, so that variable values can be substitued in the prompt at runtime.
4. [Batch Text Generation](./03_batch_generate.ipynb) - Submits a Bedrock batch job to generate multiple emails in parallel based on input prompts and parameters. The notebook handles uploading input data to S3, submitting and monitoring the batch job, and retrieving output results.

## Setup
Before running any of the labs in this section ensure you've run the [Bedrock boto3 setup notebook](../00_Intro/bedrock_boto3_setup.ipynb#Prerequisites).
Expand Down