Add timer

app.py

@@ -2,11 +2,12 @@ import streamlit as st
 from extractor import extract, FewDocumentsError
 from summarizer import summarize
 from translation import translate
-import time
+from utils.timing import Timer
 import cProfile
 from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
+from os import environ
 
 @st.cache(allow_output_mutation=True)
 def init():
@@ -26,6 +27,7 @@ def init():
 
 def main():
     search_model, summ_model, tokenizer = init()
+    Timer.reset()
 
     st.title("AutoSumm")
     st.subheader("Lucas Antunes & Matheus Vieira")
@@ -33,14 +35,18 @@ def main():
     portuguese = st.checkbox('Traduzir para o português.')
 
     if portuguese:
+        environ['PORTUGUESE'] = 'true' # work around (gambiarra)
         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:
+        environ['PORTUGUESE'] = 'false' # work around (gambiarra)
         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')
 
+    result = st.empty()
+
     if 'few_documents' not in st.session_state:
         st.session_state['few_documents'] = False
         few_documents = False
@@ -48,11 +54,9 @@ def main():
         few_documents = st.session_state['few_documents']
 
     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)
+            text = extract(query, search_model=search_model)
         except FewDocumentsError as e:
             few_documents = True
             st.session_state['few_documents'] = True
@@ -60,32 +64,26 @@ def main():
             st.session_state['msg'] = e.msg
         else:
 
-            st.info(f'(Extraction) Elapsed time: {time.time() - start_time:.2f}s')
-            with st.spinner('Gerando resumo...'):
-                summary = summarize(text, summ_model, tokenizer)
-            st.info(f'(Total) Elapsed time: {time.time() - start_time:.2f}s')
+            summary = summarize(text, summ_model, tokenizer)
 
             if portuguese:
-                st.markdown(f'Seu resumo para "{query_pt}":\n\n> {translate(summary, "en", "pt")}')
+                result.markdown(f'Seu resumo para "{query_pt}":\n\n> {translate(summary, "en", "pt")}')
             else:
-                st.markdown(f'Your summary for "{query}":\n\n> {summary}')
+                result.markdown(f'Your summary for "{query}":\n\n> {summary}')
+
+            Timer.show_total()
 
 
     if few_documents:
         st.warning(st.session_state['msg'])
         if st.button('Prosseguir'):
-            start_time = time.time()
-            with st.spinner('Extraindo textos relevantes...'):
-                text = extract(query, search_model=search_model, extracted_documents=st.session_state['documents'])
-            st.info(f'(Extraction) Elapsed time: {time.time() - start_time:.2f}s')
-            with st.spinner('Gerando resumo...'):
-                summary = summarize(text, summ_model, tokenizer)
-            st.info(f'(Total) Elapsed time: {time.time() - start_time:.2f}s')
+            text = extract(query, search_model=search_model, extracted_documents=st.session_state['documents'])
+            summary = summarize(text, summ_model, tokenizer)
 
             if portuguese:
-                st.markdown(f'Seu resumo para "{query_pt}":\n\n> {translate(summary, "en", "pt")}')
+                result.markdown(f'Seu resumo para "{query_pt}":\n\n> {translate(summary, "en", "pt")}')
             else:
-                st.markdown(f'Your summary for "{query}":\n\n> {summary}')
+                result.markdown(f'Your summary for "{query}":\n\n> {summary}')
 
             st.session_state['few_documents'] = False
             few_documents = False

extractor/extract.py

@@ -1,10 +1,12 @@
 from ._utils import FewDocumentsError
 from ._utils import document_extraction, paragraph_extraction, semantic_search
+from utils.timing import Timer
 from corpora import gen_corpus
 from nltk.corpus import stopwords
 from nltk.tokenize import word_tokenize
 import string
 
+@Timer.time_it('extração', 'extraction')
 def extract(query: str, search_model, n: int=3, extracted_documents: list=None) -> str:
     """Extract n paragraphs from the corpus using the given query.
 
@@ -16,7 +18,8 @@ def extract(query: str, search_model, n: int=3, extracted_documents: list=None)
     str: String containing the n most relevant paragraphs joined by line breaks
     """
     # Open corpus
-    corpus = gen_corpus(query)
+    with Timer('geração do corpus', 'corpus generation'):
+        corpus = gen_corpus(query)
 
     # Setup query
     stop_words = set(stopwords.words('english'))
@@ -25,36 +28,37 @@ def extract(query: str, search_model, n: int=3, extracted_documents: list=None)
     keywords = [keyword for keyword in tokens_without_sw if keyword not in string.punctuation]
 
     # Gross search
-    if not extracted_documents:
-        extracted_documents, documents_empty, documents_sizes = document_extraction(
-            dataset=corpus,
-            query=query,
-            keywords=keywords,
-            min_document_size=0,
-            min_just_one_paragraph_size=0
-        )
+    with Timer('busca exaustiva', 'exhaustive search'):
+        if not extracted_documents:
+            extracted_documents, documents_empty, documents_sizes = document_extraction(
+                dataset=corpus,
+                query=query,
+                keywords=keywords,
+                min_document_size=0,
+                min_just_one_paragraph_size=0
+            )
 
     # First semantc search (over documents)
