Finetuning

# DPO (Direct Preference Optimization) Template

# Cell 0 [markdown]: Fine-Tuning

# Cell 1: Install Dependencies

%pip install --upgrade 'sagemaker>=3.7.1,<4.0' boto3 -q  # NOTEBOOK_ONLY

# Cell 2: Setup & Credentials

import boto3
import json
from pathlib import Path
from botocore.exceptions import ClientError
from sagemaker.ai_registry.dataset import DataSet
from sagemaker.core.resources import ModelPackageGroup
from sagemaker.core.helper.session_helper import Session, get_execution_role
from sagemaker.core import Attribution, set_attribution

set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)

# Setup
sm_client = boto3.Session().client("sagemaker")
sagemaker_session = Session(sagemaker_client=sm_client)
bucket = sagemaker_session.default_bucket()

# Configuration - USER please fill in these fields with your information:

BASE_MODEL = ""  # e.g., "meta-textgeneration-llama-3-8b"
TRAINING_DATA_S3 = ""  # S3 path
S3_OUTPUT_PATH = f"s3://{bucket}/finetuning-output/"
ROLE_ARN = get_execution_role() # You can change this to a specific role.
ACCEPT_EULA = False  # Set to True to accept the base model's End-User License Agreement
MODEL_PACKAGE_GROUP_NAME = ""  # Auto-generated based on use case

# Cell 3: Create Dataset and Model Package Group

# Create Model Package Group
try:
    model_package_group = ModelPackageGroup.create(
        model_package_group_name=MODEL_PACKAGE_GROUP_NAME,
        model_package_group_description="",
    )
    print(f"Created new model package group named {MODEL_PACKAGE_GROUP_NAME}")
except ClientError as e:
    if e.response['Error']['Code'] in ('ResourceInUse', 'ValidationException'):
        model_package_group = ModelPackageGroup.get(model_package_group_name=MODEL_PACKAGE_GROUP_NAME)
        print(f"There is already a model package group with the name {MODEL_PACKAGE_GROUP_NAME}.\nIf you want to save your finetuned model under a different name, change the value of MODEL_PACKAGE_GROUP_NAME in the previous cell.")
    else:
        raise

# Create Dataset
# Register dataset in SageMaker AI Registry. This creates a versioned dataset that can be referenced by ARN
dataset = DataSet.create(
    name=MODEL_PACKAGE_GROUP_NAME,
    source=TRAINING_DATA_S3,
    wait=True
)

TRAINING_DATASET_ARN = dataset.arn
print(f"Here is your model package group ARN: {model_package_group.model_package_group_arn}\n")
print(f"Here is your training dataset ARN: {dataset.arn}")

# Cell 4: Configure Trainer

from sagemaker.train.dpo_trainer import DPOTrainer
from sagemaker.train.common import TrainingType

trainer = DPOTrainer(
    model=BASE_MODEL,
    training_type=TrainingType.LORA,
    model_package_group=model_package_group,
    training_dataset=TRAINING_DATASET_ARN,
    s3_output_path=S3_OUTPUT_PATH,
    sagemaker_session=sagemaker_session,
    #accept_eula=ACCEPT_EULA, # Uncomment for Meta models
    role=ROLE_ARN
)
print("Here are the recommended hyperparameters for the current training job:")
print(f"Batch size: {trainer.hyperparameters.global_batch_size}")
print(f"Learning rate: {trainer.hyperparameters.learning_rate}")
print(f"Adam Beta: {trainer.hyperparameters.adam_beta}") 
# Remove the following two print statements for Nova models (Nova models don't use max_epochs or lr_warmup_steps_ratio)
print(f"Number of epochs: {trainer.hyperparameters.max_epochs}")
print(f"Learning rate warmup steps ratio: {trainer.hyperparameters.lr_warmup_steps_ratio}")


# Cell 5: Hyperparameter Overrides

# To change a hyperparameter, uncomment its corresponding line, and set the value you want.

# Note: If the value you choose is not supported for your model, you will get an error indicating the allowed range.

# Uncomment the following line to change the learning rate
# trainer.hyperparameters.learning_rate = 0.0002

# Uncomment the following line to change the batch size
# trainer.hyperparameters.global_batch_size = 16

# Uncomment the following line to change the number of epochs (unavailable for Nova models)
# trainer.hyperparameters.max_epochs = 5

# Uncomment the following line to change the learning rate warmup steps ratio (unavailable for Nova models)
# trainer.hyperparameters.lr_warmup_steps_ratio = 0.05

# Uncomment the following line to change Adam Beta
# trainer.hyperparameters.adam_beta = 0.01

# Cell 6: Start Training

# Start training
training_job = trainer.train(wait=True)

print(f"Training Job Name: {training_job.training_job_name}")
print(f"Training Status: {training_job.training_job_status}")

# Save manifest
manifest_dir = Path("[PROJECT_DIR]") / "manifests"
manifest_dir.mkdir(parents=True, exist_ok=True)
manifest_path = manifest_dir / f"training-{training_job.training_job_name}.json"
manifest_path.write_text(json.dumps({
    "training_job_name": training_job.training_job_name,
    "model_package_group_name": MODEL_PACKAGE_GROUP_NAME,
}, indent=2))
print(f"Manifest saved: {manifest_path}")

# Cell 7: Plot and Display Metrics  # NOTEBOOK_ONLY_SECTION

import matplotlib.pyplot as plt
import mlflow
from mlflow.tracking import MlflowClient

run_id = training_job.mlflow_details.mlflow_run_id
mlflow.set_tracking_uri(training_job.mlflow_config.mlflow_resource_arn)
client = MlflowClient()

metrics = ["loss_per_batch", "rewards/chosen", "rewards/rejected", "rewards/margins", "acc_per_batch"]
fig, axes = plt.subplots(1, len(metrics), figsize=(4 * len(metrics), 3))
for idx, metric in enumerate(metrics):
    history = client.get_metric_history(run_id, metric)
    axes[idx].plot([h.step for h in history], [h.value for h in history], linewidth=2, marker='o', markersize=4)
    axes[idx].set_xlabel('Step')
    axes[idx].set_ylabel(metric.split('/')[-1])
    axes[idx].set_title(metric, fontweight='bold')
    axes[idx].grid(True, alpha=0.3)

plt.suptitle(f'Training Metrics: {training_job.training_job_name}', fontweight='bold')
plt.tight_layout()
plt.show()

