import gradio as gr
import torch
import transformers as trf

# Load the summarization model
summarization_model_path = 'ieuniversity/summarization_model_translator'
summarization_tokenizer = trf.AutoTokenizer.from_pretrained("it5/it5-base-news-summarization")
summarization_model = trf.AutoModelForSeq2SeqLM.from_pretrained(summarization_model_path)

# Load the translation model
translation_model_path = 'hackathon-pln-es/t5-small-finetuned-spanish-to-quechua'
translation_tokenizer = trf.AutoTokenizer.from_pretrained(translation_model_path)
translation_model = trf.AutoModelForSeq2SeqLM.from_pretrained(translation_model_path)

def summarize_and_translate(news_text):
    # Summarize the news article
    max_input_length = 512
    max_output_length = 128
    input_encoded = summarization_tokenizer.encode_plus(news_text, add_special_tokens=True,
                                                         max_length=max_input_length, pad_to_max_length=True,
                                                         return_attention_mask=True, return_tensors='pt')
    input_ids = input_encoded['input_ids']
    attention_mask = input_encoded['attention_mask']
    output_ids = summarization_model.generate(input_ids=input_ids, attention_mask=attention_mask,
                                               max_length=max_output_length)
    summary_text = summarization_tokenizer.decode(output_ids[0], skip_special_tokens=True)
    
    # Translate the summary to Quechua
    input_encoded = translation_tokenizer(summary_text, padding=True, truncation=True, max_length=512, return_tensors='pt')
    input_ids = input_encoded['input_ids']
    attention_mask = input_encoded['attention_mask']
    output_ids = translation_model.generate(input_ids=input_ids, attention_mask=attention_mask,
                                             max_length=512)
    output_text = translation_tokenizer.decode(output_ids[0], skip_special_tokens=True)
    
    return output_text

# Define the input and output interfaces for Gradio
input_interface = gr.inputs.Textbox(label="Input your News text! (Spanish)")
output_interface = gr.outputs.Textbox(label="Your Summarized News Text in Native Quechua!")

# Add poster link to the interface description
description = "This is a Spanish-to-Quechua news summarization and translation app. It uses a state-of-the-art AI model to summarize news articles in Spanish from the newspaper La Razón and translate the summary to Quechua. You can learn more about this project and its creators by visiting this <a href='https://files.fm/u/ktszfunsz'>poster</a>."

# Create and launch the Gradio app
iface = gr.Interface(fn=summarize_and_translate, inputs=input_interface, outputs=output_interface,
                     title="Spanish-to-Quechua News Summarization and Translation", description=description)
iface.launch()