News-Translator / app.py
juanesbch's picture
Update app.py
e413b5a
raw
history blame
No virus
2.84 kB
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()