# SageMaker JumpStart provides APIs as part of SageMaker SDK that allow you to deploy and fine-tune models in network isolation using scripts that SageMaker maintains.

from sagemaker.jumpstart.model import JumpStartModel


model_id = "huggingface-vlm-gemma-4-31b-it"

model = JumpStartModel(model_id=model_id)
payload = model.retrieve_example_payload()

predictor = model.deploy()

try:
    response = predictor.predict(payload)
except Exception as e:
    response = str(e)
print(response)
# Fine-tune with Direct Preference Optimization (DPO) using the SageMaker Python SDK v3.

from sagemaker.train.dpo_trainer import DPOTrainer
from sagemaker.train.common import TrainingType
from sagemaker.ai_registry.dataset import DataSet

# Dataset format: JSONL with {"prompt": [...messages...], "chosen": "...", "rejected": "..."} per line
# Download an example dataset: https://jumpstart-cache-prod-us-west-2.s3.us-west-2.amazonaws.com/datasets/oss/dpo/train_dpo_1k.jsonl
dataset = DataSet.create(
    name="my-dpo-dataset",
    source="s3://jumpstart-cache-prod-us-west-2/datasets/oss/dpo/train_dpo_1k.jsonl",
    wait=False,
)

# Create the DPOTrainer
dpo_trainer = DPOTrainer(
    model="huggingface-vlm-gemma-4-31b-it",
    training_type=TrainingType.LORA,
    training_dataset=dataset.arn,
    # You must manually accept the end-user license agreement (EULA) to fine-tune the model.
    accept_eula=False,
)

# Start training
training_job = dpo_trainer.train(wait=True)
print(f"Training Job Status: {training_job.training_job_status}")

# Fine-tune with RL from AI Feedback (RLAIF) using the SageMaker Python SDK v3.

from sagemaker.train.rlaif_trainer import RLAIFTrainer
from sagemaker.train.common import TrainingType
from sagemaker.ai_registry.dataset import DataSet

# Dataset format: JSONL with {"prompt": [...messages...], "reward_model": {"style": "llmj", "ground_truth": "..."}} per line
# Download an example dataset: https://jumpstart-cache-prod-us-west-2.s3.us-west-2.amazonaws.com/datasets/oss/rlaif/train_rlaif_1k.jsonl
dataset = DataSet.create(
    name="my-rlaif-dataset",
    source="s3://jumpstart-cache-prod-us-west-2/datasets/oss/rlaif/train_rlaif_1k.jsonl",
    wait=False,
)

# Create the RLAIFTrainer
rlaif_trainer = RLAIFTrainer(
    model="huggingface-vlm-gemma-4-31b-it",
    training_type=TrainingType.LORA,
    reward_model_id="qwen.qwen3-32b-v1:0",
    training_dataset=dataset.arn,
    # You must manually accept the end-user license agreement (EULA) to fine-tune the model.
    accept_eula=False,
)

# Start training
training_job = rlaif_trainer.train(wait=True)
print(f"Training Job Status: {training_job.training_job_status}")

# Fine-tune with RL with Verifiable Rewards (RLVR) using the SageMaker Python SDK v3.

from sagemaker.train.rlvr_trainer import RLVRTrainer
from sagemaker.train.common import TrainingType
from sagemaker.ai_registry.dataset import DataSet

# Dataset format: JSONL with {"prompt": [...messages...], "reward_model": {"ground_truth": "...", "style": "rule"}} per line
# Download an example dataset: https://jumpstart-cache-prod-us-west-2.s3.us-west-2.amazonaws.com/datasets/oss/rlvr/train_rlvr_1k.jsonl
dataset = DataSet.create(
    name="my-rlvr-dataset",
    source="s3://jumpstart-cache-prod-us-west-2/datasets/oss/rlvr/train_rlvr_1k.jsonl",
    wait=False,
)

# Create the RLVRTrainer
rlvr_trainer = RLVRTrainer(
    model="huggingface-vlm-gemma-4-31b-it",
    training_type=TrainingType.LORA,
    training_dataset=dataset.arn,
    # You must manually accept the end-user license agreement (EULA) to fine-tune the model.
    accept_eula=False,
)

# Start training
training_job = rlvr_trainer.train(wait=True)
print(f"Training Job Status: {training_job.training_job_status}")

# Fine-tune with Supervised Fine-Tuning (SFT) using the SageMaker Python SDK v3.

from sagemaker.train.sft_trainer import SFTTrainer
from sagemaker.train.common import TrainingType
from sagemaker.ai_registry.dataset import DataSet

# Dataset format: JSONL with {"prompt": "...", "completion": "..."} per line
# Download an example dataset: https://jumpstart-cache-prod-us-west-2.s3.us-west-2.amazonaws.com/datasets/oss/sft/train_sft_1k.jsonl
dataset = DataSet.create(
    name="my-sft-dataset",
    source="s3://jumpstart-cache-prod-us-west-2/datasets/oss/sft/train_sft_1k.jsonl",
    wait=False,
)

# Create the SFTTrainer
sft_trainer = SFTTrainer(
    model="huggingface-vlm-gemma-4-31b-it",
    training_type=TrainingType.LORA,
    training_dataset=dataset.arn,
    # You must manually accept the end-user license agreement (EULA) to fine-tune the model.
    accept_eula=False,
)

# Start training
training_job = sft_trainer.train(wait=True)
print(f"Training Job Status: {training_job.training_job_status}")
