improve effeciency

app.py

@@ -1,8 +1,7 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelWithLMHead
+from transformers import pipeline
 import requests
 from bs4 import BeautifulSoup
-from nltk.corpus import stopwords
 import nltk
 import string
 from streamlit.components.v1 import html
@@ -78,18 +77,19 @@ def find_source(text, docs):
 @st.experimental_singleton
 def init_models():
     nltk.download('stopwords')
+    from nltk.corpus import stopwords
     stop = set(stopwords.words('english') + list(string.punctuation))
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     question_answerer = pipeline(
         "question-answering", model='sultan/BioM-ELECTRA-Large-SQuAD2-BioASQ8B',
         device=device
     )
-    reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device=device)
-    queryexp_tokenizer = AutoTokenizer.from_pretrained("doc2query/all-with_prefix-t5-base-v1")
-    queryexp_model = AutoModelWithLMHead.from_pretrained("doc2query/all-with_prefix-t5-base-v1")
-    return question_answerer, reranker, stop, device, queryexp_model, queryexp_tokenizer
+    reranker = CrossEncoder('cross-encoder/ms-marco-TinyBERT-L-2-v2', device=device)
+    # queryexp_tokenizer = AutoTokenizer.from_pretrained("doc2query/all-with_prefix-t5-base-v1")
+    # queryexp_model = AutoModelWithLMHead.from_pretrained("doc2query/all-with_prefix-t5-base-v1")
+    return question_answerer, reranker, stop, device # uqeryexp_model, queryexp_tokenizer
 
-qa_model, reranker, stop, device, queryexp_model, queryexp_tokenizer = init_models()
+qa_model, reranker, stop, device = init_models() # queryexp_model, queryexp_tokenizer
 
 
 def clean_query(query, strict=True, clean=True):
@@ -157,27 +157,27 @@ with st.expander("Settings (strictness, context limit, top hits)"):
     use_reranking = st.radio(
         "Use Reranking? Reranking will rerank the top hits using semantic similarity of document and query.",
         ('yes', 'no'))
-    top_hits_limit = st.slider('Top hits? How many documents to use for reranking. Larger is slower but higher quality', 10, 300, 200 if torch.cuda.is_available() else 50)
-    context_lim = st.slider('Context limit? How many documents to use for answering from. Larger is slower but higher quality', 10, 300, 25 if torch.cuda.is_available() else 10)
+    top_hits_limit = st.slider('Top hits? How many documents to use for reranking. Larger is slower but higher quality', 10, 300, 200)
+    context_lim = st.slider('Context limit? How many documents to use for answering from. Larger is slower but higher quality', 10, 300, 25)
     use_query_exp = st.radio(
         "(Experimental) use query expansion? Right now it just recommends queries",
         ('yes', 'no'))
     suggested_queries = st.slider('Number of suggested queries to use', 0, 10, 5)
 
-def paraphrase(text, max_length=128):
-    input_ids = queryexp_tokenizer.encode(text, return_tensors="pt", add_special_tokens=True)
-    generated_ids = queryexp_model.generate(input_ids=input_ids, num_return_sequences=suggested_queries or 5, num_beams=suggested_queries or 5, max_length=max_length)
-    queries = set([queryexp_tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in generated_ids])
-    preds = '\n * '.join(queries)
-    return preds
+# def paraphrase(text, max_length=128):
+#     input_ids = queryexp_tokenizer.encode(text, return_tensors="pt", add_special_tokens=True)
+#     generated_ids = queryexp_model.generate(input_ids=input_ids, num_return_sequences=suggested_queries or 5, num_beams=suggested_queries or 5, max_length=max_length)
+#     queries = set([queryexp_tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in generated_ids])
+#     preds = '\n * '.join(queries)
+#     return preds
 
 def run_query(query):
-    if use_query_exp == 'yes':
-        query_exp = paraphrase(f"question2question: {query}")
-        st.markdown(f"""
-If you are not getting good results try one of:
-* {query_exp}
-""")
+#     if use_query_exp == 'yes':
+#         query_exp = paraphrase(f"question2question: {query}")
+#         st.markdown(f"""
+# If you are not getting good results try one of:
+# * {query_exp}
+# """)
     limit = top_hits_limit or 100
     context_limit = context_lim or 10
     contexts, orig_docs = search(query, limit=limit, strict=strict_mode == 'strict')