# --- PASO 0: Instalación de Librerías ---
# Descomenta y ejecuta esta celda si no tienes las librerías instaladas
# !pip install transformers datasets accelerate -U
# !pip install sentencepiece # Necesario para algunos tokenizers como el de mT5

import datasets
from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    Seq2SeqTrainingArguments,
    Seq2SeqTrainer,
    DataCollatorForSeq2Seq
)

# --- PASO 1: Preparar el Dataset (¡¡¡REEMPLAZAR CON TU DATASET REAL!!!) ---
# Este es un EJEMPLO MUY PEQUEÑO. Necesitas un dataset GRANDE y REAL.
# Tu dataset debería tener miles de ejemplos.
raw_data = {
    "train": [
        {"input_text": "El medico me receto analgesicos.", "target_text": "El médico me recetó analgésicos."},
        {"input_text": "La cancion tiene una melodia pegadiza.", "target_text": "La canción tiene una melodía pegadiza."},
        {"input_text": "La accion legal continuo.", "target_text": "La acción legal continuó."},
        {"input_text": "Manana habra reunion.", "target_text": "Mañana habrá reunión."},
        # ... Agrega MUCHOS MÁS ejemplos aquí ...
    ],
    "validation": [
        {"input_text": "La informacion es util.", "target_text": "La información es útil."},
        {"input_text": "El analisis fue complejo.", "target_text": "El análisis fue complejo."},
        # ... Agrega ejemplos de validación ...
    ]
}

# Convertir el diccionario a un objeto Dataset de Hugging Face
raw_datasets = datasets.DatasetDict({
    "train": Dataset.from_dict(raw_data["train"]),
    "validation": Dataset.from_dict(raw_data["validation"])
})

print("Dataset de ejemplo cargado:")
print(raw_datasets)

# --- PASO 2: Cargar Tokenizer y Modelo ---
# Elegimos un modelo Seq2Seq preentrenado. mT5-small es relativamente ligero.
# Para mejores resultados, considera modelos más grandes o específicos para español si están disponibles.
model_checkpoint = "google/mt5-small"

tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, legacy=False) # legacy=False recomendado para T5
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)

# --- PASO 3: Preprocesar el Dataset ---
max_input_length = 128
max_target_length = 128

def preprocess_function(examples):
    # Asegurarse que los textos son strings
    inputs = [str(ex) for ex in examples["input_text"]]
    targets = [str(ex) for ex in examples["target_text"]]

    # Tokenizar entradas
    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True, padding="max_length")

    # Tokenizar salidas (targets) como etiquetas
    # Usamos el tokenizer en modo "target" para preparar las labels correctamente para el decoder
    with tokenizer.as_target_tokenizer():
        labels = tokenizer(targets, max_length=max_target_length, truncation=True, padding="max_length")

    # Los IDs de padding en las etiquetas deben ser ignorados en la función de pérdida
    # Reemplazamos los pad_token_id por -100
    labels["input_ids"] = [
        [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
    ]

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

# Aplicar la función de preprocesamiento a todo el dataset
tokenized_datasets = raw_datasets.map(preprocess_function, batched=True)

print("\nDataset tokenizado:")
print(tokenized_datasets)

# --- PASO 4: Configurar el Entrenamiento ---

# Directorio donde se guardarán los resultados y checkpoints
output_dir = "./results-accent-corrector"

# Argumentos de entrenamiento
training_args = Seq2SeqTrainingArguments(
    output_dir=output_dir,
    evaluation_strategy="epoch",        # Evaluar después de cada época
    learning_rate=2e-5,                 # Tasa de aprendizaje
    per_device_train_batch_size=8,      # Tamaño de lote por GPU/CPU para entrenamiento (ajusta según tu memoria)
    per_device_eval_batch_size=8,       # Tamaño de lote por GPU/CPU para evaluación
    weight_decay=0.01,                  # Decaimiento de peso
    save_total_limit=3,                 # Guardar solo los últimos 3 checkpoints
    num_train_epochs=3,                 # Número de épocas (ajusta según tu dataset y convergencia)
    predict_with_generate=True,         # Necesario para generar texto durante la evaluación Seq2Seq
    fp16=False,                         # Cambiar a True si tienes GPU compatible y quieres acelerar (requiere 'accelerate')
    logging_steps=10,                   # Frecuencia de logging
    # report_to="tensorboard",          # Opcional: para visualizar en TensorBoard
    # push_to_hub=False,                # Cambiar a True para subir tu modelo al Hub de Hugging Face (necesitas login)
)

# Data Collator: prepara los lotes de datos para el modelo Seq2Seq
data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)

# Crear el objeto Trainer
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    # compute_metrics=compute_metrics, # Opcional: define una función para métricas más avanzadas
)

# --- PASO 5: Entrenar el Modelo ---
print("\nIniciando entrenamiento...")
trainer.train()
print("Entrenamiento completado.")

# --- PASO 6: Guardar el Modelo Final y Tokenizer ---
final_model_path = f"{output_dir}/final_model"
trainer.save_model(final_model_path)
tokenizer.save_pretrained(final_model_path)
print(f"Modelo final y tokenizer guardados en: {final_model_path}")

# --- PASO 7: Probar el Modelo Entrenado (Inferencia) ---
from transformers import pipeline

print("\nProbando el modelo afinado...")

# Cargar el pipeline con el modelo afinado localmente
# Si el entrenamiento fue corto o con pocos datos, el resultado puede no ser bueno.
corrector_pipe = pipeline("text2text-generation", model=final_model_path, tokenizer=final_model_path)

# Frases de ejemplo sin acentos
test_sentences = [
    "Que dia es manana?",
    "La musica clasica me relaja.",
    "El exito requiere dedicacion.",
    "accion" # El ejemplo original del usuario
]

for sentence in test_sentences:
    result = corrector_pipe(sentence, max_length=50) # Ajusta max_length si es necesario
    corrected_text = result[0]['generated_text']
    print(f"Entrada: {sentence}")
    print(f"Salida:  {corrected_text}")
    print("---")

print("\n¡Proceso terminado!")