# RLAIF (Reinforcement Learning from AI Feedback) Template — Builtin Reward Prompt

# Cell 0 [markdown]: Fine-Tuning

# Cell 1: Install Dependencies

%pip install --upgrade 'sagemaker>=3.7.1,<4.0' boto3 -q  # NOTEBOOK_ONLY

# Cell 2: Setup & Credentials

import boto3
import json
from pathlib import Path
from botocore.exceptions import ClientError
from sagemaker.ai_registry.dataset import DataSet
from sagemaker.core.resources import ModelPackageGroup
from sagemaker.core.helper.session_helper import Session, get_execution_role
from sagemaker.core import Attribution, set_attribution
from sagemaker.train.rlaif_trainer import RLAIFTrainer
from sagemaker.train.common import TrainingType

set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)

# Setup
sm_client = boto3.Session().client("sagemaker")
sagemaker_session = Session(sagemaker_client=sm_client)
bucket = sagemaker_session.default_bucket()

# Configuration - USER please fill in these fields with your information:

BASE_MODEL = "" # Sagemaker Hub model id
TRAINING_DATA_S3 = ""  # S3 path
S3_OUTPUT_PATH = f"s3://{bucket}/finetuning-output/"
ROLE_ARN = get_execution_role()  # You can change this to a specific role
ACCEPT_EULA = False  # Set to True to accept the base model's End-User License Agreement (OSS models only)
MODEL_PACKAGE_GROUP_NAME = ""  # Auto-generated based on use case

# Reward model — the Bedrock LLM used as judge
# Available models and regions: see references/rlaif_guide.md
REWARD_MODEL_ID = ""

# Builtin reward prompt value — choose one that matches your use case:
# "Builtin.Summarize", "Builtin.Faithfulness", "Builtin.ChainOfThought", "Builtin.Evaluation"
REWARD_PROMPT_VALUE = ""  

# Cell 3: Create Dataset and Model Package Group

# Create Model Package Group
try:
    model_package_group = ModelPackageGroup.create(
        model_package_group_name=MODEL_PACKAGE_GROUP_NAME,
        model_package_group_description="",
    )
    print(f"Created new model package group named {MODEL_PACKAGE_GROUP_NAME}")
except ClientError as e:
    if e.response['Error']['Code'] in ('ResourceInUse', 'ValidationException'):
        model_package_group = ModelPackageGroup.get(model_package_group_name=MODEL_PACKAGE_GROUP_NAME)
        print(f"There is already a model package group with the name {MODEL_PACKAGE_GROUP_NAME}.\nIf you want to save your finetuned model under a different name, change the value of MODEL_PACKAGE_GROUP_NAME in the previous cell.")
    else:
        raise

# Create Dataset
dataset = DataSet.create(
    name=MODEL_PACKAGE_GROUP_NAME,
    source=TRAINING_DATA_S3,
    wait=True
)
TRAINING_DATASET_ARN = dataset.arn

print(f"Here is your model package group ARN: {model_package_group.model_package_group_arn}\n")
print(f"Here is your training dataset ARN: {dataset.arn}")

# Cell 4: Configure Trainer

trainer = RLAIFTrainer(
    model=BASE_MODEL,
    model_package_group=model_package_group,
    reward_model_id=REWARD_MODEL_ID,
    reward_prompt=REWARD_PROMPT_VALUE,
    training_dataset=TRAINING_DATASET_ARN,
    s3_output_path=S3_OUTPUT_PATH,
    sagemaker_session=sagemaker_session,
    #accept_eula=ACCEPT_EULA,  # Uncomment for Meta models
    role=ROLE_ARN,
)

print("Here are the recommended hyperparameters for the current training job:")
print(f"Batch size:    {trainer.hyperparameters.global_batch_size}")
print(f"Learning rate: {trainer.hyperparameters.learning_rate}")
print(f"Epochs:        {trainer.hyperparameters.max_epochs}")

# Cell 5: Hyperparameter Overrides

# To change a hyperparameter, uncomment its corresponding line, and set the value you want.

# Note: If the value you choose is not supported for your model, you will get an error indicating the allowed range.

# Uncomment the following line to change the learning rate
# trainer.hyperparameters.learning_rate = 0.0002

# Uncomment the following line to change the batch size
# trainer.hyperparameters.global_batch_size = 16

# Uncomment the following line to change the number of epochs
# trainer.hyperparameters.max_epochs = 5

# Cell 6: Start Training

training_job = trainer.train(wait=True)

print(f"Training Job Name: {training_job.training_job_name}")
print(f"Training Status: {training_job.training_job_status}")

# Save manifest
manifest_dir = Path("[PROJECT_DIR]") / "manifests"
manifest_dir.mkdir(parents=True, exist_ok=True)
manifest_path = manifest_dir / f"training-{training_job.training_job_name}.json"
manifest_path.write_text(json.dumps({
    "training_job_name": training_job.training_job_name,
    "model_package_group_name": MODEL_PACKAGE_GROUP_NAME,
}, indent=2))
print(f"Manifest saved: {manifest_path}")

# Cell 7: Plot and Display Metrics  # NOTEBOOK_ONLY_SECTION

import matplotlib.pyplot as plt
import mlflow
from mlflow.tracking import MlflowClient

run_id = training_job.mlflow_details.mlflow_run_id
mlflow.set_tracking_uri(training_job.mlflow_config.mlflow_resource_arn)
client = MlflowClient()

metrics = [
    "critic/rewards/mean",
    "response_length/mean",
    "actor/entropy_loss",
    "actor/grad_norm",
    "critic/advantages/mean",
]

fig, axes = plt.subplots(1, len(metrics), figsize=(4 * len(metrics), 3))
for idx, metric in enumerate(metrics):
    history = client.get_metric_history(run_id, metric)
    if history:
        axes[idx].plot([h.step for h in history], [h.value for h in history], linewidth=2, marker='o', markersize=4)
        axes[idx].set_xlabel('Step')
        axes[idx].set_ylabel(metric.split('/')[-1])
        axes[idx].set_title(metric, fontweight='bold')
        axes[idx].grid(True, alpha=0.3)

plt.suptitle(f'Training Metrics: {training_job.training_job_name}', fontweight='bold')
plt.tight_layout()
plt.show()

# RLAIF (Reinforcement Learning from AI Feedback) Template — Custom Reward Prompt

