Pre-load models

app.py

@@ -3,15 +3,30 @@ from extractor import extract, FewDocumentsError
 from summarizer import summarize
 import time
 import cProfile
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
-# Dowload required NLTK resources
-from nltk import download
-download('punkt')
-download('stopwords')
+@st.cache(allow_output_mutation=True)
+def init():
+    # Dowload required NLTK resources
+    from nltk import download
+    download('punkt')
+    download('stopwords')
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # Model for semantic searches
+    search_model = SentenceTransformer('msmarco-distilbert-base-v4', device=device)
+    # Model for abstraction
+    summ_model = AutoModelForSeq2SeqLM.from_pretrained('t5-base')
+    tokenizer = AutoTokenizer.from_pretrained('t5-base')
+
+    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")
@@ -31,7 +46,7 @@ def main():
         start_time = time.time()
         try:
             with st.spinner('Extraindo textos relevantes...'):
-                text = extract(query)
+                text = extract(query, search_model=search_model)
         except FewDocumentsError as e:
             few_documents = True
             st.session_state['few_documents'] = True
@@ -41,7 +56,7 @@ def main():
 
             st.info(f'(Extraction) Elapsed time: {time.time() - start_time:.2f}s')
             with st.spinner('Gerando resumo...'):
-                summary = summarize(text)
+                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}')
@@ -52,10 +67,10 @@ def main():
         if st.button('Prosseguir'):
             start_time = time.time()
             with st.spinner('Extraindo textos relevantes...'):
-                text = extract(query, extracted_documents=st.session_state['documents'])
+                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)
+                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}')

extractor/_utils.py

@@ -4,8 +4,6 @@ import streamlit as st
 # import inflect
 import torch
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
 # p = inflect.engine()
 
 class FewDocumentsError(Exception):
@@ -90,8 +88,8 @@ def paragraph_extraction(documents, min_paragraph_size):
   return paragraphs
 
 def semantic_search(model, query, files, number_of_similar_files):
-  encoded_query = model.encode(query, device=device)
-  encoded_files = model.encode(files, device=device)
+  encoded_query = model.encode(query)
+  encoded_files = model.encode(files)
 
   model_index = nmslib.init(method='hnsw', space='angulardist')
   model_index.addDataPointBatch(encoded_files)

extractor/extract.py

@@ -1,4 +1,3 @@
-from sentence_transformers import SentenceTransformer
 from ._utils import FewDocumentsError
 from ._utils import document_extraction, paragraph_extraction, semantic_search
 from corpora import gen_corpus
@@ -6,9 +5,7 @@ from nltk.corpus import stopwords
 from nltk.tokenize import word_tokenize
 import string
 
-from ._utils import device
-
-def extract(query: str, n: int=3, extracted_documents: list=None) -> str:
+def extract(query: str, search_model, n: int=3, extracted_documents: list=None) -> str:
     """Extract n paragraphs from the corpus using the given query.
 
     Parameters:
@@ -38,8 +35,6 @@ def extract(query: str, n: int=3, extracted_documents: list=None) -> str:
         )
 
     # First semantc search (over documents)
-    # Model for semantic searches
-    search_model = SentenceTransformer('msmarco-distilbert-base-v4', device=device)
     selected_documents, documents_distances = semantic_search(
         model=search_model,
         query=query,

summarizer/summarize.py

@@ -1,14 +1,9 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-def summarize(text: str) -> str:
+def summarize(text: str, model, tokenizer) -> str:
     """
     Generate a summary based from the given text
     """
 
-    # Model for abstraction
-    model = AutoModelForSeq2SeqLM.from_pretrained('t5-base')
-    tokenizer = AutoTokenizer.from_pretrained('t5-base')
-
     input_tokens = tokenizer.encode(
         f'summarize: {text}',
         return_tensors='pt',