# Deploy a gated JumpStart model using the SageMaker Python SDK v3.
# This model requires accepting the end-user license agreement (EULA).

from sagemaker.serve.model_builder import ModelBuilder
from sagemaker.core.jumpstart.configs import JumpStartConfig
from sagemaker.core.training.configs import Compute

# You must manually accept the end-user license agreement (EULA) to use this model.
accept_eula = False

model_builder = ModelBuilder.from_jumpstart_config(
    jumpstart_config=JumpStartConfig(model_id="meta-textgeneration-llama-3-1-70b-instruct", accept_eula=accept_eula),
    compute=Compute(instance_type="ml.p5en.48xlarge"),
)
model_builder.build()
endpoint = model_builder.deploy()

# Fine-tune a gated Hugging Face Hub model using the SageMaker Python SDK v3.
# This model requires a Hugging Face token to download weights.
# See https://github.com/huggingface/transformers/tree/v4.56.2/examples/pytorch

import boto3
from sagemaker.train.model_trainer import ModelTrainer
from sagemaker.core.training.configs import Compute, SourceCode
from sagemaker.core import image_uris

training_image = image_uris.retrieve(
    framework="huggingface",
    region=boto3.Session().region_name,
    version="4.56.2",
    base_framework_version="pytorch2.8.0",
    py_version="py312",
    image_scope="training",
)

trainer = ModelTrainer(
    training_image=training_image,
    source_code=SourceCode(
        command=(
            "git clone --depth 1 --branch v4.56.2 "
            "https://github.com/huggingface/transformers.git /tmp/transformers && "
            "python /tmp/transformers/examples/pytorch/language-modeling/run_clm.py "
            "--model_name_or_path meta-llama/Llama-3.1-70B-Instruct "
            "--output_dir /opt/ml/model "
            "--do_train"
        ),
    ),
    compute=Compute(instance_type="ml.p5en.48xlarge"),
    environment={"HF_TOKEN": "<YOUR_HF_TOKEN>"},
)
trainer.train()