# Cell 0 [markdown]: Fine-Tuning

# Cell 1: Install Dependencies

%pip install --upgrade 'sagemaker>=3.7.1,<4.0' boto3 -q  # NOTEBOOK_ONLY

# Cell 2: Setup & Credentials

import boto3
import json
from pathlib import Path
from botocore.exceptions import ClientError
from sagemaker.ai_registry.dataset import DataSet
from sagemaker.ai_registry.evaluator import Evaluator
from sagemaker.ai_registry.air_constants import REWARD_PROMPT
from sagemaker.core.resources import ModelPackageGroup
from sagemaker.core.helper.session_helper import Session, get_execution_role
from sagemaker.core import Attribution, set_attribution
from sagemaker.train.rlaif_trainer import RLAIFTrainer
from sagemaker.train.common import TrainingType

set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)

# Setup
sm_client = boto3.Session().client("sagemaker")
sagemaker_session = Session(sagemaker_client=sm_client)
bucket = sagemaker_session.default_bucket()

# Configuration - USER please fill in these fields with your information:

BASE_MODEL = ""  # Sagemaker Hub model id
TRAINING_DATA_S3 = ""  # S3 path
S3_OUTPUT_PATH = f"s3://{bucket}/finetuning-output/"
ROLE_ARN = get_execution_role()  # You can change this to a specific role
ACCEPT_EULA = False  # Set to True to accept the base model's End-User License Agreement (OSS models only)
MODEL_PACKAGE_GROUP_NAME = ""  # Auto-generated based on use case

# Reward model — the Bedrock LLM used as judge
# Available models and regions: see references/rlaif_guide.md
REWARD_MODEL_ID = ""

# Cell 3: Register Custom Reward Prompt

# Insert path to the custom Jinja prompt file (usually ../scripts/custom_reward_prompt.jinja)
CUSTOM_PROMPT_PATH = ""

reward_prompt_evaluator = Evaluator.create(
    name="[GENERATE A NAME FOR THE EVALUATOR HERE]",  # lowercase alphanumeric + hyphens, max 20 chars
    type=REWARD_PROMPT,
    source=CUSTOM_PROMPT_PATH,
    sagemaker_session=sagemaker_session,
    wait=True
)
REWARD_PROMPT_ARN = reward_prompt_evaluator.arn
print(f"Reward Prompt Evaluator ARN: {REWARD_PROMPT_ARN}")

# Cell 4: Create Dataset and Model Package Group

# Create Model Package Group
try:
    model_package_group = ModelPackageGroup.create(
        model_package_group_name=MODEL_PACKAGE_GROUP_NAME,
        model_package_group_description="",
    )
    print(f"Created new model package group named {MODEL_PACKAGE_GROUP_NAME}")
except ClientError as e:
    if e.response['Error']['Code'] in ('ResourceInUse', 'ValidationException'):
        model_package_group = ModelPackageGroup.get(model_package_group_name=MODEL_PACKAGE_GROUP_NAME)
        print(f"There is already a model package group with the name {MODEL_PACKAGE_GROUP_NAME}.\nIf you want to save your finetuned model under a different name, change the value of MODEL_PACKAGE_GROUP_NAME in the previous cell.")
    else:
        raise

# Create Dataset
dataset = DataSet.create(
    name=MODEL_PACKAGE_GROUP_NAME,
    source=TRAINING_DATA_S3,
    wait=True
)
TRAINING_DATASET_ARN = dataset.arn

print(f"Here is your model package group ARN: {model_package_group.model_package_group_arn}\n")
print(f"Here is your training dataset ARN: {dataset.arn}")

# Cell 5: Configure Trainer

trainer = RLAIFTrainer(
    model=BASE_MODEL,
    model_package_group=model_package_group,
    reward_model_id=REWARD_MODEL_ID,
    reward_prompt=REWARD_PROMPT_ARN,  # ARN of the registered custom prompt evaluator
    training_dataset=TRAINING_DATASET_ARN,
    s3_output_path=S3_OUTPUT_PATH,
    sagemaker_session=sagemaker_session,
    #accept_eula=ACCEPT_EULA,  # Uncomment for Meta models
    role=ROLE_ARN,
)

print("Here are the recommended hyperparameters for the current training job:")
print(f"Batch size:    {trainer.hyperparameters.global_batch_size}")
print(f"Learning rate: {trainer.hyperparameters.learning_rate}")
print(f"Epochs:        {trainer.hyperparameters.max_epochs}")

# Cell 6: Hyperparameter Overrides

# To change a hyperparameter, uncomment its corresponding line, and set the value you want.

# Note: If the value you choose is not supported for your model, you will get an error indicating the allowed range.

# Uncomment the following line to change the learning rate
# trainer.hyperparameters.learning_rate = 0.0002

# Uncomment the following line to change the batch size
# trainer.hyperparameters.global_batch_size = 16

# Uncomment the following line to change the number of epochs
# trainer.hyperparameters.max_epochs = 5

# Cell 7: Start Training

training_job = trainer.train(wait=True)

print(f"Training Job Name: {training_job.training_job_name}")
print(f"Training Status: {training_job.training_job_status}")

# Save manifest
manifest_dir = Path("[PROJECT_DIR]") / "manifests"
manifest_dir.mkdir(parents=True, exist_ok=True)
manifest_path = manifest_dir / f"training-{training_job.training_job_name}.json"
manifest_path.write_text(json.dumps({
    "training_job_name": training_job.training_job_name,
    "model_package_group_name": MODEL_PACKAGE_GROUP_NAME,
}, indent=2))
print(f"Manifest saved: {manifest_path}")

# Cell 8: Plot and Display Metrics  # NOTEBOOK_ONLY_SECTION

import matplotlib.pyplot as plt
import mlflow
from mlflow.tracking import MlflowClient

run_id = training_job.mlflow_details.mlflow_run_id
mlflow.set_tracking_uri(training_job.mlflow_config.mlflow_resource_arn)
client = MlflowClient()

metrics = [
    "critic/rewards/mean",
    "response_length/mean",
    "actor/entropy_loss",
    "actor/grad_norm",
    "critic/advantages/mean",
]

