from transformers import pipeline, MarianMTModel, MarianTokenizer
import gradio as gr

summarizer = pipeline("summarization", model="Falconsai/text_summarization") # modelo de resumen de texto inglés a inglés

# modelo para la traducción del texto ingles a español
model_name = "Helsinki-NLP/opus-mt-en-es"
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)

def traducir_resumir(texto):
    text_summarized = summarizer(texto, max_length=1000, min_length=100, do_sample=False)
    texto_solo = ">>esp<< " + text_summarized[0]["summary_text"]
    translated_resumido = model.generate(**tokenizer(texto_solo, return_tensors="pt", padding=True))
    texto_resumido = [tokenizer.decode(t, skip_special_tokens=True) for t in translated_resumido]
    return texto_resumido[0]

def traducir(texto):
    texto_unido = ""
    texto_completo = ""
    texto_separado = texto.split(" ")
    for i in range(len(texto_separado)):
        texto_unido = texto_unido + "".join(texto_separado[i]) + " "
        if i//(len(texto_separado)-1) == 1 or i%100 == 0:
            texto_unido_traducido = model.generate(**tokenizer(texto_unido, return_tensors="pt", padding=True))
            texto_unido = [tokenizer.decode(t, skip_special_tokens=True) for t in texto_unido_traducido]
            texto_completo = texto_completo + "".join(texto_unido[0]) + " "    
            texto_unido = ""
    return texto_completo

with gr.Blocks() as demo:
    gr.Markdown("### Proyecto 1 - Trabajo con modelos de IA")
    with gr.Row():
        with gr.Column():
            text = gr.Textbox(label="Introduzca el texto a resumir y traducir")
            with gr.Row():
                with gr.Column():
                    sumbtn = gr.Button("Resumir y traducir")
                with gr.Column():
                    tradbtn = gr.Button("Traducir texto")
        with gr.Column():
            output1 = gr.Textbox(label="Texto resumido traducido")
            output2 = gr.Textbox(label="Texto completo traducido")
    with gr.Accordion("Información"):
        with gr.Row():
            with gr.Column():
                with gr.Accordion("Detalles del proyecto", open = False):
                    gr.Markdown("""El problema principal que resuelve este proyecto es para adquirir información de textos en inglés cuando no se sabe este idioma.\n
                    Esto es un problema bastante comun dado que hay mucha gente en España la cual no sabe inglés y si quiere informarse de cualquier tema el cual no está \
                    muy desarrollado en español tiene que buscarse la vida para poder traducir esos textos a su idioma natal.\n
                    Gracias a este proyecto se puede traducir estos textos rápidamente con un función añadida de hacerle un resumen para poder entenderlo mejor. \n
                    El modelo requiere como input un texto en ingles y sacara como outputs un texto resumido en español y el mismo texto completo traducido al español.
                    """
                    )
            with gr.Column():
                with gr.Accordion("Detalles de los modelos", open = False):
                    gr.Markdown("""#### Helsinki-NLP/opus-mt-en-es (Modelo traductor)
                    Este modelo es el que se usa para traducir el texto de ingles a español. Se usa desde un framework (MarianMT) el cual contiene más de 1000 modelos de lenguaje.\
                    Para usar el modelo específico de traducir de ingles a español se tiene que poner las etiquetas de idioma (en este caso en y es) junto al modelo en si \
                    (Helsinki-NLP/opus-mt-{src}-{trg}). El modelo necesita tokenizar el texto para poder traducirlo. Esto se hace con su propio tokenizador el cual es MarianTokenizer.\
                    La duración de la traducción viene dada a la longitud del texto que se quiere traducir por lo que textos más largos tardarán más en traducirse.\n
                    El principal problema del modelo es que para la versión que tenemos solo permite 512 tokens al traducir lo cual limita mucho, por eso traduzco el texto \
                    partiendolo en segmentos de misma longitud (100 en este caso)\n
                    #### Falconsai (Modelo que resume)
                    Este modelo es un fine-tuning hecho para resumir de otro modelo llamado T5 Small. Es un modelo muy sencillo el cual puede traducir textos muy extensos \
                    y de forma rápida y sencilla. La forma en la que lo uso es con el pipeline de transformer y dentro de este se le pone los parametros max_lenght y min_length \
                    los cuales permite poner la longitud máxima y mínima para el texto resumido."""
                    )
    
    sumbtn.click(fn = traducir_resumir, inputs = text, outputs = output1)
    tradbtn.click(fn = traducir, inputs = text, outputs = output2)
# demo = gr.Interface(
#     fn=traducir,
#     inputs=[gr.Textbox(label="Texto")],
#     outputs=[gr.Textbox(label="Texto completo traducido"), gr.Textbox(label="Texto resumido traducido")],
# )

demo.launch()