Add translation

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from extractor import extract, FewDocumentsError
 from summarizer import summarize
+from translation import translate
 import time
 import cProfile
 from sentence_transformers import SentenceTransformer
@@ -23,16 +24,22 @@ def init():
 
     return search_model, summ_model, tokenizer
 
-# TODO: translation
-
 def main():
     search_model, summ_model, tokenizer = init()
 
-    st.title("Trabalho de Formatura - Construindo textos para a internet")
-    st.subheader("Lucas Antunes e Matheus Vieira")
+    st.title("AutoSumm")
+    st.subheader("Lucas Antunes & Matheus Vieira")
+
+    portuguese = st.checkbox('Traduzir para o português.')
 
-    st.subheader("Digite o tópico sobre o qual você deseja gerar um resumo")
-    query = st.text_input('Digite o tópico em inglês') #text is stored in this variable
+    if portuguese:
+        st.subheader("Digite o tópico sobre o qual você deseja gerar um resumo")
+        query_pt = st.text_input('Digite o tópico') #text is stored in this variable
+        button = st.button('Gerar resumo')
+    else:
+        st.subheader("Type the desired topic to generate the summary")
+        query = st.text_input('Type your topic') #text is stored in this variable
+        button = st.button('Generate summary')
 
     if 'few_documents' not in st.session_state:
         st.session_state['few_documents'] = False
@@ -40,10 +47,9 @@ def main():
     else:
         few_documents = st.session_state['few_documents']
 
-    button1 = st.button('Gerar resumo')
-
-    if button1:
+    if button:
         start_time = time.time()
+        query = translate(query_pt, 'pt', 'en') if portuguese else query
         try:
             with st.spinner('Extraindo textos relevantes...'):
                 text = extract(query, search_model=search_model)
@@ -59,7 +65,10 @@ def main():
                 summary = summarize(text, summ_model, tokenizer)
             st.info(f'(Total) Elapsed time: {time.time() - start_time:.2f}s')
 
-            st.markdown(f'Seu resumo para "{query}":\n\n> {summary}')
+            if portuguese:
+                st.markdown(f'Seu resumo para "{query_pt}":\n\n> {translate(summary, "en", "pt")}')
+            else:
+                st.markdown(f'Your summary for "{query}":\n\n> {summary}')
 
 
     if few_documents:
@@ -73,7 +82,10 @@ def main():
                 summary = summarize(text, summ_model, tokenizer)
             st.info(f'(Total) Elapsed time: {time.time() - start_time:.2f}s')
 
-            st.markdown(f'Seu resumo para "{query}":\n\n> {summary}')
+            if portuguese:
+                st.markdown(f'Seu resumo para "{query_pt}":\n\n> {translate(summary, "en", "pt")}')
+            else:
+                st.markdown(f'Your summary for "{query}":\n\n> {summary}')
 
             st.session_state['few_documents'] = False
             few_documents = False

translation/__init__.py

@@ -0,0 +1 @@
+from .translation import translate

translation/translation.py

@@ -0,0 +1,11 @@
+from deep_translator import GoogleTranslator
+from easynmt import EasyNMT
+
+def translate(text, source_language, target_language):
+    try:
+        print("Trying to use Google Translator...")
+        return GoogleTranslator(source = source_language, target = target_language).translate(text)
+    except:
+        print("Google Translator is not working... using M2M100")
+        translation_model = EasyNMT('m2m_100_418M')
+        return translation_model.translate(text, source_lang = source_language, target_lang = target_language)