import os
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM  # Cambiado para coincidir con train.py
import asyncio
import selectors

# Parche para Python 3.13 y Gradio
class SafeSelector(selectors.DefaultSelector):
    def get_map(self):
        try:
            return super().get_map()
        except Exception:
            return {}

selector = SafeSelector()
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)

# --- CONEXIÓN CON EL MODELO ENTRENADO ---
# Usamos la ruta donde train.py guardó el modelo final
MODEL_PATH = "final_model"

# Cargamos usando AutoClasses, que es lo que usaste en el entrenamiento
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH)
# ----------------------------------------

def formal_to_dostoevsky(texto_informal):
    if not texto_informal.strip():
        return "El silencio es, a veces, la confesión más aterradora..."
    
    # Aseguramos que el dispositivo sea correcto (CPU por defecto para inferencia simple)
    inputs = tokenizer(
        texto_informal, 
        return_tensors="pt", 
        truncation=True, 
        max_length=128 # Ajustado al max_length que usaste en el entrenamiento (128)
    )
    
    output_ids = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_length=300,
        min_length=30,
        num_beams=5,
        do_sample=True,
        temperature=0.75,
        top_p=0.80,
        no_repeat_ngram_size=3,
        repetition_penalty=1.7,
        early_stopping=True
    )
    
    resultado = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    
    if resultado.strip().lower() == texto_informal.strip().lower():
        return f"¡Ah! ¿Acaso no comprendes que... {resultado}? Es la miseria del alma lo que habita en esas palabras."
            
    return resultado

demo = gr.Interface(
    fn=formal_to_dostoevsky,
    inputs=gr.Textbox(
        label="Texto de entrada", 
        placeholder="Ej: 'Me siento mal porque no tengo dinero y la ciudad es fea.'"
    ),
    outputs=gr.Textbox(label="Versión de Dostoievski"),
    title="Dostoevsky-zador",
    description="Transforma pensamientos cotidianos en pasajes de angustia existencialista rusa."
)

if __name__ == "__main__":
    demo.launch()