sagemaker_serve_huggingface_model_endpoint_deployment.py

1import sagemaker
2from sagemaker_serve import SageMakerServe
3
4# Initialize SageMaker session
5session = sagemaker.Session()
6role = sagemaker.get_execution_role()
7
8# Define model details
9model_id = "sentence-transformers/all-MiniLM-L6-v2"
10model_name = "all-minilm-l6-v2-endpoint"
11
12# Initialize SageMakerServe
13ss = SageMakerServe(
14    model_id=model_id,
15    role=role,
16    session=session
17)
18
19# Deploy the model
20# This handles the container selection, model packaging, and endpoint creation
21predictor = ss.deploy(
22    endpoint_name=model_name,
23    instance_type="ml.m5.xlarge",
24    initial_instance_count=1
25)
26
27# Run inference
28input_data = {"inputs": "This is an example sentence"}
29prediction = predictor.predict(input_data)
30
31print(f"Prediction: {prediction}")
32
33# Clean up
34# ss.delete_endpoint()