File size: 1,149 Bytes
0bd62e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
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')
|