fig, axes = plt.subplots(1, len(metrics), figsize=(4 * len(metrics), 3))
for idx, metric in enumerate(metrics):
    history = client.get_metric_history(run_id, metric)
    if history:
        axes[idx].plot([h.step for h in history], [h.value for h in history], linewidth=2, marker='o', markersize=4)
        axes[idx].set_xlabel('Step')
        axes[idx].set_ylabel(metric.split('/')[-1])
        axes[idx].set_title(metric, fontweight='bold')
        axes[idx].grid(True, alpha=0.3)

plt.suptitle(f'Training Metrics: {training_job.training_job_name}', fontweight='bold')
plt.tight_layout()
plt.show()

# RLVR (Reinforcement Learning from Verifiable Rewards) Template

# Cell 0 [markdown]: Fine-Tuning

# Cell 1: Install Dependencies

%pip install --upgrade 'sagemaker>=3.7.1,<4.0' boto3 -q  # NOTEBOOK_ONLY

# Cell 2: Setup & Credentials

import boto3
import json
from pathlib import Path
from botocore.exceptions import ClientError
from sagemaker.ai_registry.dataset import DataSet
from sagemaker.core.resources import ModelPackageGroup
from sagemaker.core.helper.session_helper import Session, get_execution_role
from sagemaker.core import Attribution, set_attribution

set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)

# Setup
sm_client = boto3.Session().client("sagemaker")
sagemaker_session = Session(sagemaker_client=sm_client)
bucket = sagemaker_session.default_bucket()

# Configuration - USER please fill in these fields with your information:

BASE_MODEL = ""  # e.g., "meta-textgeneration-llama-3-8b"
TRAINING_DATA_S3 = ""  # S3 path
S3_OUTPUT_PATH = f"s3://{bucket}/finetuning-output/"
ROLE_ARN = get_execution_role() # You can change this to a specific role.
ACCEPT_EULA = False  # Set to True to accept the base model's End-User License Agreement
MODEL_PACKAGE_GROUP_NAME = ""  # Auto-generated based on use case

# Cell 3: Register Reward Function

from sagemaker.ai_registry.evaluator import Evaluator

reward_function_path = "" # Insert path to the local reward function (usually ../scripts/lambda_function.py)

evaluator = Evaluator.create(
    name="[GENERATE A NAME FOR THE EVALUATOR HERE]",
    type="RewardFunction",
    source=reward_function_path,
)
CUSTOM_REWARD_FUNCTION = evaluator.arn
print(f"Reward Function ARN: {CUSTOM_REWARD_FUNCTION}")

# Cell 4: Create Dataset and Model Package Group

# Create Model Package Group
try:
    model_package_group = ModelPackageGroup.create(
        model_package_group_name=MODEL_PACKAGE_GROUP_NAME,
        model_package_group_description="",
    )
    print(f"Created new model package group named {MODEL_PACKAGE_GROUP_NAME}")
except ClientError as e:
    if e.response['Error']['Code'] in ('ResourceInUse', 'ValidationException'):
        model_package_group = ModelPackageGroup.get(model_package_group_name=MODEL_PACKAGE_GROUP_NAME)
        print(f"There is already a model package group with the name {MODEL_PACKAGE_GROUP_NAME}.\nIf you want to save your finetuned model under a different name, change the value of MODEL_PACKAGE_GROUP_NAME in the previous cell.")
    else:
        raise

# Create Dataset
# Register dataset in SageMaker AI Registry. This creates a versioned dataset that can be referenced by ARN
dataset = DataSet.create(
    name=MODEL_PACKAGE_GROUP_NAME,
    source=TRAINING_DATA_S3,
    wait=True
)
TRAINING_DATASET_ARN = dataset.arn

print(f"Here is your model package group ARN: {model_package_group.model_package_group_arn}\n")
print(f"Here is your training dataset ARN: {dataset.arn}")

# Cell 5: Configure Trainer

from sagemaker.train.rlvr_trainer import RLVRTrainer
from sagemaker.train.common import TrainingType


trainer = RLVRTrainer(
    model=BASE_MODEL,
    model_package_group=model_package_group,
    training_dataset=TRAINING_DATASET_ARN,
    s3_output_path=S3_OUTPUT_PATH,
    sagemaker_session=sagemaker_session,
    #accept_eula=ACCEPT_EULA, # Uncomment for Meta models
    role=ROLE_ARN,
    custom_reward_function=CUSTOM_REWARD_FUNCTION
)
print("Here are the recommended hyperparameters for the current training job:")
print(f"Batch size: {trainer.hyperparameters.global_batch_size}")
print(f"Learning rate: {trainer.hyperparameters.learning_rate}")
# Delete the following print statement for Nova models (Nova models don't use max_epochs)
print(f"Number of epochs: {trainer.hyperparameters.max_epochs}") 

# Cell 6: Hyperparameter Overrides

# To change a hyperparameter, uncomment its corresponding line, and set the value you want.

# Note: If the value you choose is not supported for your model, you will get an error indicating the allowed range.

# Uncomment the following line to change the learning rate
# trainer.hyperparameters.learning_rate = 0.0002

# Uncomment the following line to change the batch size
# trainer.hyperparameters.global_batch_size = 16

# Uncomment the following line to change the number of epochs (unavailable for Nova models)
# trainer.hyperparameters.max_epochs = 5

# Cell 7: Start Training

# Start training
training_job = trainer.train(wait=True)

print(f"Training Job Name: {training_job.training_job_name}")
print(f"Training Status: {training_job.training_job_status}")

# Save manifest
manifest_dir = Path("[PROJECT_DIR]") / "manifests"
manifest_dir.mkdir(parents=True, exist_ok=True)
manifest_path = manifest_dir / f"training-{training_job.training_job_name}.json"
manifest_path.write_text(json.dumps({
    "training_job_name": training_job.training_job_name,
    "model_package_group_name": MODEL_PACKAGE_GROUP_NAME,
}, indent=2))
print(f"Manifest saved: {manifest_path}")

# Cell 8: Plot and Display Metrics  # NOTEBOOK_ONLY_SECTION

import matplotlib.pyplot as plt
import mlflow
from mlflow.tracking import MlflowClient

run_id = training_job.mlflow_details.mlflow_run_id
mlflow.set_tracking_uri(training_job.mlflow_config.mlflow_resource_arn)
client = MlflowClient()

