JorgeIsaac/Writtertrainer

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

Cargar modelo y tokenizer preentrenado

model_name = "gpt2" model = GPT2LMHeadModel.from_pretrained(model_name) tokenizer = GPT2Tokenizer.from_pretrained(model_name)

Crear dataset a partir de los textos del escritor

def load_dataset(file_path, tokenizer, block_size=128): dataset = TextDataset( tokenizer=tokenizer, file_path=file_path, block_size=block_size ) return dataset

file_path = "ruta/al/texto/del/escritor.txt" dataset = load_dataset(file_path, tokenizer)

Crear collator para el dataset

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

Definir parámetros de entrenamiento

training_args = TrainingArguments( output_dir="./results", overwrite_output_dir=True, num_train_epochs=1, per_device_train_batch_size=2, save_steps=10_000, save_total_limit=2, )

Crear Trainer y entrenar el modelo

trainer = Trainer( model=model, args=training_args, data_collator=data_collator, train_dataset=dataset )

trainer.train()