import sagemaker
import boto3
from sagemaker.huggingface import HuggingFace
from datasets import load_dataset

# Obtém o papel do SageMaker ou cria um caso ele não exista
try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

# Carrega o dataset
dataset = load_dataset("practical-dreamer/RPGPT_PublicDomain-ShareGPT")

# Configuração dos hiperparâmetros
hyperparameters = {
    'model_name_or_path': 'unsloth/Llama-3.2-11B-Vision-Instruct',
    'dataset_name': 'practical-dreamer/RPGPT_PublicDomain-ShareGPT',
    'output_dir': '/opt/ml/model',
    'learning_rate': 5e-5,
    'per_device_train_batch_size': 4,
    'num_train_epochs': 3,
}

# Configuração do repositório Git para download do script de treinamento
git_config = {
    'repo': 'https://github.com/huggingface/transformers.git',
    'branch': 'v4.37.0'
}

# Configura o Estimador do Hugging Face
huggingface_estimator = HuggingFace(
    entry_point='train.py',
    source_dir='./path/to/script',  # atualize para o caminho correto do script
    instance_type='ml.p3.2xlarge',
    instance_count=1,
    role=role,
    git_config=git_config,
    transformers_version='4.37.0',
    pytorch_version='2.1.0',
    py_version='py310',
    hyperparameters=hyperparameters
)

# Inicia o trabalho de treinamento
huggingface_estimator.fit()