# Core RL metrics - adjust val-core metric names based on your data source and reward function
metrics = [
    "critic/rewards/mean",
    "response_length/mean",
    "actor/entropy_loss",
    "actor/grad_norm",
    "critic/advantages/mean",
]
# Note: Validation reward metrics follow the pattern: val-core/{data_source}/reward(/acc)/mean@{k}
# Add your specific val-core metrics to the list above, e.g.:
#   "val-core/my_dataset/reward/mean@1"
# ResponseQuality: Verl allows printing to a file. Check training job output for details.

fig, axes = plt.subplots(1, len(metrics), figsize=(4 * len(metrics), 3))
for idx, metric in enumerate(metrics):
    history = client.get_metric_history(run_id, metric)
    if history:
        axes[idx].plot([h.step for h in history], [h.value for h in history], linewidth=2, marker='o', markersize=4)
        axes[idx].set_xlabel('Step')
        axes[idx].set_ylabel(metric.split('/')[-1])
        axes[idx].set_title(metric, fontweight='bold')
        axes[idx].grid(True, alpha=0.3)

plt.suptitle(f'Training Metrics: {training_job.training_job_name}', fontweight='bold')
plt.tight_layout()
plt.show()

# SFT (Supervised Fine-Tuning) Template

# Cell 0 [markdown]: Fine-Tuning

# Cell 1: Install Dependencies

%pip install --upgrade 'sagemaker>=3.7.1,<4.0' boto3 -q  # NOTEBOOK_ONLY

# Cell 2: Setup & Credentials

import boto3
import json
from pathlib import Path
from botocore.exceptions import ClientError
from sagemaker.ai_registry.dataset import DataSet
from sagemaker.core.resources import ModelPackageGroup
from sagemaker.core.helper.session_helper import Session, get_execution_role
from sagemaker.core import Attribution, set_attribution

set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)

# Setup
sm_client = boto3.Session().client("sagemaker")
sagemaker_session = Session(sagemaker_client=sm_client)
bucket = sagemaker_session.default_bucket()

# Configuration - USER please fill in these fields with your information:

BASE_MODEL = ""  # e.g., "meta-textgeneration-llama-3-8b"
TRAINING_DATA_S3 = ""  # S3 path
S3_OUTPUT_PATH = f"s3://{bucket}/finetuning-output/"
ROLE_ARN = get_execution_role() # You can change this to a specific role.
ACCEPT_EULA = False  # Set to True to accept the base model's End-User License Agreement
MODEL_PACKAGE_GROUP_NAME = ""  # Auto-generated based on use case

# Cell 3: Create Dataset and Model Package Group

# Create Model Package Group
try:
    model_package_group = ModelPackageGroup.create(
        model_package_group_name=MODEL_PACKAGE_GROUP_NAME,
        model_package_group_description="",
    )
    print(f"Created new model package group named {MODEL_PACKAGE_GROUP_NAME}")
except ClientError as e:
    if e.response['Error']['Code'] in ('ResourceInUse', 'ValidationException'):
        model_package_group = ModelPackageGroup.get(model_package_group_name=MODEL_PACKAGE_GROUP_NAME)
        print(f"There is already a model package group with the name {MODEL_PACKAGE_GROUP_NAME}. If you want to save your finetuned model under a different name, change the value of MODEL_PACKAGE_GROUP_NAME in the previous cell.")
    else:
        raise

# Create Dataset
# Register dataset in SageMaker AI Registry. This creates a versioned dataset that can be referenced by ARN
dataset = DataSet.create(
    name=MODEL_PACKAGE_GROUP_NAME,
    source=TRAINING_DATA_S3,
    wait=True
)

TRAINING_DATASET_ARN = dataset.arn
print(f"Here is your model package group ARN: {model_package_group.model_package_group_arn}\n")
print(f"Here is your training dataset ARN: {dataset.arn}")

# Cell 4: Configure Trainer

from sagemaker.train.sft_trainer import SFTTrainer
from sagemaker.train.common import TrainingType


trainer = SFTTrainer(
    model=BASE_MODEL,
    training_type=TrainingType.LORA,
    model_package_group=model_package_group,
    training_dataset=TRAINING_DATASET_ARN,
    s3_output_path=S3_OUTPUT_PATH,
    sagemaker_session=sagemaker_session,
    #accept_eula=ACCEPT_EULA, # Uncomment for Meta models
    role=ROLE_ARN
)

print("Here are the recommended hyperparameters for the current training job:")
print(f"Batch size: {trainer.hyperparameters.global_batch_size}")
print(f"Learning rate: {trainer.hyperparameters.learning_rate}")
# Remove the following two print statements for Nova models (Nova models don't use max_epochs or lr_warmup_steps_ratio)
print(f"Number of epochs: {trainer.hyperparameters.max_epochs}")
print(f"Learning rate warmup steps ratio: {trainer.hyperparameters.lr_warmup_steps_ratio}")

# Cell 5: Hyperparameter Overrides

# To change a hyperparameter, uncomment its corresponding line, and set the value you want.

# Note: If the value you choose is not supported for your model, you will get an error indicating the allowed range.

# Uncomment the following line to change the learning rate
# trainer.hyperparameters.learning_rate = 0.0002

# Uncomment the following line to change the batch size
# trainer.hyperparameters.global_batch_size = 16

# Uncomment the following line to change the number of epochs (unavailable for Nova models)
# trainer.hyperparameters.max_epochs = 5

# Uncomment the following line to change the learning rate warmup steps ratio (unavailable for Nova models)
# trainer.hyperparameters.lr_warmup_steps_ratio = 0.05

# Cell 6: Start Training

# Start training
training_job = trainer.train(wait=True)

print(f"Training Job Name: {training_job.training_job_name}")
print(f"Training Status: {training_job.training_job_status}")

# Save manifest
manifest_dir = Path("[PROJECT_DIR]") / "manifests"
manifest_dir.mkdir(parents=True, exist_ok=True)
manifest_path = manifest_dir / f"training-{training_job.training_job_name}.json"
manifest_path.write_text(json.dumps({
    "training_job_name": training_job.training_job_name,
    "model_package_group_name": MODEL_PACKAGE_GROUP_NAME,
}, indent=2))
print(f"Manifest saved: {manifest_path}")

# Cell 7: Plot and Display Metrics  # NOTEBOOK_ONLY_SECTION

import matplotlib.pyplot as plt
import mlflow
from mlflow.tracking import MlflowClient

run_id = training_job.mlflow_details.mlflow_run_id
mlflow.set_tracking_uri(training_job.mlflow_config.mlflow_resource_arn)
client = MlflowClient()

