from sagemaker.huggingface import HuggingFaceModel
import sagemaker

role = sagemaker.get_execution_role()

# Replace this with your merged model's HF repo name
hub_model = "SamuelJaja/llama-3.1-8b-construction-merged"

# Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    env={
        'HF_MODEL_ID': hub_model,
        'HF_TASK': 'text-generation',
        'TRANSFORMERS_CACHE': '/tmp'  # helps avoid cache storage issues
    },
    role=role,
    transformers_version="4.39.1",
    pytorch_version="2.1.2",
    py_version="py310",
    image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference:2.1.2-transformers4.39.1-gpu-py310-cu121",
    instance_type="ml.g5.2xlarge",
)

# Deploy the model
predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.2xlarge"
)