from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments, TextDataset, DataCollatorForLanguageModeling

# Carregar o tokenizer e o modelo pré-treinado
tokenizer = GPT2Tokenizer.from_pretrained('pierreguillou/gpt2-small-portuguese')
model = GPT2LMHeadModel.from_pretrained('pierreguillou/gpt2-small-portuguese')

# Preparar o dataset
train_dataset = TextDataset(
    tokenizer=tokenizer,
    file_path='dataset.txt',
    block_size=128
)

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=False,
)

# Configurar os parâmetros de treinamento
training_args = TrainingArguments(
    output_dir='./modelo_treinado',
    overwrite_output_dir=True,
    num_train_epochs=3,
    per_device_train_batch_size=4,
    save_steps=10_000,
    save_total_limit=2,
)

# Instanciar o Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=train_dataset,
)

# Iniciar o treinamento
trainer.train()

# Salvar o modelo
trainer.save_model('./modelo_treinado')
tokenizer.save_pretrained('./modelo_treinado')