Skip to content

Instantly share code, notes, and snippets.

@mrmaheshrajput

mrmaheshrajput/SentenceTransformers SageMaker.py

Created May 12, 2024 15:15
Show Gist options
  • Save mrmaheshrajput/4fd67b2ef88bb3870fc746ecec17978a to your computer and use it in GitHub Desktop.
Save mrmaheshrajput/4fd67b2ef88bb3870fc746ecec17978a to your computer and use it in GitHub Desktop.
Download ZIP
Raw
SentenceTransformers SageMaker.py
from sagemaker.huggingface.model import HuggingFaceModel
ENDPOINT_NAME = "sbert-embeddings-minilml6" # Change this as desired
role = "" # SageMaker execution role ARN
hub = {
"HF_MODEL_ID": "sentence-transformers/all-MiniLM-L6-v2", # Change to your model
"HF_TASK": "feature-extraction"
}
# Deploy Hugging Face Model
predictor = HuggingFaceModel(
env=hub, # configuration for loading model from Hub
role=role, # iam role with permissions to create an Endpoint
transformers_version="4.26",
pytorch_version="1.13",
py_version="py39",
).deploy(
initial_instance_count=1,
instance_type="ml.m5.xlarge", # You can go for smaller instances too
endpoint_name=ENDPOINT_NAME,
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment