sagemaker_serve_huggingface_model_local_test_and_endpoint_deploy.py

1import sagemaker
2from sagemaker_serve import SageMakerModel, InferenceSpec
3from transformers import pipeline
4
5# Define your inference logic
6class MyInferenceSpec(InferenceSpec):
7    def load(self, model_dir):
8        # Load model using Hugging Face pipeline
9        return pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
10
11    def predict(self, model, data):
12        # Run inference
13        return model(data)
14
15# Initialize the SageMakerModel with the custom InferenceSpec
16model = SageMakerModel(
17    inference_spec=MyInferenceSpec(),
18    model_description="My Sentiment Analysis Model"
19)
20
21# Test the model locally
22local_prediction = model.predict("I love using SageMaker Serve!")
23print(f"Local Prediction: {local_prediction}")
24
25# Deploy the model to an AWS SageMaker Endpoint
26predictor = model.deploy(
27    instance_type="ml.m5.xlarge",
28    initial_instance_count=1
29)
30
31# Run a prediction against the remote endpoint
32remote_prediction = predictor.predict("This is an amazing tool for deployment.")
33print(f"Remote Prediction: {remote_prediction}")
34
35# Clean up
36predictor.delete_endpoint()