-    selected_documents, documents_distances = semantic_search(
-        model=search_model,
-        query=query,
-        files=extracted_documents,
-        number_of_similar_files=10
-    )
+    with Timer('busca semantica nos documentos', 'semantic search over documents'):
+        selected_documents, documents_distances = semantic_search(
+            model=search_model,
+            query=query,
+            files=extracted_documents,
+            number_of_similar_files=10
+        )
 
     # Second semantic search (over paragraphs)
-    paragraphs = paragraph_extraction(
-        documents=selected_documents,
-        min_paragraph_size=20,
-    )
-    
-    # Model for the second semantic search
-    selected_paragraphs, paragraphs_distances = semantic_search(
-        model=search_model,
-        query=query,
-        files=paragraphs,
-        number_of_similar_files=10
-    )
+    with Timer('busca semantica nos parágrafos', 'semantic search over paragraphs'):
+        paragraphs = paragraph_extraction(
+            documents=selected_documents,
+            min_paragraph_size=20,
+        )
+        selected_paragraphs, paragraphs_distances = semantic_search(
+            model=search_model,
+            query=query,
+            files=paragraphs,
+            number_of_similar_files=10
+        )
 
     text = '\n'.join(selected_paragraphs[:n])
 

summarizer/summarize.py

@@ -1,4 +1,6 @@
+from utils.timing import Timer
 
+@Timer.time_it('abstração', 'abstraction')
 def summarize(text: str, model, tokenizer) -> str:
     """
     Generate a summary based from the given text

translation/translation.py

@@ -1,6 +1,8 @@
 from deep_translator import GoogleTranslator
 from easynmt import EasyNMT
+from utils.timing import Timer
 
+@Timer.time_it('tradução', 'translation')
 def translate(text, source_language, target_language):
     try:
         print("Trying to use Google Translator...")

utils/timing.py

@@ -0,0 +1,79 @@
+import time
+import streamlit as st
+from os import environ
+
+class Timer():
+    total = 0
+    expander = None
+    def __init__(self, pt_name, en_name):
+        self.pt_name = pt_name
+        self.en_name = en_name
+        if environ['PORTUGUESE'] == 'true':
+            self.portuguese = True
+        elif environ['PORTUGUESE'] == 'false':
+            self.portuguese = False
+        else:
+            raise EnvironmentError
+        if not Timer.expander:
+            if self.portuguese:
+                Timer.expander = st.expander('Ver progresso')
+            else:
+                Timer.expander = st.expander('See progress')
+        self.display = Timer.expander.empty()
+            
+    def __enter__(self):
+        if self.portuguese:
+            self.display.info(f'Executando "{self.pt_name}"...')
+        else:
+            self.display.info(f'Running "{self.en_name}"...')
+        self.start_time = time.time()
+
+    def __exit__(self, type, value, traceback):
+        end_time = time.time()
+        elapsed_time = end_time - self.start_time
+        Timer.total += elapsed_time
+        self.display.empty()
+        if self.portuguese:
+            Timer.expander.info(f'"{self.pt_name}" terminou em {elapsed_time:.2f} s')
+        else:
+            Timer.expander.info(f'"{self.en_name}" finished in {elapsed_time:.2f} s')
+
+    # for manually starting the timer
+    def start(self):
+        if self.portuguese:
+            self.display.warning(f'Executando "{self.pt_name}"...')
+        else:
+            self.display.warning(f'Running "{self.en_name}"...')
+        self.start_time = time.time()
+
+    # for manually stopping the timer
+    def stop(self):
+        end_time = time.time()
+        elapsed_time = end_time - self.start_time
+        Timer.total += elapsed_time
+        self.display.empty()
+        if self.portuguese:
+            Timer.expander.warning(f'"{self.pt_name}" terminou em {elapsed_time:.2f} s')
+        else:
+            Timer.expander.warning(f'"{self.en_name}" finished in {elapsed_time:.2f} s')
+
+    def reset():
+        Timer.total = 0
+        Timer.expander = None
+
+    def show_total():
+        if environ['PORTUGUESE'] == 'true':
+            Timer.expander.success(f'Tempo de execução total: {Timer.total:.2f} s')
+        elif environ['PORTUGUESE'] == 'false':
+            Timer.expander.success(f'Total elapsed time: {Timer.total:.2f} s')
+
+    def time_it(pt_name, en_name):
+        def decorator(func):
+            def wrapper(*args, **kwargs):
+                timer = Timer(pt_name, en_name)
+                timer.start()
+                result = func(*args, **kwargs)
+                timer.stop()
+                return result
+            return wrapper
+        return decorator