fig, axes = plt.subplots(1, 2, figsize=(12, 3))
for idx, metric in enumerate(["total_loss", "val_eval_total_loss"]):
    history = client.get_metric_history(run_id, metric)
    axes[idx].plot([h.step for h in history], [h.value for h in history], linewidth=2, marker='o', markersize=4)
    axes[idx].set_xlabel('Step')
    axes[idx].set_ylabel('Loss')
    axes[idx].set_title(metric, fontweight='bold')
    axes[idx].grid(True, alpha=0.3)

plt.suptitle(f'Training Metrics: {training_job.training_job_name}', fontweight='bold')
plt.tight_layout()
plt.show()

# scripts/mlflow_reference.py
# Reference for querying MLflow metrics from a training job.
# The agent reads this to understand the pattern, then writes
# its own code adapted to what the user needs.

import os
os.environ['AWS_DEFAULT_REGION'] = '[REGION]'

from sagemaker.core.resources import TrainingJob
import mlflow
from mlflow.tracking import MlflowClient

# Connect to MLflow via the training job
tj = TrainingJob.get(training_job_name='[TRAINING_JOB_NAME]')
mlflow.set_tracking_uri(tj.mlflow_config.mlflow_resource_arn)
client = MlflowClient()
run_id = tj.mlflow_details.mlflow_run_id

# List available metrics
run = client.get_run(run_id)
print(run.data.metrics.keys())

# Get full history for a metric
history = client.get_metric_history(run_id, '[METRIC_NAME]')
for h in history:
    print(f"step={h.step}, value={h.value:.4f}")

"""
Provide your custom reward function code below. Learn about the available libraries and templates that you can use
at: https://docs.aws.amazon.com/sagemaker/latest/dg/customize-model.html.

- You must add your evaluation logic in the reward_function() function
- Do not remove the lambda_handler() function or modify its schema as it is required to create the reward function
"""

import json  # For JSON parsing - adjust imports based on your use case
import re    # For pattern matching and validation
from typing import Dict, Any, List, Optional, Union # For type hints
# Add any other imports your use case requires

# ========================================================================================
#  NOTE: INITIAL SUGGESTION ONLY - MUST BE CUSTOMIZED
#
#     YOU MUST:
#     1. Review and update each section per YOUR use case
#     2. Customize the logic for YOUR SPECIFIC requirements
#     3. Replace example values (field names, thresholds, etc.) with your actual values
#     4. Test thoroughly before using
#
#     DO NOT use this code as-is. It will not work until you uncomment and customize it.
# =========================================================================================


# =========================================================================================
# SECTION 1: Helper function — content normalization
# =========================================================================================
# Nova messages use content as a string, a list of {"type":"text","text":"..."} chunks,
# or a dict with a "text" key. This helper normalizes all forms to a plain string.
def content_to_text(content: Any) -> str:
    """
    Normalize Nova message content to a plain string.

    Args:
        content: String, list of text chunks, or dict with "text" key

    Returns:
        Plain text string
    """
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: List[str] = []
        for item in content:
            if isinstance(item, str):
                parts.append(item)
            elif isinstance(item, dict) and "text" in item:
                parts.append(item["text"])
            else:
                parts.append(str(item))
        return "".join(parts)
    if isinstance(content, dict) and "text" in content:
        return content["text"]
    return str(content)


# =========================================================================================
# SECTION 2: Helper function — ground truth extraction
# =========================================================================================
# Nova reference_answer can be a dict with flexible keys (answer, label, sentiment, etc.),
# a JSON string, or a plain string.
def coerce_ground_truth(ground_truth: Union[str, Dict[str, Any], Any]) -> Optional[str]:
    """
    Extract the ground-truth answer as a string from reference_answer.

    Args:
        ground_truth: Dict, JSON string, or plain string

    Returns:
        Ground truth string, or None if not found
    """
    if ground_truth is None:
        return None

    if isinstance(ground_truth, str):
        s = ground_truth.strip()
        if not s:
            return None
        if s.startswith("{") or s.startswith("["):
            try:
                ground_truth = json.loads(s)
            except Exception:
                return s
        else:
            return s

    if isinstance(ground_truth, dict):
        for key in ("ground_truth", "answer", "label", "sentiment", "polarity", "target"):
            if key in ground_truth and ground_truth[key] is not None:
                return str(ground_truth[key])
        if len(ground_truth) == 1:
            only_val = next(iter(ground_truth.values()))
            if only_val is not None:
                return str(only_val)
        return None

    return str(ground_truth)


# =========================================================================================
# SECTION 3: Helper function — number extraction
# =========================================================================================
# TODO: UPDATE or REMOVE the helper function as per YOUR use case
# Note the below lines of code are examples and will not work for your use case
# You MUST update them to match YOUR use case
def extract_number(text: str) -> Optional[float]:
    """
    Extract numerical answer from text.
    Looks for numbers after answer keywords, or returns the last number found.

    Args:
        text: Text containing a numerical answer

    Returns:
        Extracted number as float, or None if no number found
    """
    if not text:
        return None

    # Try to find numbers after common answer keywords
    answer_patterns = [
        r'(?:equals|is|answer is|result is|=)\s*(-?\d+\.?\d*)',
        r'(?:answer|result|solution):\s*(-?\d+\.?\d*)',
    ]

    for pattern in answer_patterns:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            try:
                return float(match.group(1))
            except ValueError:
                pass

    # Fallback: find all numbers and return the last one (likely the answer)
    pattern = r'-?\d+\.?\d*'
    matches = re.findall(pattern, text)

    if matches:
        try:
            return float(matches[-1])
        except ValueError:
            return None

    return None


# =========================================================================================
# SECTION 4: Helper function — reasoning quality
# =========================================================================================
# TODO: UPDATE or REMOVE the helper function as per YOUR use case
# Note the below lines of code are examples and will not work for your use case
# You MUST update them to match YOUR use case
def compute_reasoning_quality(response: str) -> float:
    """
    Compute reasoning quality score based on response characteristics.
    This is a simple heuristic - customize based on your needs.

    Args:
        response: The model's response text

    Returns:
        Quality score between 0.0 and 1.0
    """
    if not response:
        return 0.0

    score = 0.0

    # Check for reasoning indicators (customize these for your use case)
    reasoning_indicators = [
        'because', 'therefore', 'thus', 'since', 'so',
        'first', 'second', 'then', 'finally',
        'step', 'calculate', 'compute', 'equals'
    ]

    response_lower = response.lower()

    # Award points for reasoning indicators (max 0.55)
    indicator_count = sum(1 for indicator in reasoning_indicators if indicator in response_lower)
    score += min(indicator_count * 0.11, 0.55)

    # Award points for response length (indicates detailed reasoning, max 0.25)
    if len(response) > 30:
        score += 0.05
    if len(response) > 60:
        score += 0.1
    if len(response) > 120:
        score += 0.1

    # Award points for structured response (max 0.2)
    if '\n' in response or '.' in response:
        score += 0.2

    return min(score, 1.0)


# =========================================================================================
# SECTION 5: Helper function — answer extraction
# =========================================================================================
# TODO: UPDATE or REMOVE the helper function as per YOUR use case
# Note the below lines of code are examples and will not work for your use case
# You MUST update them to match YOUR use case
def extract_answer(response: str) -> Optional[str]:
    """
    Extract the answer from a Nova model response.
    Looks for <|begin_of_solution|>...<|end_of_solution|> blocks and \\boxed{} patterns.

    Args:
        response: The model's response text

    Returns:
        Extracted answer string, or None if not found
    """
    if not response:
        return None

    # Try solution block first
    solution_match = re.search(
        r"<\|begin_of_solution\|>(.*?)<\|end_of_solution\|>",
        response,
        re.DOTALL,
    )
    if solution_match:
        boxed = re.findall(r"\\boxed\{([^}]+)\}", solution_match.group(1))
        if boxed:
            return boxed[-1].strip()

    # Fallback: boxed anywhere
    boxed = re.findall(r"\\boxed\{([^}]+)\}", response)
    if boxed:
        return boxed[-1].strip()

    return None


# =========================================================================================
# SECTION 6: Sample reward function
# =========================================================================================
# TODO: UPDATE or REMOVE the reward function as per YOUR use case
# Note the below lines of code are examples and will not work for your use case
# You MUST update them to match YOUR use case
def reward_function(sample: Dict[str, Any], index: int) -> Dict[str, Any]:
    """
    Args:
        sample: Dictionary containing messages and reference_answer
        index: Sample index in batch

    Returns:
        Dictionary with reward scores and metrics
    """
    # ========================================================================
    # SECTION 7: Parse input
    # ========================================================================
    # TODO: UPDATE logic to parse the input as per YOUR use case
    # Note the below lines of code are examples and will not work for your use case
    # You MUST update them to match YOUR use case
    messages = sample.get('messages', [])
    ground_truth = sample.get('reference_answer', {})

    # Get the assistant's response (last message with role assistant or nova_assistant)
    response = ""
    for msg in messages:
        role = msg.get('role', '')
        if role in ('assistant', 'nova_assistant'):
            response = content_to_text(msg.get('content', ''))

    # Extract numerical answers
    predicted = extract_number(response)
    expected_str = coerce_ground_truth(ground_truth)
    expected = extract_number(expected_str) if expected_str else None

    # Compute metrics
    exact_match = 0.0
    answer_present = 0.0
    reasoning_quality = compute_reasoning_quality(response)

    if predicted is not None and expected is not None:
        exact_match = 1.0 if abs(predicted - expected) < 1e-6 else 0.0
        answer_present = 1.0

    # ========================================================================
    # SECTION 8: Compute reward scores
    # ========================================================================
    # TODO: UPDATE logic to compute aggregate score
    # Note the below lines of code are examples and will not work for your use case
    # You MUST update them to match YOUR use case
    aggregate_reward = 0.7 * exact_match + 0.3 * reasoning_quality

    # ========================================================================
    # SECTION 9: Form the metrics list
    # ========================================================================
    # TODO: UPDATE logic to compute metrics list
    # Note the below lines of code are examples and will not work for your use case
    # You MUST update them to match YOUR use case
    metrics = [
        {
            'name': 'exact_match',
            'value': float(exact_match),
            'type': 'Reward'
        },
        {
            'name': 'answer_present',
            'value': float(answer_present),
            'type': 'Metric'
        },
        {
            'name': 'reasoning_quality',
            'value': float(reasoning_quality),
            'type': 'Metric'
        }
    ]

    # ========================================================================
    # SECTION 10: Return output
    # ========================================================================
    # TODO: UPDATE the return statement to return YOUR output
    # UPDATE the key before creating the evaluator
    # Note the below lines of code are examples and will not work for your use case
    # You MUST update them to match YOUR use case

    return {
        'id': str(sample.get('id', f'sample-{index:03d}')),
        'aggregate_reward_score': float(aggregate_reward),
        'metrics_list': metrics
    }

def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """
    AWS Lambda Handler for reward function.
    SageMaker Nova RLVR invokes this with a bare list of samples and expects
    a bare list of {id, aggregate_reward_score, ...} dicts in return.
    """
    # Event is a bare list of samples
    batch = event if isinstance(event, list) else []

    results = []
    for i, sample in enumerate(batch):
        try:
            result = reward_function(sample, i)
            results.append(result)
        except Exception as e:
            results.append({
                'id': str(sample.get('id', f'sample-{i:03d}') if isinstance(sample, dict) else f'sample-{i:03d}'),
                'aggregate_reward_score': 0.0,
                'metrics_list': []
            })

    return results

"""
Provide your custom reward function code below. Learn about the available libraries and templates that you can use
at: https://docs.aws.amazon.com/sagemaker/latest/dg/customize-model.html.

- You must add your evaluation logic in the reward_function() function
- Do not remove the lambda_handler() function or modify its schema as it is required to create the reward function
"""

import json  # For JSON parsing - adjust imports based on your use case
import re    # For pattern matching and validation
from typing import Dict, Any, List, Optional # For type hints
# Add any other imports your use case requires

# ========================================================================================
#  NOTE: INITIAL SUGGESTION ONLY - MUST BE CUSTOMIZED
#
#     YOU MUST:
#     1. Review and update each section per YOUR use case
#     2. Customize the logic for YOUR SPECIFIC requirements
#     3. Replace example values (field names, thresholds, etc.) with your actual values
#     4. Test thoroughly before using
#
#     DO NOT use this code as-is. It will not work until you uncomment and customize it.
# =========================================================================================


