| 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') | |