import os

from datasets import Dataset
from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments

# Formateamos las preguntas y respuestas en el formato adecuado
data = [
    {
        "input": "¿Cómo se conocieron Emmanuel y Rebeca?",
        "output": "Rebeca y Emmanuel se conocieron en su primer día de trabajo cuando se estrecharon la mano..."
    },
    # Agrega más pares pregunta-respuesta aquí
]

# Crear el dataset
dataset = Dataset.from_dict({
    'input': [entry['input'] for entry in data],
    'output': [entry['output'] for entry in data]
})

# Cargar el tokenizer de GPT-2
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# Asignar el token de padding al eos_token
tokenizer.pad_token = tokenizer.eos_token

# Tokenizar el dataset
def tokenize_function(examples):
    return tokenizer(examples['input'], padding='max_length', truncation=True, max_length=512)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

# Cargar el modelo GPT-2 Small
model = GPT2LMHeadModel.from_pretrained('gpt2')

# Configurar los parámetros de entrenamiento
output_dir = "./app"

training_args = TrainingArguments(
    output_dir=output_dir,  # Directorio para guardar el modelo
    num_train_epochs=3,  # Número de épocas (ajustable según el tamaño de tu dataset)
    per_device_train_batch_size=2,  # Tamaño del batch
    per_device_eval_batch_size=2,  # Tamaño del batch de evaluación
    warmup_steps=500,  # Número de pasos de calentamiento
    weight_decay=0.01,  # Decaimiento del peso
    logging_dir='./logs',  # Directorio para los logs
    logging_steps=10,
    save_steps=500,  # Guardar el modelo cada 500 pasos
)

# Definir el trainer
trainer = Trainer(
    model=model,  # El modelo de GPT-2
    args=training_args,  # Los parámetros de entrenamiento
    train_dataset=tokenized_datasets,  # El dataset tokenizado
    eval_dataset=tokenized_datasets,  # El dataset de evaluación (puedes dividir tu dataset en entrenamiento y validación)
)


# Iniciar el entrenamiento
trainer.train()

# Guardar el modelo entrenado


trainer.save_model('./app/gpt2-small-finetuned')