import gradio as gr
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast

# Cargar el modelo y el tokenizador
model = MBartForConditionalGeneration.from_pretrained("eslamxm/MBART-finetuned-Spanish")
tokenizer = MBart50TokenizerFast.from_pretrained("eslamxm/MBART-finetuned-Spanish")

# Función para generar el resumen
def resumir_texto(texto):
    inputs = tokenizer(texto, return_tensors="pt", max_length=1024, truncation=True)
    resumen_ids = model.generate(
        inputs["input_ids"],
        max_length=150,
        min_length=40,
        length_penalty=2.0,
        num_beams=4,
        early_stopping=True
    )
    resumen = tokenizer.decode(resumen_ids[0], skip_special_tokens=True)
    return resumen

# Crear la interfaz con Gradio
with gr.Blocks() as demo:
    gr.Markdown("## 📝 Resumen de Textos en Español")
    gr.Markdown("Introduce un texto en español y obtén un resumen generado por el modelo MBART.")

    with gr.Row():
        with gr.Column():
            entrada = gr.Textbox(
                label="Texto de entrada",
                placeholder="Escribe o pega aquí el texto que deseas resumir...",
                lines=10
            )
            boton = gr.Button("Generar Resumen")
        with gr.Column():
            salida = gr.Textbox(
                label="Resumen generado",
                placeholder="El resumen aparecerá aquí...",
                lines=10
            )

    boton.click(fn=resumir_texto, inputs=entrada, outputs=salida)

demo.launch()