# =========================================================================================
# SECTION 1: Helper function 1
# =========================================================================================
# TODO: UPDATE or REMOVE the helper function as per YOUR use case
# Note the below lines of code are examples and will not work for your use case
# You MUST update them to match YOUR use case
def extract_number(text: str) -> Optional[float]:
    """
    Extract numerical answer from text.
    Looks for numbers after answer keywords, or returns the last number found.
    
    Args:
        text: Text containing a numerical answer
        
    Returns:
        Extracted number as float, or None if no number found
    """
    if not text:
        return None

    # Try to find numbers after common answer keywords
    answer_patterns = [
        r'(?:equals|is|answer is|result is|=)\s*(-?\d+\.?\d*)',
        r'(?:answer|result|solution):\s*(-?\d+\.?\d*)',
    ]

    for pattern in answer_patterns:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            try:
                return float(match.group(1))
            except ValueError:
                pass

    # Fallback: find all numbers and return the last one (likely the answer)
    pattern = r'-?\d+\.?\d*'
    matches = re.findall(pattern, text)

    if matches:
        try:
            return float(matches[-1])  # Return last number instead of first
        except ValueError:
            return None

    return None

# =========================================================================================
# SECTION 2: Helper function 2
# =========================================================================================
# TODO: UPDATE or REMOVE the helper function as per YOUR use case
# Note the below lines of code are examples and will not work for your use case
# You MUST update them to match YOUR use case
def compute_reasoning_quality(response: str) -> float:
    """
    Compute reasoning quality score based on response characteristics.
    This is a simple heuristic - customize based on your needs.
    
    Args:
        response: The model's response text
        
    Returns:
        Quality score between 0.0 and 1.0
    """
    if not response:
        return 0.0

    score = 0.0

    # Check for reasoning indicators (customize these for your use case)
    reasoning_indicators = [
        'because', 'therefore', 'thus', 'since', 'so',
        'first', 'second', 'then', 'finally',
        'step', 'calculate', 'compute', 'equals'
    ]

    response_lower = response.lower()

    # Award points for reasoning indicators (max 0.55)
    indicator_count = sum(1 for indicator in reasoning_indicators if indicator in response_lower)
    score += min(indicator_count * 0.11, 0.55)

    # Award points for response length (indicates detailed reasoning, max 0.25)
    if len(response) > 30:
        score += 0.05
    if len(response) > 60:
        score += 0.1
    if len(response) > 120:
        score += 0.1

    # Award points for structured response (max 0.2)
    if '\n' in response or '.' in response:
        score += 0.2

    return min(score, 1.0)

# =========================================================================================
# SECTION 3: Sample reward function
# =========================================================================================
# TODO: UPDATE or REMOVE the reward function as per YOUR use case
# Note the below lines of code are examples and will not work for your use case
# You MUST update them to match YOUR use case
def reward_function(sample: Dict[str, Any], index: int) -> Dict[str, Any]:
    """
    Args:
        sample: Dictionary containing messages and reference_answer
        index: Sample index in batch
    
    Returns:
        Dictionary with reward scores and metrics
    """
    # ========================================================================
    # SECTION 4: Parse input
    # ========================================================================
    # TODO: UPDATE logic to parse the input as per YOUR use case
    # Note the below lines of code are examples and will not work for your use case
    # You MUST update them to match YOUR use case
    # Extract the response and reference
    messages = sample.get('messages', sample.get('prompt', []))
    reference_answer = sample.get('reference_answer', {}).get('text', '') or sample.get('reward_model', {}).get('ground_truth', '')

    # Get the question and assistant's response
    question = ""
    response = ""
    for msg in messages:
        if msg.get('role') == 'user':
            question = msg.get('content', '')
        elif msg.get('role') == 'assistant':
            response = msg.get('content', '')

    # Extract numerical answers
    predicted = extract_number(response)
    expected = extract_number(reference_answer)

    # Compute metrics
    exact_match = 0.0
    answer_present = 0.0
    reasoning_quality = compute_reasoning_quality(response)

    if predicted is not None and expected is not None:
        exact_match = 1.0 if abs(predicted - expected) < 1e-6 else 0.0
        answer_present = 1.0

    # ========================================================================
    # SECTION 5: Compute reward scores
    # ========================================================================
    # TODO: UPDATE logic to compute aggregate score 
    # Note the below lines of code are examples and will not work for your use case
    # You MUST update them to match YOUR use case
    # Aggregate reward computation
    aggregate_reward = 0.7 * exact_match + 0.3 * reasoning_quality

    # ========================================================================
    # SECTION 6: Form the metrics list
    # ========================================================================
    # TODO: UPDATE logic to compute metrics list
    # Note the below lines of code are examples and will not work for your use case
    # You MUST update them to match YOUR use case
    metrics = [
        {
            'name': 'exact_match',
            'value': float(exact_match),
            'type': 'Reward'
        },
        {
            'name': 'answer_present',
            'value': float(answer_present),
            'type': 'Metric'
        },
        {
            'name': 'reasoning_quality',
            'value': float(reasoning_quality),
            'type': 'Metric'
        }
    ]

    # ========================================================================
    # SECTION 7: Return output
    # ========================================================================
    # TODO: UPDATE the return statement to return YOUR outout
    # UPDATE the key before creating the evaluator
    # Note the below lines of code are examples and will not work for your use case
    # You MUST update them to match YOUR use case

    return {
        'id': str(sample.get('my_key', f'sample-{index:03d}')),  # Use formatted index as fallback
        'aggregate_reward_score': float(aggregate_reward),
        'metrics_list': metrics
    }

def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """
    AWS Lambda Handler for reward function
    """
    try:
        # Extract batch from event
        batch = event.get('input', event) if isinstance(event, dict) else event
        if 'batch' in event:
            batch = event.get('batch', [])
        elif 'body' in event:
            body = json.loads(event.get('body', '{}'))
            batch = body.get('batch', [])

        if not batch:
            return {"error":"Missing or empty batch"}

        # Process each sample
        results = []
        for i, sample in enumerate(batch):
            try:
                result = reward_function(sample, i)
                results.append(result)
            except Exception as e:
                return {"error": str(e)}

        return {
            'statusCode': 200,
            'headers': {'Content-Type': 'application/json'},
            'body': json.dumps(results)
        }
    except Exception as e:
        return {
            'statusCode': 400,
            'body': json.dumps({"error": str(e)})
        }