import gradio as gr
import torch
from transformers import BertTokenizerFast, EncoderDecoderModel

# Set the model name to use
model_name= "LeoCordoba/beto2beto-cc-news-es-titles"
# Identify the device
device = 'cuda' if torch.cuda.is_available() else 'cpu'

def get_summary(text):
    # Load the tokenizer and the EncoderDecoder model
    tokenizer = BertTokenizerFast.from_pretrained(model_name)
    model = EncoderDecoderModel.from_pretrained(model_name).to(device)
    # Tokenize the input text
    inputs = tokenizer([text], padding="max_length",
                       truncation=True, max_length=512, return_tensors="pt")
    input_ids = inputs.input_ids.to(device)
    attention_mask = inputs.attention_mask.to(device)
    # Generate the model output
    output = model.generate(input_ids, attention_mask=attention_mask)
    # Return the decoded output
    return tokenizer.decode(output[0], skip_special_tokens=True)

theme = "darkhuggingface"
title = "A Summarizer for News in Spanish"
description = "Insert an article in spanish and press Submit / Inserta  un artículo en español y pulsa Submit"
article= "This is an Space App for a mini-project on an Encoder Decoder model in a non-english language by edumunozsala - based on the model LeoCordoba/beto2beto-cc-news-es-titles by LeoCordoba."
examples = [["El consejero de Educación y Deporte, Javier Imbroda, ha subrayado este martes, tras el acuerdo alcanzado con los sindicatos CCOO, CSIF y ANPE para ampliar en 6.949 docentes la dotación adicional de profesores de la educación pública para el curso 2020-21, que nunca tuvo la educación pública tanto apoyo, tanto refuerzo, ni tanto Presupuesto para un curso."],["Las observaciones coinciden perfectamente con los tiempos de difusión de la epidemia y las decisiones del gobierno chino, que el 23 de enero cerró el transporte desde y hacia Wuhan, el foco epidémico inicial, y los negocios locales, para reducir la expansión del virus, en lo que fue la primera de varias severas cuarentenas dispuestas en China y en otros lugares del mundo. La primera imagen, antes de la cuarentena , coincide con el inicio del Año Nuevo Lunar en China y otras partes de Asia, cuando suelen observarse caídas en los niveles de dióxido de nitrógeno."]]

interface = gr.Interface(fn=get_summary, inputs=[gr.inputs.Textbox(
    lines=5, label="Article / Artículo")], outputs=[gr.outputs.Textbox(
    type='str', label="Summary / Resumen")],
    theme=theme, title=title, description=description, examples=examples, article=article, layout = 'vertical', enable_queue=True)

interface.launch(inline=False, debug=False)