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!")

# 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")
iface.launch()