sagemaker_modelbuilder_huggingface_local_test_and_endpoint_deploy.py

1import sagemaker
2from sagemaker.serve.builder.model_builder import ModelBuilder
3from sagemaker.serve.builder.schema_builder import SchemaBuilder
4from transformers import pipeline
5
6# 1. Setup: Define your model (e.g., a Hugging Face sentiment analysis pipeline)
7model = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
8
9# 2. Define Schema: Provide sample input and output to define the model interface
10sample_input = "I love using SageMaker!"
11sample_output = model(sample_input)
12schema_builder = SchemaBuilder(sample_input, sample_output)
13
14# 3. Build Model: Create the ModelBuilder object
15# This automatically handles container selection and serialization
16model_builder = ModelBuilder(
17    model=model,
18    schema_builder=schema_builder,
19    role_arn=sagemaker.get_execution_role(),
20)
21
22# 4. Local Testing: Verify the model works in a local container environment
23local_model = model_builder.build_for_local()
24local_model.create()
25print(f"Local prediction: {local_model.predict(sample_input)}")
26
27# 5. Deploy: Deploy the model to a production SageMaker Endpoint
28predictor = model_builder.deploy(initial_instance_count=1, instance_type="ml.m5.xlarge")
29
30# 6. Predict: Run inference against the live endpoint
31print(f"Endpoint prediction: {predictor.predict(sample_input)}")
32
33# Clean up
34# local_model.delete_model()
35# predictor.delete_endpoint()