# 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 = JumpStartModel(model_id="huggingface-reasoning-nvidia-nemotron-3-nano-30b-a3b-bf16")
example_payloads = model.retrieve_all_examples()

predictor = model.deploy()

for payload in example_payloads:
    response = predictor.predict(payload.body)
    print("Input:\n", payload.body[payload.prompt_key])
    print("Output:\n", response[0]["generated_text"], "\n\n===============\n")

# 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-reasoning-nvidia-nemotron-3-nano-30b-a3b-bf16",
    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-reasoning-nvidia-nemotron-3-nano-30b-a3b-bf16",
    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-reasoning-nvidia-nemotron-3-nano-30b-a3b-bf16",
    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}")
